diff --git a/__init__.py b/__init__.py index 81184b5..572db8c 100644 --- a/__init__.py +++ b/__init__.py @@ -12,7 +12,8 @@ bl_info = { } import bpy # noqa: E402 -from .utils import classes # noqa: E402 +from .utils import classes, cleanup_preview_collections # noqa: E402 +from .utils import Comfy_RenderedImageItem # noqa: E402 def register(): @@ -30,9 +31,21 @@ def register(): ], default='1024', ) + bpy.types.Scene.sna_rendered_images = bpy.props.CollectionProperty( + type=Comfy_RenderedImageItem, + ) + bpy.types.Scene.sna_rendered_images_index = bpy.props.IntProperty( + name="Active Rendered Image Index", + default=0, + ) def unregister(): + cleanup_preview_collections() + if hasattr(bpy.types.Scene, "sna_rendered_images_index"): + del bpy.types.Scene.sna_rendered_images_index + if hasattr(bpy.types.Scene, "sna_rendered_images"): + del bpy.types.Scene.sna_rendered_images if hasattr(bpy.types.Scene, "sna_render_resolution"): del bpy.types.Scene.sna_render_resolution if hasattr(bpy.types.Scene, "sna_activeobject_name"): @@ -43,3 +56,4 @@ def unregister(): del bpy.types.Scene.sna_activeobject for cls in reversed(classes): bpy.utils.unregister_class(cls) + diff --git a/utils.py b/utils.py index 868ccd3..7baf4d7 100644 --- a/utils.py +++ b/utils.py @@ -2,6 +2,35 @@ 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) @@ -132,37 +161,67 @@ def _append_from_materials_blend(data_type, name): 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 _render_depth(operator): """Render depth map using Depth_Mat compositor NodeTree.""" scene = bpy.context.scene output_dir = _ensure_output_dir() - # Append the Depth_Mat compositor node tree + # Force fresh Depth_Mat by removing any cached version + if 'Depth_Mat' in bpy.data.node_groups: + bpy.data.node_groups.remove(bpy.data.node_groups['Depth_Mat']) + + # Append fresh Depth_Mat from Materials.blend depth_tree = _append_from_materials_blend('NodeTree', 'Depth_Mat') if depth_tree is None: operator.report({'ERROR'}, message='Failed to load Depth_Mat from Materials.blend') return False - # Save original compositor state + # Save original render settings orig_use_nodes = scene.use_nodes orig_comp_group = scene.compositing_node_group + orig_filepath = scene.render.filepath + orig_format = scene.render.image_settings.file_format + orig_color_mode = scene.render.image_settings.color_mode # Enable compositor and assign Depth_Mat as the compositing node group scene.use_nodes = True scene.compositing_node_group = depth_tree + scene.view_layers["ViewLayer"].use_pass_z = True # Setup output - scene.render.filepath = os.path.join(output_dir, 'depth.png') + out_path = os.path.join(output_dir, 'depth.png') + scene.render.filepath = out_path scene.render.image_settings.file_format = 'PNG' scene.render.image_settings.color_mode = 'RGB' - scene.view_layers["ViewLayer"].use_pass_z = True # Render bpy.ops.render.render(write_still=True) - # Restore original compositor state + # Restore all settings scene.compositing_node_group = orig_comp_group scene.use_nodes = orig_use_nodes + scene.render.filepath = orig_filepath + scene.render.image_settings.file_format = orig_format + scene.render.image_settings.color_mode = orig_color_mode + + # Add to rendered images list + _add_rendered_image(scene, 'depth', out_path) return True @@ -184,31 +243,41 @@ def _render_normal(operator): if converted_obj is None or converted_obj.type != 'MESH': converted_obj = obj - # Save original materials + # Save original materials and render settings orig_materials = [slot.material for slot in converted_obj.material_slots] + orig_use_nodes = scene.use_nodes + orig_filepath = scene.render.filepath + orig_format = scene.render.image_settings.file_format + orig_color_mode = scene.render.image_settings.color_mode # Clear and assign Normal_Mat converted_obj.data.materials.clear() converted_obj.data.materials.append(normal_mat) # Disable compositor for normal pass (direct render) - orig_use_nodes = scene.use_nodes scene.use_nodes = False # Setup output - scene.render.filepath = os.path.join(output_dir, 'normal.png') + out_path = os.path.join(output_dir, 'normal.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 original materials + # Restore original materials and settings converted_obj.data.materials.clear() for mat in orig_materials: converted_obj.data.materials.append(mat) scene.use_nodes = orig_use_nodes + scene.render.filepath = orig_filepath + scene.render.image_settings.file_format = orig_format + scene.render.image_settings.color_mode = orig_color_mode + + # Add to rendered images list + _add_rendered_image(scene, 'normal', out_path) return True @@ -220,12 +289,20 @@ class Comfy_OT_Mvgen_B6229(bpy.types.Operator): 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 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: @@ -272,6 +349,88 @@ class Comfy_OT_Mvgen_B6229(bpy.types.Operator): 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 (active object that is NOT the original) + orig_name = scene.sna_activeobject_name + if orig_name: + # The converted copy is typically named with .001 suffix + for obj in list(bpy.data.objects): + if obj.name != orig_name and obj.name.startswith(orig_name) and obj.type == 'MESH': + bpy.data.objects.remove(obj, 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 + 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) + # Make it active and selected + 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() + + return {'FINISHED'} + + class Comfy_PT_ComfyProjectorPanel(bpy.types.Panel): bl_label = "ComfyProjector" bl_idname = "Comfy_PT_comfyprojector" @@ -281,12 +440,79 @@ class Comfy_PT_ComfyProjectorPanel(bpy.types.Panel): def draw(self, context): layout = self.layout - layout.prop(context.scene, "sna_render_resolution", text="Resolution") - layout.operator("comfyui.mvgen_b6229", text="生成多视角视图") + 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, )