From 66a3e779bbf4f9f8f5d3aacada2ee94346b9d4b2 Mon Sep 17 00:00:00 2001 From: rickard Date: Sat, 5 Oct 2024 11:19:18 +0200 Subject: [PATCH] add option to not return image to Blender --- sd_render/__init__.py | 12 ++++++++++++ sd_render/camera_utils.py | 14 +------------- sd_render/generate_texture.py | 16 +++++++++------- sd_render/image_gen/automatic1111.py | 3 --- sd_render/image_gen/comfy_ui.py | 3 +-- 5 files changed, 23 insertions(+), 25 deletions(-) diff --git a/sd_render/__init__.py b/sd_render/__init__.py index d4d936f..85eb1d7 100644 --- a/sd_render/__init__.py +++ b/sd_render/__init__.py @@ -155,6 +155,16 @@ class StableDiffusionProperties(PropertyGroup): description="Bake the generated image to the selected object", default=True ) + return_image: BoolProperty( + name="Return image to Blender", + description="Whether to process the image in Blender", + default=True + ) + output_image_name: StringProperty( + name="Output image name", + description="Name of the image to save", + default="sd_output.png" + ) class Automatic1111Properties(StableDiffusionProperties): auto1111_samplers = [ @@ -302,6 +312,8 @@ def draw_generation_settings(layout, generate_image_properties): layout.prop(generate_image_properties, "sd_port") layout.separator() layout.label(text="Baking settings") + layout.prop(generate_image_properties, "output_image_name", text="Output image") + layout.prop(generate_image_properties, "return_image") layout.prop(generate_image_properties, "delete_projector") layout.prop(generate_image_properties, "bake_texture") diff --git a/sd_render/camera_utils.py b/sd_render/camera_utils.py index 356f844..052ee4a 100644 --- a/sd_render/camera_utils.py +++ b/sd_render/camera_utils.py @@ -27,21 +27,9 @@ def project_uvs(projector): bpy.ops.uv.project_from_view() bpy.ops.object.mode_set(mode = 'OBJECT') bpy.context.active_object.select_set(False) + bpy.ops.object.select_all(action='DESELECT') apply_viewport_position(view_params) -def project_uv_from_active_camera(obj): - active_camera = bpy.context.scene.camera - - if active_camera is None: - print("No active camera found.") - return - - bpy.context.view_layer.objects.active = obj - bpy.context.scene.tool_settings.uv_texture_add_operator_projector = active_camera - - with bpy.context.temp_override(edit_object=bpy.context.edit_object): - bpy.ops.uv.project_from_view(correct_aspect=True, camera_bounds=True) - """ Source: https://b3d.interplanety.org/en/saving-and-restoring-the-3d-viewport-position/ diff --git a/sd_render/generate_texture.py b/sd_render/generate_texture.py index f21b68d..ebcafec 100644 --- a/sd_render/generate_texture.py +++ b/sd_render/generate_texture.py @@ -3,7 +3,7 @@ from sd_render.image_gen.comfy_ui import ComfyUi from sd_render.image_gen.automatic1111 import Automatic1111 from sd_render.camera_utils import render_viewport, project_uvs -from sd_render.image_utils import bake_from_active, create_projector_objects, set_projector_position_and_orientation, setup_projector_material, assign_material_to_projector, remove_projector, create_projector_object_from_list +from sd_render.image_utils import bake_from_active, create_projector_objects, set_projector_position_and_orientation, setup_projector_material, assign_material_to_projector, remove_projector logger = logging.getLogger(__name__) logger.setLevel(logging.DEBUG) @@ -23,7 +23,7 @@ def render(props): return "Invalid backend selected." depth_map = f'{output_folder}/viewer_node.png' #generated_path = "/tmp/sd_output.png" - generated_path = generator.generate_image(prompt=props.prompt, + image_data = generator.generate_image(prompt=props.prompt, negative_prompt=props.negative_prompt, seed=props.seed, sampler=props.sampler, @@ -36,10 +36,12 @@ def render(props): depth_map=depth_map, scheduler=scheduler, model=props.model) - try: - return bpy.data.images.load(f'{output_folder}/sd_output.png') - except: - return "Failed to generate texture." + if props.return_image: + generator.convert_image(image_data, props.output_image_name) + try: + return bpy.data.images.load(f'{output_folder}/{props.output_image_name}') + except: + return "Failed to generate texture." def bake(selected_objects, texture_image, delete_projector: bool) -> str: @@ -78,7 +80,7 @@ def execute(): return texture_image # Error message selected_objects = None - if props.bake_texture: + if props.bake_texture and props.return_image: selected_objects = bpy.context.selected_objects return bake(selected_objects, texture_image=texture_image, delete_projector=props.delete_projector) return '' diff --git a/sd_render/image_gen/automatic1111.py b/sd_render/image_gen/automatic1111.py index 54c1a99..9671cdc 100644 --- a/sd_render/image_gen/automatic1111.py +++ b/sd_render/image_gen/automatic1111.py @@ -16,9 +16,6 @@ def generate_image(self, prompt: str, depth_map: str, negative_prompt: str = "te if model: self.set_model(model) image_data = self.send_request(prompt, depth_map, negative_prompt, seed, sampler, steps, cfg_scale, width, height, cn_weight, cn_guidance) - if not image_data: - return None - self.convert_image(image_data, "sd_output") return image_data diff --git a/sd_render/image_gen/comfy_ui.py b/sd_render/image_gen/comfy_ui.py index 019fcff..3304be5 100644 --- a/sd_render/image_gen/comfy_ui.py +++ b/sd_render/image_gen/comfy_ui.py @@ -22,8 +22,7 @@ def generate_image(self, prompt: str, depth_map: str, negative_prompt: str = "te image_data = self.send_request(prompt, depth_map, negative_prompt, seed, sampler, steps, cfg_scale, width, height, cn_weight, cn_guidance, scheduler, model, cn_start, cn_end) if not image_data: return None - self.convert_image(image_data[0], "sd_output") - return image_data + return image_data[0] def send_request(self, prompt, depth_map, negative_prompt: str, seed: int, sampler: str, steps: int, cfg_scale: int, width: int, height: int, cn_weight: float, cn_guidance: float, scheduler: str = '', model: str = '', cn_start: float = 1.0, cn_end: float = 1.0) -> bytes: