Nathan
411 lines
15 KiB
Python
411 lines
15 KiB
Python
import bpy
|
|
from bpy.types import Panel, Operator, PropertyGroup
|
|
from bpy.props import StringProperty, IntProperty, BoolProperty, EnumProperty, PointerProperty
|
|
from bpy.utils import register_class, unregister_class
|
|
|
|
# Module-level caches prevent enum item lists being garbage-collected mid-draw,
|
|
# which would crash Blender when dynamic EnumProperty items are used.
|
|
_socket_items_cache = [('0', 'Select a node', '', 0)]
|
|
_uv_items_cache = [('__active__', '(Active UV)', '', 0)]
|
|
|
|
|
|
def _get_socket_items(self, context):
|
|
global _socket_items_cache
|
|
|
|
node = getattr(context, 'active_node', None) if context else None
|
|
if node is None:
|
|
_socket_items_cache = [('0', 'Select a node', '', 0)]
|
|
return _socket_items_cache
|
|
|
|
items = []
|
|
for i, socket in enumerate(node.outputs):
|
|
if (socket.bl_idname not in ('NodeSocketShader', 'NodeSocketVirtual')
|
|
and not socket.hide
|
|
and socket.enabled):
|
|
items.append((str(i), socket.name, socket.name, len(items)))
|
|
|
|
if not items:
|
|
_socket_items_cache = [('0', 'No bakeable outputs', '', 0)]
|
|
return _socket_items_cache
|
|
|
|
_socket_items_cache = items
|
|
return _socket_items_cache
|
|
|
|
|
|
def _get_uv_items(self, context):
|
|
global _uv_items_cache
|
|
obj = context.active_object if context else None
|
|
items = [('__active__', '(Active UV)', 'Use whichever UV map is currently active on the object', 0)]
|
|
if obj and obj.type == 'MESH':
|
|
for i, layer in enumerate(obj.data.uv_layers):
|
|
items.append((layer.name, layer.name, layer.name, i + 1))
|
|
_uv_items_cache = items
|
|
return _uv_items_cache
|
|
|
|
|
|
class SimpleBakeNodeBakeProps(PropertyGroup):
|
|
socket: EnumProperty(
|
|
name="Socket",
|
|
description="Output socket to bake",
|
|
items=_get_socket_items,
|
|
)
|
|
image_name: StringProperty(
|
|
name="Image Name",
|
|
default="NodeBake",
|
|
description="Name for the baked image in the blend file",
|
|
)
|
|
width: IntProperty(name="Width", default=1024, min=1, max=16384)
|
|
height: IntProperty(name="Height", default=1024, min=1, max=16384)
|
|
float_buffer: BoolProperty(
|
|
name="32-bit Float",
|
|
default=False,
|
|
description="Use a 32-bit float image buffer (higher precision, larger memory footprint)",
|
|
)
|
|
uv_map: EnumProperty(
|
|
name="UV Map",
|
|
description="UV map to bake to. Defaults to whichever map is currently active on the object",
|
|
items=_get_uv_items,
|
|
)
|
|
replace_with_image: BoolProperty(
|
|
name="Replace with Image Texture",
|
|
default=False,
|
|
description=(
|
|
"After baking, insert an Image Texture node and reconnect it "
|
|
"to wherever this socket was feeding"
|
|
),
|
|
)
|
|
|
|
|
|
def _find_active_mat_output(nodes):
|
|
"""Return the active Material Output node, or any Material Output, or None."""
|
|
active = next((n for n in nodes if n.bl_idname == 'ShaderNodeOutputMaterial' and n.is_active_output), None)
|
|
if active:
|
|
return active
|
|
return next((n for n in nodes if n.bl_idname == 'ShaderNodeOutputMaterial'), None)
|
|
|
|
|
|
class SIMPLEBAKE_OT_NodeBake(Operator):
|
|
bl_idname = "simplebake.node_bake"
|
|
bl_label = "Bake Node Output"
|
|
bl_description = (
|
|
"Bake the selected output socket of the active node to an image texture. "
|
|
"The render engine is temporarily switched to Cycles for baking"
|
|
)
|
|
|
|
@classmethod
|
|
def poll(cls, context):
|
|
sd = getattr(context, 'space_data', None)
|
|
if sd is None or sd.type != 'NODE_EDITOR':
|
|
return False
|
|
if sd.tree_type != 'ShaderNodeTree':
|
|
return False
|
|
if context.active_object is None:
|
|
return False
|
|
node = context.active_node
|
|
if node is None:
|
|
return False
|
|
return any(
|
|
s.bl_idname not in ('NodeSocketShader', 'NodeSocketVirtual')
|
|
for s in node.outputs
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Material helpers
|
|
# ------------------------------------------------------------------
|
|
|
|
def _silence_material(self, mat):
|
|
"""
|
|
Temporarily wire a zero-emission shader into mat's Material Output so
|
|
it contributes nothing to an EMIT bake. Returns a restore tuple.
|
|
"""
|
|
nt = mat.node_tree
|
|
mat_out = _find_active_mat_output(nt.nodes)
|
|
if mat_out is None:
|
|
return None
|
|
|
|
orig = mat_out.inputs['Surface'].links[0].from_socket if mat_out.inputs['Surface'].is_linked else None
|
|
|
|
emit = nt.nodes.new('ShaderNodeEmission')
|
|
emit.inputs['Color'].default_value = (0.0, 0.0, 0.0, 1.0)
|
|
emit.inputs['Strength'].default_value = 0.0
|
|
nt.links.new(emit.outputs['Emission'], mat_out.inputs['Surface'])
|
|
|
|
return (nt, mat_out, emit, orig)
|
|
|
|
def _restore_material(self, restore_data):
|
|
if restore_data is None:
|
|
return
|
|
nt, mat_out, emit_node, orig_socket = restore_data
|
|
nt.nodes.remove(emit_node)
|
|
if orig_socket is not None:
|
|
nt.links.new(orig_socket, mat_out.inputs['Surface'])
|
|
|
|
# ------------------------------------------------------------------
|
|
# Operator execute
|
|
# ------------------------------------------------------------------
|
|
|
|
def execute(self, context):
|
|
props = context.scene.SimpleBakeNodeBake
|
|
space = context.space_data
|
|
node = context.active_node
|
|
node_location = (node.location[0], node.location[1])
|
|
|
|
# Must be at the root of the node tree, not inside a group
|
|
if space.edit_tree is not space.node_tree:
|
|
self.report({'ERROR'}, "Exit the node group before baking")
|
|
return {'CANCELLED'}
|
|
|
|
mat = space.id
|
|
if not isinstance(mat, bpy.types.Material):
|
|
self.report({'ERROR'}, "Could not determine the active material")
|
|
return {'CANCELLED'}
|
|
|
|
obj = context.active_object
|
|
mat_names = {s.material.name for s in obj.material_slots if s.material}
|
|
if mat.name not in mat_names:
|
|
self.report({'ERROR'}, f"'{mat.name}' is not assigned to the active object")
|
|
return {'CANCELLED'}
|
|
|
|
if not obj.data.uv_layers:
|
|
self.report({'ERROR'}, "The active object has no UV map — add one before baking")
|
|
return {'CANCELLED'}
|
|
|
|
# Activate the chosen UV map, remember original to restore later
|
|
chosen_uv = props.uv_map
|
|
orig_uv_index = obj.data.uv_layers.active_index
|
|
if chosen_uv != '__active__':
|
|
idx = next((i for i, l in enumerate(obj.data.uv_layers) if l.name == chosen_uv), None)
|
|
if idx is None:
|
|
self.report({'ERROR'}, f"UV map '{chosen_uv}' not found on object")
|
|
return {'CANCELLED'}
|
|
obj.data.uv_layers.active_index = idx
|
|
|
|
# Resolve the target socket
|
|
socket_idx = int(props.socket)
|
|
if socket_idx >= len(node.outputs):
|
|
self.report({'ERROR'}, "Socket index out of range — refresh by reselecting the node")
|
|
return {'CANCELLED'}
|
|
target_socket = node.outputs[socket_idx]
|
|
if target_socket.bl_idname in ('NodeSocketShader', 'NodeSocketVirtual'):
|
|
self.report({'ERROR'}, "Cannot bake a Shader or Virtual socket")
|
|
return {'CANCELLED'}
|
|
|
|
# Remember downstream links for the optional reconnect.
|
|
# Store names/identifiers rather than live Python wrapper references —
|
|
# those can become stale after node tree modifications during the bake.
|
|
downstream = [(lnk.to_node.name, lnk.to_socket.identifier) for lnk in target_socket.links]
|
|
|
|
# Create the output image
|
|
image_name = props.image_name.strip() or f"{node.name}_{target_socket.name}"
|
|
existing = bpy.data.images.get(image_name)
|
|
if existing:
|
|
bpy.data.images.remove(existing)
|
|
image = bpy.data.images.new(
|
|
image_name,
|
|
width=props.width,
|
|
height=props.height,
|
|
alpha=False,
|
|
float_buffer=props.float_buffer,
|
|
)
|
|
|
|
# --- Set up the target material ---
|
|
nt = mat.node_tree
|
|
nodes = nt.nodes
|
|
links = nt.links
|
|
|
|
mat_out = _find_active_mat_output(nodes)
|
|
created_mat_out = False
|
|
if mat_out is None:
|
|
mat_out = nodes.new('ShaderNodeOutputMaterial')
|
|
mat_out.location = (300.0, 300.0)
|
|
created_mat_out = True
|
|
|
|
orig_surface = mat_out.inputs['Surface'].links[0].from_socket if mat_out.inputs['Surface'].is_linked else None
|
|
|
|
# Temporary nodes
|
|
emit_node = nodes.new('ShaderNodeEmission')
|
|
emit_node.location = (mat_out.location[0] - 220, mat_out.location[1] - 80)
|
|
|
|
img_node = nodes.new('ShaderNodeTexImage')
|
|
img_node.name = "SB_NodeBake_Image"
|
|
img_node.image = image
|
|
img_node.location = (mat_out.location[0] - 420, mat_out.location[1] - 160)
|
|
|
|
# Make img_node the active image node (Cycles writes bake output here)
|
|
for n in nodes:
|
|
n.select = False
|
|
img_node.select = True
|
|
nodes.active = img_node
|
|
|
|
links.new(target_socket, emit_node.inputs['Color'])
|
|
links.new(emit_node.outputs['Emission'], mat_out.inputs['Surface'])
|
|
|
|
# --- Silence all other materials so they don't pollute the EMIT bake ---
|
|
silence_restores = []
|
|
for slot in obj.material_slots:
|
|
if slot.material and slot.material.name != mat.name:
|
|
rd = self._silence_material(slot.material)
|
|
if rd:
|
|
silence_restores.append(rd)
|
|
|
|
# --- Save scene state ---
|
|
orig_engine = context.scene.render.engine
|
|
orig_active_obj = context.view_layer.objects.active
|
|
orig_selected = list(context.selected_objects)
|
|
orig_mode = context.mode
|
|
|
|
# Configure for bake
|
|
context.scene.render.engine = 'CYCLES'
|
|
for o in list(context.selected_objects):
|
|
o.select_set(False)
|
|
obj.select_set(True)
|
|
context.view_layer.objects.active = obj
|
|
|
|
if orig_mode != 'OBJECT':
|
|
bpy.ops.object.mode_set(mode='OBJECT')
|
|
|
|
# --- Bake ---
|
|
success = False
|
|
try:
|
|
bpy.ops.object.bake(type='EMIT', use_clear=True)
|
|
success = True
|
|
except Exception as exc:
|
|
self.report({'ERROR'}, f"Bake failed: {exc}")
|
|
|
|
# --- Restore UV map ---
|
|
obj.data.uv_layers.active_index = orig_uv_index
|
|
|
|
# --- Restore scene state ---
|
|
context.scene.render.engine = orig_engine
|
|
context.view_layer.objects.active = orig_active_obj
|
|
for o in orig_selected:
|
|
try:
|
|
o.select_set(True)
|
|
except Exception:
|
|
pass
|
|
|
|
# --- Restore silenced materials ---
|
|
for rd in silence_restores:
|
|
self._restore_material(rd)
|
|
|
|
# --- Restore target material ---
|
|
nodes.remove(emit_node)
|
|
if created_mat_out:
|
|
nodes.remove(mat_out)
|
|
elif orig_surface is not None:
|
|
links.new(orig_surface, mat_out.inputs['Surface'])
|
|
|
|
# --- Handle image node ---
|
|
# Always remove the bake-helper image node — it was only needed for Cycles
|
|
# to know where to write the result. If replacing, a fresh node is created
|
|
# below at the correct position, avoiding any stale Python wrapper issues.
|
|
bake_helper = nodes.get("SB_NodeBake_Image") or img_node
|
|
nodes.remove(bake_helper)
|
|
|
|
if not success:
|
|
bpy.data.images.remove(image)
|
|
return {'CANCELLED'}
|
|
|
|
if props.replace_with_image:
|
|
final_node = nodes.new('ShaderNodeTexImage')
|
|
final_node.image = image
|
|
final_node.parent = node.parent
|
|
|
|
image_node_height = max(
|
|
float(getattr(final_node, "height", 0.0) or 0.0),
|
|
float(final_node.dimensions[1] or 0.0),
|
|
)
|
|
if image_node_height <= 0.0:
|
|
image_node_height = 260.0
|
|
final_node.location = (node_location[0], node_location[1] + image_node_height + 80.0)
|
|
|
|
if chosen_uv != '__active__':
|
|
uv_node = nodes.new('ShaderNodeUVMap')
|
|
uv_node.uv_map = chosen_uv
|
|
uv_node.parent = final_node.parent
|
|
uv_node.location = (final_node.location[0] - 220, final_node.location[1] - 100)
|
|
links.new(uv_node.outputs['UV'], final_node.inputs['Vector'])
|
|
|
|
for node_name, sock_id in downstream:
|
|
to_node = nodes.get(node_name)
|
|
if to_node is None:
|
|
continue
|
|
to_socket = next((s for s in to_node.inputs if s.identifier == sock_id), None)
|
|
if to_socket is not None:
|
|
links.new(final_node.outputs['Color'], to_socket)
|
|
|
|
self.report({'INFO'}, f"Baked '{target_socket.name}' → '{image_name}'")
|
|
return {'FINISHED'}
|
|
|
|
|
|
class SIMPLEBAKE_PT_NodeBake(Panel):
|
|
bl_label = "Node Bake"
|
|
bl_idname = "SIMPLEBAKE_PT_NodeBake"
|
|
bl_space_type = 'NODE_EDITOR'
|
|
bl_region_type = 'UI'
|
|
bl_category = "SimpleBake"
|
|
|
|
@classmethod
|
|
def poll(cls, context):
|
|
sd = context.space_data
|
|
return (
|
|
sd.type == 'NODE_EDITOR'
|
|
and sd.tree_type == 'ShaderNodeTree'
|
|
and context.active_object is not None
|
|
)
|
|
|
|
def draw(self, context):
|
|
layout = self.layout
|
|
props = context.scene.SimpleBakeNodeBake
|
|
node = context.active_node
|
|
|
|
if node is None:
|
|
layout.label(text="Select a node to bake", icon='INFO')
|
|
return
|
|
|
|
valid_outputs = [
|
|
s for s in node.outputs
|
|
if s.bl_idname not in ('NodeSocketShader', 'NodeSocketVirtual')
|
|
]
|
|
if not valid_outputs:
|
|
layout.label(text="No bakeable outputs on this node", icon='ERROR')
|
|
return
|
|
|
|
layout.label(text=node.name, icon='NODE')
|
|
|
|
layout.prop(props, "socket")
|
|
layout.prop(props, "uv_map")
|
|
layout.prop(props, "image_name")
|
|
|
|
row = layout.row(align=True)
|
|
row.prop(props, "width", text="W")
|
|
row.prop(props, "height", text="H")
|
|
|
|
layout.prop(props, "float_buffer")
|
|
layout.prop(props, "replace_with_image")
|
|
|
|
layout.separator()
|
|
|
|
row = layout.row()
|
|
row.scale_y = 1.4
|
|
row.operator("simplebake.node_bake", text="Bake Node Output", icon='RENDER_STILL')
|
|
|
|
|
|
classes = [
|
|
SimpleBakeNodeBakeProps,
|
|
SIMPLEBAKE_OT_NodeBake,
|
|
SIMPLEBAKE_PT_NodeBake,
|
|
]
|
|
|
|
|
|
def register():
|
|
for cls in classes:
|
|
register_class(cls)
|
|
bpy.types.Scene.SimpleBakeNodeBake = PointerProperty(type=SimpleBakeNodeBakeProps)
|
|
|
|
|
|
def unregister():
|
|
del bpy.types.Scene.SimpleBakeNodeBake
|
|
for cls in reversed(classes):
|
|
unregister_class(cls)
|