Files
ComfyProjector/utils.py
小煜 05d8702373 feat(渲染): 新增线稿渲染功能
添加线稿渲染功能,用于生成 ControlNet Canny 所需的黑白线稿图。实现通过 Freestyle 渲染引擎,在白色背景上生成黑色线条,并临时替换材质与世界设置以确保输出符合要求。同时更新了清理函数以移除新增的临时白色材质。
2026-02-13 16:51:46 +08:00

708 lines
26 KiB
Python

import bpy
import bpy.utils.previews
from mathutils import Vector
import os
import subprocess
# Global preview collection for rendered image thumbnails
_preview_collections = {}
def _get_preview_collection():
"""Get or create the preview collection for rendered images."""
if 'rendered_images' not in _preview_collections:
_preview_collections['rendered_images'] = bpy.utils.previews.new()
return _preview_collections['rendered_images']
def _load_preview(filepath):
"""Load an image file as a preview icon, returns the preview icon_id."""
pcoll = _get_preview_collection()
name = os.path.basename(filepath)
if name in pcoll:
# Force reload
del pcoll[name]
preview = pcoll.load(name, filepath, 'IMAGE')
return preview.icon_id
class Comfy_RenderedImageItem(bpy.types.PropertyGroup):
"""A single rendered image entry."""
name: bpy.props.StringProperty(name="Name")
filepath: bpy.props.StringProperty(name="File Path", subtype='FILE_PATH')
icon_id: bpy.props.IntProperty(name="Icon ID", default=0)
def property_exists(prop_path, glob, loc):
try:
eval(prop_path, glob, loc)
return True
except: # noqa: E722
return False
class Comfy_OT_Aovbake_Ad980(bpy.types.Operator):
bl_idname = "comfyui.aovbake_ad980"
bl_label = "aovbake"
bl_description = ""
bl_options = {"REGISTER", "UNDO"}
@classmethod
def poll(cls, _context):
if bpy.app.version >= (3, 0, 0) and True:
cls.poll_message_set(
"Please save your Blender file before running the script."
)
return not ("" == bpy.data.filepath)
def execute(self, _context):
if "" == bpy.data.filepath:
self.report(
{"ERROR"},
message="Please save your Blender file before running the script.",
)
else:
if bpy.context.view_layer.objects.active.type == "MESH":
bpy.data.scenes["Scene"].render.film_transparent = True
bpy.data.scenes["Scene"].render.engine = "CYCLES"
bpy.data.scenes["Scene"].view_layers["ViewLayer"].use_pass_z = True
bpy.data.scenes["Scene"].cycles.samples = 16
res = int(bpy.context.scene.sna_render_resolution)
bpy.data.scenes["Scene"].render.resolution_x = res
bpy.data.scenes["Scene"].render.resolution_y = res
bpy.data.scenes["Scene"].view_layers["ViewLayer"].use_pass_normal = True
bpy.context.scene.sna_activeobject = (
bpy.context.view_layer.objects.active
)
bpy.data.scenes["Scene"].use_nodes = True
import math
def show_message_box(message="", title="Message", icon="INFO"):
"""Shows a pop-up message in Blender."""
def draw(self, context):
self.layout.label(text=message)
bpy.context.window_manager.popup_menu(draw, title=title, icon=icon)
def create_camera_for_object(obj):
if obj.type != "MESH":
show_message_box("Selected object must be a mesh.", "Error")
return None
bpy.ops.object.camera_add()
cam = bpy.context.object
cam.name = "ObjectCamera"
bpy.context.scene.camera = cam
bpy.context.scene.sna_bakecam = bpy.context.scene.camera
bbox_corners = [
obj.matrix_world @ Vector(corner) for corner in obj.bound_box
]
min_x = min(corner.x for corner in bbox_corners)
max_x = max(corner.x for corner in bbox_corners)
min_y = min(corner.y for corner in bbox_corners)
max_y = max(corner.y for corner in bbox_corners)
min_z = min(corner.z for corner in bbox_corners)
max_z = max(corner.z for corner in bbox_corners)
center = (
Vector((min_x + max_x, min_y + max_y, min_z + max_z)) / 2
).to_3d()
size = max(max_x - min_x, max_y - min_y, max_z - min_z) * 1.5
cam.location = center + Vector((0, -size, 0))
cam.data.type = "ORTHO"
cam.data.ortho_scale = size
cam.rotation_euler = (math.radians(90), 0, 0)
bpy.context.scene.sna_activeobject.select_set(state=True)
bpy.ops.view3d.camera_to_view_selected()
return cam
def main():
obj = bpy.context.active_object
if obj is None:
show_message_box("No active object selected.", "Error")
return
create_camera_for_object(obj)
main()
bpy.context.scene.sna_activeobject.select_set(
state=True,
)
bpy.context.view_layer.objects.active = (
bpy.context.scene.sna_activeobject
)
bpy.data.scenes["Scene"].use_nodes = False
return {"FINISHED"}
def invoke(self, context, _event):
return self.execute(context)
def _ensure_output_dir():
"""Ensure output directory exists next to the blend file and return its path."""
blend_dir = os.path.dirname(bpy.data.filepath)
output_dir = os.path.join(blend_dir, 'output')
os.makedirs(output_dir, exist_ok=True)
return output_dir
def _append_from_materials_blend(data_type, name):
"""Append a data-block from Materials.blend if it doesn't already exist.
data_type: 'NodeTree' or 'Material'
"""
collection_map = {
'NodeTree': bpy.data.node_groups,
'Material': bpy.data.materials,
}
collection = collection_map[data_type]
if name in collection:
return collection[name]
blend_path = os.path.join(os.path.dirname(__file__), 'assets', 'Materials.blend')
bpy.ops.wm.append(
directory=blend_path + '\\' + data_type + '\\',
filename=name,
link=False,
)
return collection.get(name)
def _add_rendered_image(scene, name, filepath):
"""Add a rendered image to the scene's rendered image list."""
# Remove existing entry with the same name
for i, item in enumerate(scene.sna_rendered_images):
if item.name == name:
scene.sna_rendered_images.remove(i)
break
entry = scene.sna_rendered_images.add()
entry.name = name
entry.filepath = filepath
if os.path.isfile(filepath):
entry.icon_id = _load_preview(filepath)
# Set the active index to the new item
scene.sna_rendered_images_index = len(scene.sna_rendered_images) - 1
def _create_depth_material(near, far):
"""Create a material that renders depth as grayscale (near=white, far=black)."""
mat_name = '__ComfyDepthMat__'
if mat_name in bpy.data.materials:
bpy.data.materials.remove(bpy.data.materials[mat_name])
mat = bpy.data.materials.new(name=mat_name)
mat.use_nodes = True
nodes = mat.node_tree.nodes
links = mat.node_tree.links
nodes.clear()
# Camera Data → View Z Depth
cam_data = nodes.new('ShaderNodeCameraData')
cam_data.location = (0, 0)
# Map Range: [near..far] → [1..0] (near=white, far=black)
map_range = nodes.new('ShaderNodeMapRange')
map_range.location = (200, 0)
map_range.inputs['From Min'].default_value = near
map_range.inputs['From Max'].default_value = far
map_range.inputs['To Min'].default_value = 1.0
map_range.inputs['To Max'].default_value = 0.0
map_range.clamp = True
# Emission shader (unlit)
emission = nodes.new('ShaderNodeEmission')
emission.location = (400, 0)
# Output
output = nodes.new('ShaderNodeOutputMaterial')
output.location = (600, 0)
links.new(cam_data.outputs['View Z Depth'], map_range.inputs['Value'])
links.new(map_range.outputs['Result'], emission.inputs['Color'])
links.new(emission.outputs['Emission'], output.inputs['Surface'])
return mat
def _get_converted_obj(scene, operator):
"""Find the converted copy object using saved name."""
converted_name = scene.get('_comfy_converted_name', '')
converted_obj = bpy.data.objects.get(converted_name)
if converted_obj is None or converted_obj.type != 'MESH':
operator.report({'ERROR'}, message='Cannot find converted copy: ' + converted_name)
return None
return converted_obj
def _get_depth_range(scene, converted_obj):
"""Calculate the min/max view-Z depth from the camera to the object bounding box."""
cam = scene.camera
if cam is None:
return 0.1, 10.0
cam_inv = cam.matrix_world.inverted()
min_z = float('inf')
max_z = float('-inf')
for corner in converted_obj.bound_box:
world_pt = converted_obj.matrix_world @ Vector(corner)
local_pt = cam_inv @ world_pt
z = -local_pt.z # View Z Depth is positive distance along camera -Z
min_z = min(min_z, z)
max_z = max(max_z, z)
# Add padding
padding = max(0.01, (max_z - min_z) * 0.1)
return max(0.001, min_z - padding), max_z + padding
def _render_with_material(operator, mat, out_name, label):
"""Generic render: apply material to converted copy, render, restore."""
scene = bpy.context.scene
output_dir = _ensure_output_dir()
converted_obj = _get_converted_obj(scene, operator)
if converted_obj is None:
return False
# Save original state
orig_materials = [slot.material for slot in converted_obj.material_slots]
orig_use_nodes = scene.use_nodes
orig_use_compositing = scene.render.use_compositing
orig_filepath = scene.render.filepath
orig_format = scene.render.image_settings.file_format
orig_color_mode = scene.render.image_settings.color_mode
# Replace ALL material slots with the render material
if len(converted_obj.material_slots) == 0:
converted_obj.data.materials.append(mat)
else:
for i in range(len(converted_obj.material_slots)):
converted_obj.material_slots[i].material = mat
# Fully disable compositor (both flags needed in Blender 5.0)
scene.use_nodes = False
scene.render.use_compositing = False
# Setup output
out_path = os.path.join(output_dir, out_name)
scene.render.filepath = out_path
scene.render.image_settings.file_format = 'PNG'
scene.render.image_settings.color_mode = 'RGB'
# Render
bpy.ops.render.render(write_still=True)
# Restore original materials per slot
if len(orig_materials) == 0:
converted_obj.data.materials.clear()
else:
for i, mat_orig in enumerate(orig_materials):
if i < len(converted_obj.material_slots):
converted_obj.material_slots[i].material = mat_orig
scene.use_nodes = orig_use_nodes
scene.render.use_compositing = orig_use_compositing
scene.render.filepath = orig_filepath
scene.render.image_settings.file_format = orig_format
scene.render.image_settings.color_mode = orig_color_mode
_add_rendered_image(scene, label, out_path)
return True
def _render_depth(operator):
"""Render depth map using a programmatic depth material."""
scene = bpy.context.scene
converted_obj = _get_converted_obj(scene, operator)
if converted_obj is None:
return False
near, far = _get_depth_range(scene, converted_obj)
depth_mat = _create_depth_material(near, far)
return _render_with_material(operator, depth_mat, 'depth.png', 'depth')
def _render_normal(operator):
"""Render normal map using Normal_Mat material."""
normal_mat = _append_from_materials_blend('Material', 'Normal_Mat')
if normal_mat is None:
operator.report({'ERROR'}, message='Failed to load Normal_Mat from Materials.blend')
return False
return _render_with_material(operator, normal_mat, 'normal.png', '法向')
def _create_white_material():
"""Create a flat white emission material for line art background."""
mat_name = '__ComfyWhiteMat__'
if mat_name in bpy.data.materials:
return bpy.data.materials[mat_name]
mat = bpy.data.materials.new(name=mat_name)
mat.use_nodes = True
nodes = mat.node_tree.nodes
links = mat.node_tree.links
nodes.clear()
emission = nodes.new('ShaderNodeEmission')
emission.location = (0, 0)
emission.inputs['Color'].default_value = (1, 1, 1, 1)
emission.inputs['Strength'].default_value = 1.0
output = nodes.new('ShaderNodeOutputMaterial')
output.location = (200, 0)
links.new(emission.outputs['Emission'], output.inputs['Surface'])
return mat
def _render_lineart(operator):
"""Render line art using Freestyle (black lines on white) for ControlNet Canny."""
scene = bpy.context.scene
output_dir = _ensure_output_dir()
converted_obj = _get_converted_obj(scene, operator)
if converted_obj is None:
return False
# Save original state
vl = scene.view_layers["ViewLayer"]
orig_materials = [slot.material for slot in converted_obj.material_slots]
orig_use_freestyle = scene.render.use_freestyle
orig_vl_use_freestyle = vl.use_freestyle
orig_line_thickness = scene.render.line_thickness
orig_use_nodes = scene.use_nodes
orig_use_compositing = scene.render.use_compositing
orig_filepath = scene.render.filepath
orig_format = scene.render.image_settings.file_format
orig_color_mode = scene.render.image_settings.color_mode
orig_film_transparent = scene.render.film_transparent
orig_world = scene.world
# Apply white emission material to all slots
white_mat = _create_white_material()
if len(converted_obj.material_slots) == 0:
converted_obj.data.materials.append(white_mat)
else:
for i in range(len(converted_obj.material_slots)):
converted_obj.material_slots[i].material = white_mat
# Enable Freestyle
scene.render.use_freestyle = True
vl.use_freestyle = True
scene.render.line_thickness = 1.0
# Configure lineset: black lines, all relevant edge types
fs = vl.freestyle_settings
# Save original linesets state
orig_lineset_settings = []
for ls in fs.linesets:
orig_lineset_settings.append({
'color': tuple(ls.linestyle.color),
'thickness': ls.linestyle.thickness,
})
if len(fs.linesets) == 0:
ls = fs.linesets.new('LineSet')
else:
ls = fs.linesets[0]
ls.linestyle.color = (0, 0, 0) # Black lines
ls.linestyle.thickness = 1.5
ls.select_silhouette = True
ls.select_border = True
ls.select_crease = True
ls.select_edge_mark = True
# White world background
white_world = bpy.data.worlds.new('__ComfyWhiteWorld__')
white_world.use_nodes = True
bg_node = white_world.node_tree.nodes.get('Background')
if bg_node:
bg_node.inputs['Color'].default_value = (1, 1, 1, 1)
bg_node.inputs['Strength'].default_value = 1.0
scene.world = white_world
scene.render.film_transparent = False
# Disable compositor
scene.use_nodes = False
scene.render.use_compositing = False
# Setup output
out_path = os.path.join(output_dir, 'lineart.png')
scene.render.filepath = out_path
scene.render.image_settings.file_format = 'PNG'
scene.render.image_settings.color_mode = 'RGB'
# Render
bpy.ops.render.render(write_still=True)
# Restore materials
if len(orig_materials) == 0:
converted_obj.data.materials.clear()
else:
for i, mat_orig in enumerate(orig_materials):
if i < len(converted_obj.material_slots):
converted_obj.material_slots[i].material = mat_orig
# Restore lineset settings
for i, ls_data in enumerate(orig_lineset_settings):
if i < len(fs.linesets):
fs.linesets[i].linestyle.color = ls_data['color']
fs.linesets[i].linestyle.thickness = ls_data['thickness']
# Restore all settings
scene.render.use_freestyle = orig_use_freestyle
vl.use_freestyle = orig_vl_use_freestyle
scene.render.line_thickness = orig_line_thickness
scene.use_nodes = orig_use_nodes
scene.render.use_compositing = orig_use_compositing
scene.render.filepath = orig_filepath
scene.render.image_settings.file_format = orig_format
scene.render.image_settings.color_mode = orig_color_mode
scene.render.film_transparent = orig_film_transparent
scene.world = orig_world
# Remove temporary world
if '__ComfyWhiteWorld__' in bpy.data.worlds:
bpy.data.worlds.remove(bpy.data.worlds['__ComfyWhiteWorld__'])
_add_rendered_image(scene, '线稿', out_path)
return True
class Comfy_OT_Mvgen_B6229(bpy.types.Operator):
bl_idname = "comfyui.mvgen_b6229"
bl_label = "生成多视角视图"
bl_description = "Generate multi-view maps including depth and normal"
bl_options = {"REGISTER", "UNDO"}
@classmethod
def poll(cls, context):
if "" == bpy.data.filepath:
if bpy.app.version >= (3, 0, 0):
cls.poll_message_set(
"Please save your Blender file before running the script."
)
return False
if len(context.scene.sna_rendered_images) > 0:
if bpy.app.version >= (3, 0, 0):
cls.poll_message_set(
"Please cleanup existing results first (click trash icon)."
)
return False
return True
def execute(self, _context):
if "" == bpy.data.filepath:
self.report(
{"ERROR"},
message="Please save your Blender file before running the script.",
)
else:
if (None is not bpy.context.view_layer.objects.active):
if bpy.context.view_layer.objects.active.type == 'MESH':
bpy.context.scene.sna_activeobject = bpy.context.view_layer.objects.active
bpy.context.scene.sna_activeobject_name = bpy.context.scene.sna_activeobject.name
bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
if property_exists("bpy.data.node_groups['UV_New']", globals(), locals()):
pass
else:
before_data = list(bpy.data.node_groups)
bpy.ops.wm.append(directory=os.path.join(os.path.dirname(__file__), 'assets', 'geonodesuvpack.blend') + r'\NodeTree', filename='UV_New', link=False)
new_data = list(filter(lambda d: d not in before_data, list(bpy.data.node_groups)))
appended_4ED0B = None if not new_data else new_data[0]
bpy.context.view_layer.objects.active = bpy.context.scene.sna_activeobject
bpy.context.scene.sna_activeobject.select_set(state=True, )
modifier_5D814 = bpy.context.scene.sna_activeobject.modifiers.new(name='MV_Nodes', type='NODES', )
node_group = bpy.data.node_groups.get('UV_New')
if node_group is None:
self.report({'ERROR'}, message='Missing node group: UV_New')
return {"CANCELLED"}
bpy.context.scene.sna_activeobject.modifiers['MV_Nodes'].node_group = node_group
bpy.ops.object.convert(target='MESH', keep_original=True)
# Save the converted copy name immediately (active object after convert)
bpy.context.scene['_comfy_converted_name'] = bpy.context.active_object.name
bpy.data.objects[bpy.context.scene.sna_activeobject_name].hide_render = True
bpy.ops.comfyui.aovbake_ad980('INVOKE_DEFAULT', )
bpy.data.objects[bpy.context.scene.sna_activeobject_name].modifiers.clear()
bpy.context.view_layer.objects[bpy.context.scene.sna_activeobject_name].hide_set(state=True, )
# --- Render Depth Map ---
_render_depth(self)
# --- Render Normal Map ---
_render_normal(self)
# --- Render Line Art ---
_render_lineart(self)
return {"FINISHED"}
def invoke(self, context, _event):
return self.execute(context)
class Comfy_OT_RemoveRenderedImage(bpy.types.Operator):
"""Remove a rendered image from the list."""
bl_idname = "comfyui.remove_rendered_image"
bl_label = "Remove"
bl_description = "Remove this image from the list"
bl_options = {"REGISTER", "UNDO"}
index: bpy.props.IntProperty()
def execute(self, context):
images = context.scene.sna_rendered_images
if 0 <= self.index < len(images):
images.remove(self.index)
# Adjust active index
if context.scene.sna_rendered_images_index >= len(images):
context.scene.sna_rendered_images_index = max(0, len(images) - 1)
return {'FINISHED'}
class Comfy_OT_OpenRenderedFolder(bpy.types.Operator):
"""Open the folder containing the rendered image."""
bl_idname = "comfyui.open_rendered_folder"
bl_label = "Open Folder"
bl_description = "Open the folder containing this image"
filepath: bpy.props.StringProperty()
def execute(self, context):
folder = os.path.dirname(self.filepath)
if os.path.isdir(folder):
subprocess.Popen(['explorer', folder])
return {'FINISHED'}
class Comfy_OT_CleanupMvgen(bpy.types.Operator):
"""Remove generated multi-view objects and restore original."""
bl_idname = "comfyui.cleanup_mvgen"
bl_label = "Cleanup"
bl_description = "Delete the generated multi-view copy, camera, and restore original object"
bl_options = {"REGISTER", "UNDO"}
def execute(self, context):
scene = context.scene
# Delete the converted copy using saved name
converted_name = scene.get('_comfy_converted_name', '')
if converted_name and converted_name in bpy.data.objects:
bpy.data.objects.remove(bpy.data.objects[converted_name], do_unlink=True)
# Delete the bake camera
bake_cam = scene.sna_bakecam
if bake_cam is not None:
bpy.data.objects.remove(bake_cam, do_unlink=True)
scene.sna_bakecam = None
# Also remove any camera named ObjectCamera
if 'ObjectCamera' in bpy.data.objects:
bpy.data.objects.remove(bpy.data.objects['ObjectCamera'], do_unlink=True)
# Restore original object visibility
orig_name = scene.sna_activeobject_name
if orig_name and orig_name in bpy.data.objects:
orig_obj = bpy.data.objects[orig_name]
orig_obj.hide_render = False
orig_obj.hide_set(False)
bpy.context.view_layer.objects.active = orig_obj
orig_obj.select_set(True)
# Clear rendered images list
scene.sna_rendered_images.clear()
scene.sna_rendered_images_index = 0
# Clear preview collection
pcoll = _get_preview_collection()
pcoll.clear()
# Remove temp materials
for mat_name in ('__ComfyDepthMat__', '__ComfyWhiteMat__'):
if mat_name in bpy.data.materials:
bpy.data.materials.remove(bpy.data.materials[mat_name])
# Clear converted name
if '_comfy_converted_name' in scene:
del scene['_comfy_converted_name']
return {'FINISHED'}
class Comfy_PT_ComfyProjectorPanel(bpy.types.Panel):
bl_label = "ComfyProjector"
bl_idname = "Comfy_PT_comfyprojector"
bl_space_type = "VIEW_3D"
bl_region_type = "UI"
bl_category = "ComfyProjector"
def draw(self, context):
layout = self.layout
scene = context.scene
# --- ControlNet Layers section ---
header_box = layout.box()
header_box.label(text="ControlNet Layers", icon='RENDERLAYERS')
# Resolution selector
row = header_box.row(align=True)
row.prop(scene, "sna_render_resolution", text="Resolution")
# Main generate button + cleanup button (only after generation)
has_results = len(scene.sna_rendered_images) > 0
btn_row = header_box.row(align=True)
btn_row.operator(
"comfyui.mvgen_b6229",
text="Generate Multiview Maps",
icon='RENDER_RESULT',
)
if has_results:
btn_row.operator(
"comfyui.cleanup_mvgen",
text="",
icon='TRASH',
)
# --- Rendered Images List ---
if has_results:
img_box = layout.box()
img_box.label(text="Rendered Images", icon='IMAGE_DATA')
for i, item in enumerate(scene.sna_rendered_images):
item_box = img_box.box()
row = item_box.row()
# Left: thumbnail preview
pcoll = _get_preview_collection()
preview_name = os.path.basename(item.filepath)
if preview_name in pcoll:
row.template_icon(
icon_value=pcoll[preview_name].icon_id,
scale=4.0,
)
else:
row.label(text="", icon='IMAGE_DATA')
# Right: name + filepath, vertically centered
col_info = row.column()
col_info.separator(factor=2.0)
col_info.label(text=item.name, icon='IMAGE_DATA')
info_row2 = col_info.row(align=True)
info_row2.label(text=item.filepath, icon='DISK_DRIVE')
op2 = info_row2.operator(
"comfyui.open_rendered_folder",
text="",
icon='FILEBROWSER',
emboss=False,
)
op2.filepath = item.filepath
def cleanup_preview_collections():
"""Remove all preview collections on unregister."""
for pcoll in _preview_collections.values():
bpy.utils.previews.remove(pcoll)
_preview_collections.clear()
classes = (
Comfy_RenderedImageItem,
Comfy_OT_Aovbake_Ad980,
Comfy_OT_Mvgen_B6229,
Comfy_OT_RemoveRenderedImage,
Comfy_OT_OpenRenderedFolder,
Comfy_OT_CleanupMvgen,
Comfy_PT_ComfyProjectorPanel,
)