From b73159dd1de8db678091cba7363570a41f33413f Mon Sep 17 00:00:00 2001
From: Rodz Labs <rodzilla@free.fr>
Date: Tue, 18 Jul 2023 07:44:12 +0200
Subject: [PATCH] Updated ambient occlusion map generation to use
 RenderingDevice

---
 addons/material_maker/bvh_generator.gd        | 449 ++++++++++++++++++
 .../engine/pipeline/compute_shader.gd         |   3 +
 .../engine/pipeline/pipeline.gd               |   2 +-
 .../engine/pipeline/rendering_pipeline.gd     |  37 +-
 .../engine/{ => pipeline}/texture.gd          |   0
 .../map_renderer/ao_fragment.tres             | 263 ++++++++++
 .../map_renderer/ao_vertex.tres               |  35 ++
 .../map_renderer/denoise_compute.tres         |  34 ++
 .../map_renderer/map_generator.gd             |  55 ++-
 ...mal_map_vertex.tres => normal_vertex.tres} |   0
 ...n_map_vertex.tres => position_vertex.tres} |   0
 ...nt_map_vertex.tres => tangent_vertex.tres} |   0
 .../tools/map_renderer/map_renderer.gd        |  16 +-
 .../tools/map_renderer/map_renderer.tscn      |  17 +-
 14 files changed, 872 insertions(+), 39 deletions(-)
 create mode 100644 addons/material_maker/bvh_generator.gd
 rename addons/material_maker/engine/{ => pipeline}/texture.gd (100%)
 create mode 100644 addons/material_maker/map_renderer/ao_fragment.tres
 create mode 100644 addons/material_maker/map_renderer/ao_vertex.tres
 create mode 100644 addons/material_maker/map_renderer/denoise_compute.tres
 rename addons/material_maker/map_renderer/{normal_map_vertex.tres => normal_vertex.tres} (100%)
 rename addons/material_maker/map_renderer/{position_map_vertex.tres => position_vertex.tres} (100%)
 rename addons/material_maker/map_renderer/{tangent_map_vertex.tres => tangent_vertex.tres} (100%)

diff --git a/addons/material_maker/bvh_generator.gd b/addons/material_maker/bvh_generator.gd
new file mode 100644
index 000000000..e3e7d5329
--- /dev/null
+++ b/addons/material_maker/bvh_generator.gd
@@ -0,0 +1,449 @@
+extends Object
+class_name MMBvhGenerator
+
+static func generate(mesh: Mesh) -> ImageTexture:
+	var b_mesh := MeshDataTool.new()
+	if not mesh is ArrayMesh:
+		b_mesh.create_from_surface(mesh.create_outline(0.0), 0)
+	else:
+		b_mesh.create_from_surface(mesh, 0)
+
+	if b_mesh.get_vertex_count() == 0 or b_mesh.get_face_count() == 0:
+		return ImageTexture.new()
+
+	var triangles := []
+	var vertices := []
+	#print("%d vertices" % b_mesh.get_vertex_count())
+	for i in b_mesh.get_vertex_count():
+		vertices.append(b_mesh.get_vertex(i))
+	#print("%d faces" % b_mesh.get_face_count())
+	for i in b_mesh.get_face_count():
+		triangles.append({
+			0: b_mesh.get_face_vertex(i, 0 if mesh is ArrayMesh else 2),
+			1: b_mesh.get_face_vertex(i, 1),
+			2: b_mesh.get_face_vertex(i, 2 if mesh is ArrayMesh else 0),
+			3: i
+		})
+
+	var bvh := BVHNode.new()
+	bvh.vertices = vertices
+	bvh.triangles = triangles
+	print("Generating BVH...")
+	var time := Time.get_ticks_msec()
+	var max_node_level := bvh.generate()
+	print("BVH Generated!")
+	print("Maximum Node Level: %d" % max_node_level)
+	print("Time to generate BVH: %f" % ((Time.get_ticks_msec() - time) / 1000.0))
+
+	return ImageTexture.create_from_image(bvh.image)
+
+
+# The BVH (Bounding Volume Hierarchy) is used to making ray tracing on the GPU much more efficient.
+class BVHNode:
+	const MAX_TRIANGLES = 8
+	const DATA_HEADER_SIZE = 1
+
+	var node_debug: Node3D
+
+	var aabb := AABB()
+	var center: Vector3
+	var level := 0
+
+	var root := self
+	var parent: BVHNode
+	var left_node: BVHNode
+	var right_node: BVHNode
+
+	var triangles := []
+	var vertices := []
+
+	# The properties that follow are only made in the root node.
+	var image: Image
+
+	var _node_ids: Dictionary
+	var _id_nodes: Dictionary
+	var _node_data: Array
+	var _data: Array
+
+	# Gets all the nodes triangles, including its children.
+	func get_triangles() -> Array:
+		var tris := triangles.duplicate()
+		if root == self:
+			return tris
+
+		if left_node:
+			tris += left_node.get_triangles()
+			tris += right_node.get_triangles()
+		return tris
+
+	# Takes the triangles and splits itself into smaller nodes until
+	# each node has less than the maximum triangle count.
+	# Returns the maximum level that the bvh generates down to.
+	func generate() -> int:
+		calculate_aabb()
+
+		if root == self:
+			_id_nodes = {}
+			_node_ids = {}
+			_node_ids[self] = 0
+			_id_nodes[0] = self
+
+			_data = [0, 0, 0, 0]
+			_node_data = []
+		else:
+			var id := root._node_ids.size()
+			root._node_ids[self] = id
+			root._id_nodes[id] = self
+
+		var data_offset := root._node_data.size()
+# warning-ignore:integer_division
+		_append_4(root._data, data_offset / 4)
+		_append_4(root._node_data, aabb.position.x, aabb.position.y, aabb.position.z, level)
+		_append_4(root._node_data, aabb.end.x, aabb.end.y, aabb.end.z, 0)
+
+		# Stop splitting if there's not enough triangles.
+		if triangles.size() <= MAX_TRIANGLES:
+			root._node_data[-1] = triangles.size()
+			for tri in triangles:
+				var vert_a: Vector3 = vertices[tri[0]]
+				var vert_b: Vector3 = vertices[tri[1]]
+				var vert_c: Vector3 = vertices[tri[2]]
+				_append_4(root._node_data, vert_a.x, vert_a.y, vert_a.z)
+				_append_4(root._node_data, vert_b.x, vert_b.y, vert_b.z)
+				_append_4(root._node_data, vert_c.x, vert_c.y, vert_c.z)
+			if root == self:
+				_finalize_data()
+			return level
+
+		# Generate triangle centroids if there's none.
+		if not triangles[0].has("centroid"):
+			for i in triangles.size():
+				var tri: Dictionary = triangles[i]
+				var tri_aabb := AABB()
+				tri_aabb.position = vertices[tri[0]]
+				tri_aabb.size = Vector3.ZERO
+				tri_aabb = tri_aabb.expand(vertices[tri[1]])
+				tri_aabb = tri_aabb.expand(vertices[tri[2]])
+				triangles[i]["centroid"] = tri_aabb.position + tri_aabb.size * 0.5
+
+		# Split the triangles between the 6 axes directions.
+		var left_nodes := [[], [], []]
+		var right_nodes := [[], [], []]
+		for i in triangles.size():
+			for j in 3:
+				if triangles[i].centroid[j] < center[j]:
+					left_nodes[j].append(i)
+				else:
+					right_nodes[j].append(i)
+
+		# Check which splits have failed.
+		# If they all have, we can't split any further.
+		var split_failed = [false, false, false]
+		split_failed[0] = left_nodes[0].is_empty() or right_nodes[0].is_empty()
+		split_failed[1] = left_nodes[1].is_empty() or right_nodes[1].is_empty()
+		split_failed[2] = left_nodes[2].is_empty() or right_nodes[2].is_empty()
+		if split_failed[0] and split_failed[1] and split_failed[2]:
+			root._node_data[-1] = triangles.size()
+			for tri in triangles:
+				var vert_a: Vector3 = vertices[tri[0]]
+				var vert_b: Vector3 = vertices[tri[1]]
+				var vert_c: Vector3 = vertices[tri[2]]
+				_append_4(root._node_data, vert_a.x, vert_a.y, vert_a.z)
+				_append_4(root._node_data, vert_b.x, vert_b.y, vert_b.z)
+				_append_4(root._node_data, vert_c.x, vert_c.y, vert_c.z)
+			if root == self:
+				_finalize_data()
+			return level
+
+		var split_order := [0, 1, 2]
+		split_order.sort_custom(self._sort_by_extents)
+
+		var left_triangles := []
+		var right_triangles := []
+		for split in split_order:
+			if not split_failed[split]:
+				for i in left_nodes[split]:
+					left_triangles.append(triangles[i])
+				for i in right_nodes[split]:
+					right_triangles.append(triangles[i])
+				break
+		if root != self:
+			triangles.clear()
+
+		_append_4(root._node_data)
+
+		# Generate the new children nodes.
+		left_node = BVHNode.new()
+		left_node.root = root
+		left_node.parent = self
+		left_node.level = level + 1
+		left_node.vertices = vertices
+		left_node.triangles = left_triangles
+		var left_max_level := left_node.generate()
+		root._node_data[data_offset + 8] = root._node_ids[left_node]
+
+		right_node = BVHNode.new()
+		right_node.root = root
+		right_node.parent = self
+		right_node.level = level + 1
+		right_node.vertices = vertices
+		right_node.triangles = right_triangles
+		var right_max_level := right_node.generate()
+		root._node_data[data_offset + 9] = root._node_ids[right_node]
+
+		if root == self:
+			_finalize_data()
+		return int(max(left_max_level, right_max_level))
+
+	# This function is a one-to-one translation of the shader function.
+	# Not really used in this project, but makes debugging much easier.
+	func traverse(ray_start: Vector3, ray_dir: Vector3) -> float:
+		if _data.is_empty():
+			return 65536.0
+
+		var offset_to_nodes: int = _get_data(0)[0]
+		var root_data_0 = _get_data(_get_data(1)[0] + offset_to_nodes)
+		var root_data_1 = _get_data(_get_data(1)[0] + offset_to_nodes + 1)
+
+		var min_root := Vector3(root_data_0[0], root_data_0[1], root_data_0[2])
+		var max_root := Vector3(root_data_1[0], root_data_1[1], root_data_1[2])
+
+		var t := _intersect_aabb(ray_start, ray_dir, min_root, max_root)
+		if t == -1:
+			return 65536.0
+
+		var prev_hit := t
+		t = 65536.0 # Set to large number
+		var tri := [Vector3.ZERO, Vector3.ZERO, Vector3.ZERO]
+		var _min_node_idx := 0
+
+		var stack_point := 0
+		var node_stack := [] # ivec3
+		node_stack.resize(128)
+
+		var curr_node_idx := 0
+		var moving_up := false
+
+		for i in 1024:
+			if moving_up and stack_point <= 0:
+				break
+			var node_data_off = _get_data(DATA_HEADER_SIZE + curr_node_idx)[0] + offset_to_nodes
+			var node_data_0 = _get_data(node_data_off)
+			var node_data_1 = _get_data(node_data_off + 1)
+			var level = node_data_0[3]
+
+			if not moving_up: # Moving down node hierarchy
+				if node_data_1[3] > 0: # Is a leaf node
+					for j in range(node_data_off + 2, node_data_off + 2 + node_data_1[3] * 3, 3):
+						var tri_a = _get_data(j)
+						var tri_b = _get_data(j + 1)
+						var tri_c = _get_data(j + 2)
+						var tri_t = Geometry3D.ray_intersects_triangle(ray_start, ray_dir,
+							Vector3(tri_a[0], tri_a[1], tri_a[2]),
+							Vector3(tri_b[0], tri_b[1], tri_b[2]),
+							Vector3(tri_c[0], tri_c[1], tri_c[2])
+						)
+						if tri_t:
+							tri_t = tri_t.distance_to(ray_start)
+							if tri_t < t:
+								tri = [
+									Vector3(tri_a[0], tri_a[1], tri_a[2]),
+									Vector3(tri_b[0], tri_b[1], tri_b[2]),
+									Vector3(tri_c[0], tri_c[1], tri_c[2])
+								]
+#								print(curr_node_idx)
+								_min_node_idx = curr_node_idx
+							t = min(t, tri_t)
+
+					stack_point -= 1
+					if stack_point <= 0:
+						break
+					if node_stack[stack_point][1] == level: # next node in stack is sibling
+						if t < node_stack[stack_point][0]: # no chance to get better hit from sibling
+							stack_point -= 1
+							moving_up = true
+					else:
+						moving_up = true
+					prev_hit = node_stack[stack_point][0]
+					curr_node_idx = node_stack[stack_point][2]
+				else:
+
+					# Push self onto stack
+					node_stack[stack_point] = Vector3(prev_hit, level, curr_node_idx)
+					stack_point += 1
+
+					var child_indices = _get_data(node_data_off + 2)
+					var left_data_off = _get_data(1 + child_indices[0])[0] + offset_to_nodes
+					var left_data_0 = _get_data(left_data_off)
+					var left_data_1 = _get_data(left_data_off + 1)
+					var right_data_off = _get_data(1 + child_indices[1])[0] + offset_to_nodes
+					var right_data_0 = _get_data(right_data_off)
+					var right_data_1 = _get_data(right_data_off + 1)
+
+					var min_left = Vector3(left_data_0[0], left_data_0[1], left_data_0[2])
+					var max_left = Vector3(left_data_1[0], left_data_1[1], left_data_1[2])
+					var min_right = Vector3(right_data_0[0], right_data_0[1], right_data_0[2])
+					var max_right = Vector3(right_data_1[0], right_data_1[1], right_data_1[2])
+
+					var t_left = _intersect_aabb(ray_start, ray_dir, min_left, max_left, true)
+					var t_right = _intersect_aabb(ray_start, ray_dir, min_right, max_right, true)
+
+					if t_right == -1 and t_left != -1: # only left node hit
+						prev_hit = t_left
+						curr_node_idx = child_indices[0]
+					elif t_left == -1 and t_right != -1: # only right node hit
+						prev_hit = t_right
+						curr_node_idx = child_indices[1]
+					elif t_left < t_right and t_left != -1: # left node hits closer
+						node_stack[stack_point] = Vector3(t_right, right_data_0[3], child_indices[1])
+						stack_point += 1
+						prev_hit = t_left
+						curr_node_idx = child_indices[0]
+					elif t_right <= t_left and t_right != -1: # right node hits closer
+						node_stack[stack_point] = Vector3(t_left, left_data_0[3], child_indices[0])
+						stack_point += 1
+						prev_hit = t_right
+						curr_node_idx = child_indices[1]
+					else: # no hit
+						stack_point -= 2
+						if stack_point <= 0:
+							break
+						if node_stack[stack_point][1] == level: # next node in stack is sibling
+							if t < node_stack[stack_point][0]: # no chance to get better hit from sibling
+								stack_point -= 1
+								moving_up = true
+						else:
+							moving_up = true
+						prev_hit = node_stack[max(stack_point, 0)][0]
+						curr_node_idx = node_stack[max(stack_point, 0)][2]
+
+			else: # Moving up hierarchy
+				stack_point -= 1
+				if stack_point <= 0:
+					break
+				if node_stack[stack_point][1] == level: # next node in stack is sibling
+					if t < node_stack[stack_point][0]: # no chance to get better hit from sibling
+						stack_point -= 1
+					else:
+						moving_up = false
+				prev_hit = node_stack[max(stack_point, 0)][0]
+				curr_node_idx = node_stack[max(stack_point, 0)][2]
+
+		var tri_normal: Vector3 = (tri[2] - tri[0]).cross(tri[1] - tri[0])
+		var imm = node_debug.get_node_or_null("../DebugBVHHit")
+		if imm != null:
+			imm.clear()
+			imm.begin(Mesh.PRIMITIVE_TRIANGLES)
+			imm.set_normal(tri_normal)
+			imm.add_vertex(tri[0])
+			imm.add_vertex(tri[1])
+			imm.add_vertex(tri[2])
+			imm.end()
+	#		print("min: " + str(min_node_idx))
+
+		return t
+
+	func calculate_aabb() -> void:
+		aabb.position = vertices[triangles[0][0]]
+		aabb.size = Vector3.ZERO
+		for tri in triangles:
+			for j in 3:
+				aabb = aabb.expand(vertices[tri[j]])
+		center = aabb.position + aabb.size * 0.5
+
+	func debug_aabb() -> Array:
+		var aabbs := []
+
+		var aabb_shape := CollisionShape3D.new()
+		aabb_shape.name = "BVHNode-" + str(root._node_ids[self])
+#		if parent:
+#			aabb_shape.name = "Left-" if parent.left_node == self else "Right-"
+#		else:
+#			aabb_shape.name = "Root-"
+#		aabb_shape.name += str(level)
+
+		aabb_shape.shape = BoxShape3D.new()
+		aabb_shape.shape.size = aabb.size * 0.5
+		aabb_shape.position = aabb.position + aabb.size * 0.5
+		aabbs.append(aabb_shape)
+
+		if left_node:
+			var left_aabbs := left_node.debug_aabb()
+			aabbs += left_aabbs
+			aabb_shape.add_child(left_aabbs[0])
+			left_aabbs[0].position -= aabb_shape.position
+		if right_node:
+			var right_aabbs := right_node.debug_aabb()
+			aabbs += right_aabbs
+			aabb_shape.add_child(right_aabbs[0])
+			right_aabbs[0].position -= aabb_shape.position
+
+		return aabbs
+
+	func _get_data(index: int) -> Array:
+		return [
+			_data[index * 4],
+			_data[index * 4 + 1],
+			_data[index * 4 + 2],
+			_data[index * 4 + 3],
+		]
+
+	# This is faster than appending a whole array to the other array.
+	func _append_4(array: Array, val0=0, val1=0, val2=0, val3=0) -> void:
+		array.append(val0)
+		array.append(val1)
+		array.append(val2)
+		array.append(val3)
+
+	# All the stuff in _data and _node_data are appended and put into image.
+	func _finalize_data() -> void:
+#		var time := OS.get_ticks_msec()
+		_data[0] = _node_ids.size() + DATA_HEADER_SIZE
+
+		_data += _node_data
+		
+# warning-ignore:integer_division
+		var data_size = _data.size() / 4
+		var img_width = min(data_size, 16384)
+		var img_height = data_size / 16384 + 1
+		prints(data_size, img_width, img_height)
+
+		image = Image.create(img_width, img_height, false, Image.FORMAT_RGBAF)
+		var i = 0
+		for y in img_height:
+			for x in img_width:
+				if i > data_size - 1:
+					break
+				image.set_pixel(x, y, Color(
+					_data[i * 4],
+					_data[i * 4 + 1],
+					_data[i * 4 + 2],
+					_data[i * 4 + 3]
+				))
+				i += 1
+#		print((OS.get_ticks_msec() - time) / 1000.0)
+		image.save_png("d:/bvh.png")
+
+	func _sort_by_extents(a: int, b: int) -> bool:
+		return aabb.size[a] > aabb.size[b]
+
+	func _intersect_aabb(ray_origin: Vector3, ray_dir: Vector3, box_min: Vector3, box_max: Vector3, solid:=false) -> float:
+		var tMin : Vector3 = (box_min - ray_origin) / ray_dir
+		var tMax : Vector3 = (box_max - ray_origin) / ray_dir
+		var t1 : Vector3 = Vector3(min(tMin.x, tMax.x), min(tMin.y, tMax.y), min(tMin.z, tMax.z))
+		var t2 : Vector3 = Vector3(max(tMin.x, tMax.x), max(tMin.y, tMax.y), max(tMin.z, tMax.z))
+		var tNear : float = max(max(t1.x, t1.y), t1.z)
+		var tFar : float = min(min(t2.x, t2.y), t2.z)
+
+		if tNear > tFar or (tFar < 0.0 and tNear < 0.0):
+			return -1.0
+
+		if tNear < 0.0:
+			var temp := tNear
+			tNear = tFar
+			tFar = temp
+
+			if solid:
+				return 0.0
+
+		return tNear
diff --git a/addons/material_maker/engine/pipeline/compute_shader.gd b/addons/material_maker/engine/pipeline/compute_shader.gd
index ff1c98c9e..ced02dc77 100644
--- a/addons/material_maker/engine/pipeline/compute_shader.gd
+++ b/addons/material_maker/engine/pipeline/compute_shader.gd
@@ -183,6 +183,9 @@ func do_render(rd : RenderingDevice, output_tex : RID, size : Vector2i, rids : R
 	var uniform_set_1 : RID = RID()
 	if parameter_values.size() > 0:
 		uniform_set_1 = get_parameter_uniforms(rd, shader, rids)
+		if ! uniform_set_1.is_valid():
+			print("Failed to create valid uniform for parameters")
+			return false
 	
 	var uniform_set_2 = RID()
 	if !textures.is_empty():
diff --git a/addons/material_maker/engine/pipeline/pipeline.gd b/addons/material_maker/engine/pipeline/pipeline.gd
index daa204635..74051896f 100644
--- a/addons/material_maker/engine/pipeline/pipeline.gd
+++ b/addons/material_maker/engine/pipeline/pipeline.gd
@@ -116,7 +116,7 @@ func set_parameter(name : String, value) -> void:
 		p.value = value
 		match p.type:
 			"float":
-				if value is float:
+				if value is float or value is int:
 					parameter_values.encode_float(p.offset, value)
 					return
 			"int":
diff --git a/addons/material_maker/engine/pipeline/rendering_pipeline.gd b/addons/material_maker/engine/pipeline/rendering_pipeline.gd
index 7511839bb..62ff5306f 100644
--- a/addons/material_maker/engine/pipeline/rendering_pipeline.gd
+++ b/addons/material_maker/engine/pipeline/rendering_pipeline.gd
@@ -1,6 +1,7 @@
 extends MMPipeline
 class_name MMRenderingPipeline
 
+var shader : RID = RID()
 
 var index_count : int = 0
 var clearColors : PackedColorArray = PackedColorArray([Color.TRANSPARENT])
@@ -11,11 +12,6 @@ func create_framebuffer(rd : RenderingDevice, texture_rid : RID) -> RID:
 	
 	return framebuffer
 
-func compile_shader(rd : RenderingDevice, vertex_source : String, fragment_source : String, rids : RIDs, replaces : Dictionary = {}) -> RID:
-	var shader = do_compile_shader(rd, { vertex=vertex_source, fragment=fragment_source }, replaces)
-	rids.add(shader)
-	return shader
-
 func bind_buffer_uniforms(rd : RenderingDevice, draw_list : int, shader : RID, buffers : Array[PackedByteArray], set : int, rids : RIDs):
 	var uniform_set : RID = rd.uniform_set_create(create_buffers_uniform_list(rd, buffers, rids), shader, set)
 	rids.add(uniform_set)
@@ -23,13 +19,17 @@ func bind_buffer_uniforms(rd : RenderingDevice, draw_list : int, shader : RID, b
 
 func draw_list_extra_setup(rd : RenderingDevice, draw_list : int, shader : RID, rids : RIDs):
 	pass
-	
-func render(vertex_shader : String, fragment_shader : String, size : Vector2i, texture_type : int, target_texture : MMTexture, replaces : Dictionary = {}):
+
+func set_shader(vertex_source : String, fragment_source : String, replaces : Dictionary = {}):
+	replaces["@DECLARATIONS"] = get_uniform_declarations()+"\n"+get_texture_declarations()
+	var rd : RenderingDevice = await mm_renderer.request_rendering_device(self)
+	shader = do_compile_shader(rd, { vertex=vertex_source, fragment=fragment_source }, replaces)
+	mm_renderer.release_rendering_device(self)
+
+func render(size : Vector2i, texture_type : int, target_texture : MMTexture):
 	var rd : RenderingDevice = await mm_renderer.request_rendering_device(self)
 	var rids : RIDs = RIDs.new()
 	
-	var shader = compile_shader(rd, vertex_shader, fragment_shader, rids)
-	
 	var target_texture_id : RID = create_output_texture(rd, size, texture_type, true)
 	var framebuffer : RID = create_framebuffer(rd, target_texture_id)
 	rids.add(framebuffer)
@@ -53,6 +53,25 @@ func render(vertex_shader : String, fragment_shader : String, size : Vector2i, t
 		clearColors)
 	rd.draw_list_bind_render_pipeline(draw_list, pipeline)
 	
+	var uniform_set_1 : RID = RID()
+	if parameter_values.size() > 0:
+		uniform_set_1 = get_parameter_uniforms(rd, shader, rids)
+		if ! uniform_set_1.is_valid():
+			print("Failed to create valid uniform for parameters")
+			return false
+	
+	var uniform_set_2 = RID()
+	if !textures.is_empty():
+		uniform_set_2 = get_texture_uniforms(rd, shader, rids)
+		if ! uniform_set_2.is_valid():
+			print("Failed to create valid uniform for textures")
+			return false
+	
+	if uniform_set_1.is_valid():
+		rd.draw_list_bind_uniform_set(draw_list, uniform_set_1, 1)
+	if uniform_set_2.is_valid():
+		rd.draw_list_bind_uniform_set(draw_list, uniform_set_2, 2)
+	
 	draw_list_extra_setup(rd, draw_list, shader, rids)
 	
 	rd.draw_list_draw(draw_list, index_count == 0, 1, index_count)
diff --git a/addons/material_maker/engine/texture.gd b/addons/material_maker/engine/pipeline/texture.gd
similarity index 100%
rename from addons/material_maker/engine/texture.gd
rename to addons/material_maker/engine/pipeline/texture.gd
diff --git a/addons/material_maker/map_renderer/ao_fragment.tres b/addons/material_maker/map_renderer/ao_fragment.tres
new file mode 100644
index 000000000..94111f51e
--- /dev/null
+++ b/addons/material_maker/map_renderer/ao_fragment.tres
@@ -0,0 +1,263 @@
+[gd_resource type="Resource" script_class="TextResource" load_steps=2 format=3 uid="uid://cuveny5onbhr2"]
+
+[ext_resource type="Script" path="res://addons/material_maker/engine/text_resource.gd" id="1_oqegj"]
+
+[resource]
+script = ExtResource("1_oqegj")
+text = "#version 450
+
+layout(location = 0) out vec4 outColor;
+layout(location = 1) in vec3 model_vert;
+layout(location = 2) in vec3 normal;
+layout(location = 3) in vec2 uv;
+
+@DECLARATIONS
+
+const float PI = 3.1415928;
+
+float intersect_aabb(vec3 ray_origin, vec3 ray_dir, vec3 box_min, vec3 box_max, bool solid) {
+	vec3 tMin = (box_min - ray_origin) / ray_dir;
+	vec3 tMax = (box_max - ray_origin) / ray_dir;
+	vec3 t1 = min(tMin, tMax);
+	vec3 t2 = max(tMin, tMax);
+	float tNear = max(max(t1.x, t1.y), t1.z);
+	float tFar = min(min(t2.x, t2.y), t2.z);
+
+	if(tNear > tFar || (tFar < 0.0 && tNear < 0.0)) {
+		return -1.0;
+	}
+	if(tNear < 0.0) {
+		float temp = tNear;
+		tNear = tFar;
+		tFar = temp;
+		tNear *= float(!solid);
+	}
+
+	return tNear;
+}
+
+float intersects_triangle(vec3 ray_start, vec3 ray_dir, vec3 v0, vec3 v1, vec3 v2) {
+	vec3 e1 = v1 - v0;
+	vec3 e2 = v2 - v0;
+	vec3 h = cross(ray_dir, e2);
+	float a = dot(e1, h);
+	if (abs(a) < 0.000001) { // Parallel test.
+		return -1.0;
+	}
+
+	float f = 1.0 / a;
+
+	vec3 s = ray_start - v0;
+	float u = f * dot(s, h);
+
+	if (u < 0.0 || u > 1.0) {
+		return -1.0;
+	}
+
+	vec3 q = cross(s, e1);
+	float v = f * dot(ray_dir, q);
+
+	if (v < 0.0 || u + v > 1.0) {
+		return -1.0;
+	}
+
+	// At this stage we can compute t to find out where
+	// the intersection point is on the line.
+	float t = f * dot(e2, q);
+	return t > 0.0000001 ? t : -1.0;
+}
+
+vec4 get_data(int index) {
+	ivec2 data_size = textureSize(bvh_data, 0);
+	return texelFetch(bvh_data, ivec2(
+		index % data_size.x,
+		index / data_size.x
+	), 0);
+}
+
+float intersects_bvh(vec3 ray_start, vec3 ray_dir, out vec3 normal_hit) {
+	int offset_to_nodes = int(get_data(0)[0]);
+	vec4 root_data_0 = get_data(int(get_data(1)[0]) + offset_to_nodes);
+	vec4 root_data_1 = get_data(int(get_data(1)[0]) + offset_to_nodes + 1);
+
+	float t = intersect_aabb(ray_start, ray_dir, root_data_0.xyz, root_data_1.xyz, false);
+	if(t == -1.0) {
+		return 65536.0;
+	}
+
+	float prev_hit = t;
+	t = 65536.0; // Set to large number
+	vec3 tri[3];// = vec3[];
+	int min_node_idx = 0;
+
+	int stack_point = 0;
+	ivec3 node_stack[128];// = ivec3[]; // ivec3
+
+	int curr_node_idx = 0;
+	bool moving_up = false;
+
+	for(int i = 0; i < 256; i++) {
+		if(moving_up && stack_point <= 0) {
+			break;
+		}
+		int node_data_off = int(get_data(1 + curr_node_idx)[0]) + offset_to_nodes;
+		vec4 node_data_0 = get_data(node_data_off);
+		vec4 node_data_1 = get_data(node_data_off + 1);
+		int level = int(node_data_0[3]);
+
+		if(!moving_up) { // Moving down node hierarchy
+			if(node_data_1[3] > 0.0) { // Is a leaf node
+				for(int j = node_data_off + 2; j < node_data_off + 2 + int(node_data_1[3]) * 3; j+=3) {
+					vec3 tri_a = get_data(j).xyz;
+					vec3 tri_b = get_data(j + 1).xyz;
+					vec3 tri_c = get_data(j + 2).xyz;
+					float tri_t = intersects_triangle(ray_start, ray_dir,
+						tri_a, tri_b, tri_c
+					);
+					if(tri_t != -1.0) {
+						if(tri_t < t) {
+							tri[0] = tri_a;
+							tri[1] = tri_b;
+							tri[2] = tri_c;
+//							print(curr_node_idx)
+							min_node_idx = curr_node_idx;
+						}
+						t = min(t, tri_t);
+					}
+				}
+
+				stack_point -= 1;
+				if(stack_point <= 0) {
+					break;
+				}
+				if(node_stack[stack_point][1] == level) { // next node in stack is sibling
+					if(t < intBitsToFloat(node_stack[stack_point][0])) { // no chance to get better hit from sibling
+						stack_point -= 1;
+						moving_up = true;
+					}
+				} else {
+					moving_up = true;
+				}
+				prev_hit = intBitsToFloat(node_stack[stack_point][0]);
+				curr_node_idx = node_stack[stack_point][2];
+			} else {
+
+				// Push self onto stack
+				node_stack[stack_point] = ivec3(floatBitsToInt(prev_hit), level, curr_node_idx);
+				stack_point += 1;
+
+				ivec2 child_indices = ivec2(get_data(node_data_off + 2).xy);
+				int left_data_off = int(get_data(1 + child_indices[0])[0]) + offset_to_nodes;
+				vec4 left_data_0 = get_data(left_data_off);
+				vec4 left_data_1 = get_data(left_data_off + 1);
+				int right_data_off = int(get_data(1 + child_indices[1])[0]) + offset_to_nodes;
+				vec4 right_data_0 = get_data(right_data_off);
+				vec4 right_data_1 = get_data(right_data_off + 1);
+
+				float t_left = intersect_aabb(ray_start, ray_dir, left_data_0.xyz, left_data_1.xyz, true);
+				float t_right = intersect_aabb(ray_start, ray_dir, right_data_0.xyz, right_data_1.xyz, true);
+
+				if(t_right == -1.0 && t_left != -1.0) { // only left node hit
+					prev_hit = t_left;
+					curr_node_idx = child_indices[0];
+				} else if(t_left == -1.0 && t_right != -1.0) { // only right node hit
+					prev_hit = t_right;
+					curr_node_idx = child_indices[1];
+				} else if(t_left < t_right && t_left != -1.0) { // left node hits closer
+					node_stack[stack_point] = ivec3(floatBitsToInt(t_right), int(right_data_0[3]), child_indices[1]);
+					stack_point += 1;
+					prev_hit = t_left;
+					curr_node_idx = child_indices[0];
+				} else if(t_right <= t_left && t_right != -1.0) { // right node hits closer
+					node_stack[stack_point] = ivec3(floatBitsToInt(t_left), int(left_data_0[3]), child_indices[0]);
+					stack_point += 1;
+					prev_hit = t_right;
+					curr_node_idx = child_indices[1];
+				} else { // no hit
+					stack_point -= 2;
+					if(stack_point <= 0) {
+						break;
+					}
+					if(node_stack[stack_point][1] == level) { // next node in stack is sibling
+						if(t < intBitsToFloat(node_stack[stack_point][0])) { // no chance to get better hit from sibling
+							stack_point -= 1;
+							moving_up = true;
+						}
+					} else {
+						moving_up = true;
+					}
+					prev_hit = intBitsToFloat(node_stack[max(stack_point, 0)][0]);
+					curr_node_idx = node_stack[max(stack_point, 0)][2];
+				}
+			}
+		} else { // Moving up hierarchy
+			stack_point -= 1;
+			if(stack_point <= 0) {
+				break;
+			}
+			if(node_stack[stack_point][1] == level) { // next node in stack is sibling
+				if(t < intBitsToFloat(node_stack[stack_point][0])) { // no chance to get better hit from sibling
+					stack_point -= 1;
+				} else {
+					moving_up = false;
+				}
+			}
+			prev_hit = intBitsToFloat(node_stack[max(stack_point, 0)][0]);
+			curr_node_idx = node_stack[max(stack_point, 0)][2];
+		}
+	}
+
+	normal_hit = normalize(cross(tri[2] - tri[0], tri[1] - tri[0]));
+	return t;
+}
+
+vec3 random_hemi_point(vec3 rand, vec3 norm) {
+	float ang1 = (rand.x * 2.0) * PI; // [0..1) -> [0..2*PI)
+	float u = rand.y * 2.0 - 1.0; // [0..1), cos and acos(2v-1) cancel each other out, so we arrive at [-1..1)
+	float u2 = u * u;
+	float sqrt1MinusU2 = sqrt(1.0 - u2);
+	float x = sqrt1MinusU2 * cos(ang1);
+	float y = sqrt1MinusU2 * sin(ang1);
+	float z = u;
+	vec3 v = vec3(x, y, z);
+
+	return v * sign(dot(v, norm));
+}
+
+float rand(vec2 co){
+    return fract(sin(dot(co.xy ,vec2(12.9898,78.233))) * 43758.5453);
+}
+
+void main() {
+	vec3 model_norm = normalize(normal) * sign(max_dist);
+/*
+	vec3 ddx_vert = dFdx(model_vert);
+	vec3 ddy_vert = dFdy(model_vert);
+	vec3 true_normal = normalize(cross(ddy_vert, ddx_vert)) * sign(max_dist);
+*/
+	vec3 true_normal = model_norm;
+
+	vec3 ray_start = model_vert + model_norm * bias_dist;
+
+	vec3 ray_dir;
+	for(int i = 0; i < 5; i++) {
+		ray_dir = random_hemi_point(vec3(
+			rand(uv + float(iteration)),
+			rand(-uv.yx + float(iteration)),
+			0.0
+		), model_norm);
+		if(dot(ray_dir, true_normal) > 0.0) {
+			break;
+		}
+	}
+	vec3 normal_hit;
+	float hit = intersects_bvh(ray_start, ray_dir, normal_hit);
+	if(hit == 65536.0 || hit < 0.0  || hit > max_dist) {
+		hit = 1.0;
+	} else{
+		hit = 1.0 - clamp(dot(model_norm, ray_dir),0.0,1.0);
+	}
+
+	outColor = vec4(mix(texture(prev_iteration_tex, uv).rgb, vec3(hit), 1.0/float(iteration)), 1.0);
+}
+"
diff --git a/addons/material_maker/map_renderer/ao_vertex.tres b/addons/material_maker/map_renderer/ao_vertex.tres
new file mode 100644
index 000000000..5036c8820
--- /dev/null
+++ b/addons/material_maker/map_renderer/ao_vertex.tres
@@ -0,0 +1,35 @@
+[gd_resource type="Resource" script_class="TextResource" load_steps=2 format=3 uid="uid://cabcobt82hrux"]
+
+[ext_resource type="Script" path="res://addons/material_maker/engine/text_resource.gd" id="1_0bvqj"]
+
+[resource]
+script = ExtResource("1_0bvqj")
+text = "#version 450
+
+layout(location = 1) out vec3 model_vert;
+layout(location = 2) out vec3 normal;
+layout(location = 3) out vec2 uv;
+
+layout(binding = 0, std430) buffer restrict readonly Positions {
+	float positions[];
+};
+layout(binding = 1, std430) buffer restrict readonly AABB {
+	float aabb[];
+};
+layout(binding = 2, std430) buffer restrict readonly Normals {
+	float normals[];
+};
+layout(binding = 3, std430) buffer restrict readonly Tangents {
+	float tangents[];
+};
+layout(binding = 4, std430) buffer restrict readonly UVs {
+	vec2 uvs[];
+};
+
+void main() {
+	model_vert = vec3(positions[3*gl_VertexIndex], positions[3*gl_VertexIndex+1], positions[3*gl_VertexIndex+2]);
+	normal = vec3(normals[3*gl_VertexIndex], normals[3*gl_VertexIndex+1], normals[3*gl_VertexIndex+2]);
+	uv = uvs[gl_VertexIndex];
+	gl_Position = vec4(uv*2.0-vec2(1.0), 0.0, 1.0);
+}
+"
diff --git a/addons/material_maker/map_renderer/denoise_compute.tres b/addons/material_maker/map_renderer/denoise_compute.tres
new file mode 100644
index 000000000..56764d011
--- /dev/null
+++ b/addons/material_maker/map_renderer/denoise_compute.tres
@@ -0,0 +1,34 @@
+[gd_resource type="Resource" script_class="TextResource" load_steps=2 format=3 uid="uid://dyduptoh12rkt"]
+
+[ext_resource type="Script" path="res://addons/material_maker/engine/text_resource.gd" id="1_puend"]
+
+[resource]
+script = ExtResource("1_puend")
+text = "#version 450
+
+layout(local_size_x = @LOCAL_SIZE, local_size_y = 1, local_size_z = 1) in;
+
+layout(set = 0, binding = 0, @OUT_TEXTURE_TYPE) uniform image2D OUTPUT_TEXTURE;
+
+@DECLARATIONS
+
+void main() {
+	vec2 pixel = gl_GlobalInvocationID.xy;
+	vec2 image_size = imageSize(OUTPUT_TEXTURE);
+	vec2 uv = pixel/image_size;
+
+	if (texture(tex, uv).a == 0.0) {
+		return;
+	}
+
+	vec3 color = vec3(0.0);
+	float weight = 0.0;
+	for(int y = -radius; y < radius+1; y++) {
+		for(int x = -radius; x < radius+1; x++) {
+			vec4 col = texture(tex, uv + vec2(ivec2(x, y)) / image_size, 0);
+			color += col.rgb * col.a;
+			weight += col.a;
+		}
+	}
+	imageStore(OUTPUT_TEXTURE, ivec2(gl_GlobalInvocationID.xy), vec4(color / weight, 1.0));
+}"
diff --git a/addons/material_maker/map_renderer/map_generator.gd b/addons/material_maker/map_renderer/map_generator.gd
index 30521ec56..c137d9414 100644
--- a/addons/material_maker/map_renderer/map_generator.gd
+++ b/addons/material_maker/map_renderer/map_generator.gd
@@ -1,25 +1,50 @@
 extends Object
 class_name MMMapGenerator
 
+
+const SHADERS_PATH : String = "res://addons/material_maker/map_renderer"
+const SHADERS : Dictionary = {
+	position = { vertex = "position_vertex", fragment = "common_fragment" },
+	normal = { vertex = "normal_vertex", fragment = "normal_fragment" },
+	tangent = { vertex = "tangent_vertex", fragment = "normal_fragment" },
+	ambient_occlusion = { vertex = "ao_vertex", fragment = "ao_fragment" }
+}
+
+
 static func generate(mesh : Mesh, map : String, size : int, texture : MMTexture):
-	var pixels : int = size/2
-	
+	var pixels : int = size/4
 	var mesh_pipeline : MMMeshRenderingPipeline = MMMeshRenderingPipeline.new()
 	mesh_pipeline.mesh = mesh
-	var vertex_shader : String
-	var fragment_shader : String
+	var shaders : Dictionary = SHADERS[map]
+	var vertex_shader : String = load(SHADERS_PATH+"/"+shaders.vertex+".tres").text
+	var fragment_shader : String = load(SHADERS_PATH+"/"+shaders.fragment+".tres").text
 	match map:
-		"position":
-			vertex_shader = load("res://addons/material_maker/map_renderer/position_map_vertex.tres").text
-			fragment_shader = load("res://addons/material_maker/map_renderer/common_fragment.tres").text
-		"normal":
-			vertex_shader = load("res://addons/material_maker/map_renderer/normal_map_vertex.tres").text
-			fragment_shader = load("res://addons/material_maker/map_renderer/normal_fragment.tres").text
-		"tangent":
-			vertex_shader = load("res://addons/material_maker/map_renderer/tangent_map_vertex.tres").text
-			fragment_shader = load("res://addons/material_maker/map_renderer/normal_fragment.tres").text
-	await mesh_pipeline.render(vertex_shader, fragment_shader, Vector2i(size, size), 3, texture)
-	
+		"position", "normal", "tangent":
+			await mesh_pipeline.set_shader(vertex_shader, fragment_shader)
+			await mesh_pipeline.render(Vector2i(size, size), 3, texture)
+		"ambient_occlusion":
+			var bvh : MMTexture = MMTexture.new()
+			bvh.set_texture(MMBvhGenerator.generate(mesh))
+			mesh_pipeline.add_parameter_or_texture("bvh_data", "sampler2D", bvh)
+			mesh_pipeline.add_parameter_or_texture("prev_iteration_tex", "sampler2D", texture)
+			mesh_pipeline.add_parameter_or_texture("max_dist", "float", 50)
+			mesh_pipeline.add_parameter_or_texture("bias_dist", "float", 0.1)
+			mesh_pipeline.add_parameter_or_texture("iteration", "int", 1)
+			await mesh_pipeline.set_shader(vertex_shader, fragment_shader)
+			for i in range(128):
+				mesh_pipeline.set_parameter("iteration", i+1)
+				mesh_pipeline.set_parameter("prev_iteration_tex", texture)
+				await mesh_pipeline.render(Vector2i(size, size), 3, texture)
+			
+			# Denoise
+			var compute_pipeline : MMComputeShader = MMComputeShader.new()
+			compute_pipeline.clear()
+			compute_pipeline.add_parameter_or_texture("tex", "sampler2D", texture)
+			compute_pipeline.add_parameter_or_texture("radius", "int", 3)
+			await compute_pipeline.set_shader(load("res://addons/material_maker/map_renderer/denoise_compute.tres").text, 3)
+			await compute_pipeline.render(texture, size)
+
+	# Extend the map past seams
 	if pixels > 0:
 		var compute_pipeline : MMComputeShader = MMComputeShader.new()
 		compute_pipeline.clear()
diff --git a/addons/material_maker/map_renderer/normal_map_vertex.tres b/addons/material_maker/map_renderer/normal_vertex.tres
similarity index 100%
rename from addons/material_maker/map_renderer/normal_map_vertex.tres
rename to addons/material_maker/map_renderer/normal_vertex.tres
diff --git a/addons/material_maker/map_renderer/position_map_vertex.tres b/addons/material_maker/map_renderer/position_vertex.tres
similarity index 100%
rename from addons/material_maker/map_renderer/position_map_vertex.tres
rename to addons/material_maker/map_renderer/position_vertex.tres
diff --git a/addons/material_maker/map_renderer/tangent_map_vertex.tres b/addons/material_maker/map_renderer/tangent_vertex.tres
similarity index 100%
rename from addons/material_maker/map_renderer/tangent_map_vertex.tres
rename to addons/material_maker/map_renderer/tangent_vertex.tres
diff --git a/material_maker/tools/map_renderer/map_renderer.gd b/material_maker/tools/map_renderer/map_renderer.gd
index 6caec2821..3aab2710f 100644
--- a/material_maker/tools/map_renderer/map_renderer.gd
+++ b/material_maker/tools/map_renderer/map_renderer.gd
@@ -17,7 +17,6 @@ extends Node
 @export var seams_pass2: ShaderMaterial
 
 func gen_new(mesh: Mesh, map : String, renderer_method : String, arguments : Array, map_size = 512) -> void:
-	var map_generator : MMMeshRenderingPipeline = MMMeshRenderingPipeline.new()
 	var texture : MMTexture = MMTexture.new()
 	await MMMapGenerator.generate(mesh, map, map_size, texture)
 	match renderer_method:
@@ -69,17 +68,20 @@ func gen(mesh: Mesh, map : String, renderer_method : String, arguments : Array,
 		passes.first.set_shader_parameter("bvh_data", bvh_data)
 		passes.first.set_shader_parameter("max_dist", ray_distance)
 		passes.first.set_shader_parameter("bias_dist", ao_ray_bias)
+		ray_count = 1
 		for i in ray_count:
 			progress_dialog.set_progress(float(i)/ray_count)
 			passes.first.set_shader_parameter("iteration", i+1)
 			viewport.render_target_update_mode = SubViewport.UPDATE_ONCE
 			await get_tree().process_frame
-		mesh_instance.set_surface_override_material(0, denoise_pass)
-		denoise_pass.set_shader_parameter("size", map_size)
-		denoise_pass.set_shader_parameter("radius", denoise_radius)
-		viewport.render_target_update_mode = SubViewport.UPDATE_ONCE
-		await get_tree().process_frame
-		await get_tree().process_frame
+			await get_tree().process_frame
+		if false:
+			mesh_instance.set_surface_override_material(0, denoise_pass)
+			denoise_pass.set_shader_parameter("size", map_size)
+			denoise_pass.set_shader_parameter("radius", denoise_radius)
+			viewport.render_target_update_mode = SubViewport.UPDATE_ONCE
+			await get_tree().process_frame
+			await get_tree().process_frame
 		progress_dialog.queue_free()
 	else:
 		passes.first.set_shader_parameter("position", aabb.position)
diff --git a/material_maker/tools/map_renderer/map_renderer.tscn b/material_maker/tools/map_renderer/map_renderer.tscn
index 7e4c02a10..3770a6f12 100644
--- a/material_maker/tools/map_renderer/map_renderer.tscn
+++ b/material_maker/tools/map_renderer/map_renderer.tscn
@@ -1,8 +1,8 @@
-[gd_scene load_steps=32 format=3 uid="uid://cv706ip72t48u"]
+[gd_scene load_steps=33 format=3 uid="uid://cv706ip72t48u"]
 
 [ext_resource type="Script" path="res://material_maker/tools/map_renderer/curvature_generator.gd" id="1"]
 [ext_resource type="Script" path="res://material_maker/tools/map_renderer/map_renderer.gd" id="2"]
-[ext_resource type="Script" path="res://material_maker/tools/map_renderer/bvh_generator.gd" id="3"]
+[ext_resource type="Script" path="res://addons/material_maker/bvh_generator.gd" id="3"]
 
 [sub_resource type="Shader" id="4"]
 code = "shader_type spatial;
@@ -122,13 +122,13 @@ uniform int iteration = 1;
 
 uniform sampler2D prev_iteration_tex;
 
-varying vec3 normal;
 varying vec3 model_vert;
+varying vec3 normal;
 
 void vertex() {
 	model_vert = VERTEX;
 	normal = NORMAL;
-	VERTEX = vec3(UV, 0.5);
+	VERTEX = vec3(UV.x, 1.0-UV.y, 0.5);
 }
 
 float intersect_aabb(vec3 ray_origin, vec3 ray_dir, vec3 box_min, vec3 box_max, bool solid) {
@@ -375,7 +375,7 @@ void fragment() {
 "
 
 [sub_resource type="ViewportTexture" id="11"]
-viewport_path = NodePath(".")
+viewport_path = NodePath("Viewport")
 
 [sub_resource type="ShaderMaterial" id="ShaderMaterial_0j43v"]
 resource_local_to_scene = true
@@ -649,7 +649,7 @@ void fragment() {
 }"
 
 [sub_resource type="ViewportTexture" id="27"]
-viewport_path = NodePath(".")
+viewport_path = NodePath("Viewport")
 
 [sub_resource type="ShaderMaterial" id="ShaderMaterial_esyq5"]
 resource_local_to_scene = true
@@ -689,13 +689,16 @@ void fragment() {
 	ALBEDO = color / weight;
 }"
 
+[sub_resource type="ViewportTexture" id="ViewportTexture_cyhhm"]
+viewport_path = NodePath("Viewport")
+
 [sub_resource type="ShaderMaterial" id="ShaderMaterial_nu2ml"]
 resource_local_to_scene = true
 render_priority = 0
 shader = SubResource("13")
 shader_parameter/radius = 0
 shader_parameter/size = 512.0
-shader_parameter/tex = SubResource("11")
+shader_parameter/tex = SubResource("ViewportTexture_cyhhm")
 
 [sub_resource type="Shader" id="15"]
 code = "shader_type canvas_item;