From 7e8555399c531316476d69110592678ab5f48a9a Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Thu, 2 Mar 2023 17:40:24 -0500 Subject: [PATCH 01/28] Initial controlnet support --- generator_process/__init__.py | 1 + generator_process/actions/control_net.py | 303 +++++++++++++++++++ generator_process/actions/huggingface_hub.py | 57 ++-- generator_process/actions/prompt_to_image.py | 18 +- operators/dream_texture.py | 6 + property_groups/dream_prompt.py | 15 +- ui/panels/dream_texture.py | 7 +- 7 files changed, 377 insertions(+), 30 deletions(-) create mode 100644 generator_process/actions/control_net.py diff --git a/generator_process/__init__.py b/generator_process/__init__.py index 2677884a..52268c36 100644 --- a/generator_process/__init__.py +++ b/generator_process/__init__.py @@ -11,6 +11,7 @@ class Generator(Actor): from .actions.outpaint import outpaint from .actions.upscale import upscale from .actions.depth_to_image import depth_to_image + from .actions.control_net import control_net from .actions.huggingface_hub import hf_snapshot_download, hf_list_models, hf_list_installed_models from .actions.ocio_transform import ocio_transform from .actions.convert_original_stable_diffusion_to_diffusers import convert_original_stable_diffusion_to_diffusers diff --git a/generator_process/actions/control_net.py b/generator_process/actions/control_net.py new file mode 100644 index 00000000..09756773 --- /dev/null +++ b/generator_process/actions/control_net.py @@ -0,0 +1,303 @@ +from typing import Union, Generator, Callable, List, Optional, Dict, Any +from contextlib import nullcontext + +from numpy.typing import NDArray +import numpy as np +import random +from .prompt_to_image import Scheduler, Optimizations, StepPreviewMode, ImageGenerationResult, _configure_model_padding, model_snapshot_folder, load_pipe +from ..models import Pipeline +from .detect_seamless import SeamlessAxes + +def control_net( + self, + pipeline: Pipeline, + + model: str, + + scheduler: Scheduler, + + optimizations: Optimizations, + + control_net: str, + control: NDArray | None, + controlnet_conditioning_scale: float, + image: NDArray | str | None, + strength: float, + prompt: str | list[str], + steps: int, + seed: int, + + width: int | None, + height: int | None, + + cfg_scale: float, + use_negative_prompt: bool, + negative_prompt: str, + + seamless_axes: SeamlessAxes | str | bool | tuple[bool, bool] | None, + + step_preview_mode: StepPreviewMode, + + **kwargs +) -> Generator[NDArray, None, None]: + match pipeline: + case Pipeline.STABLE_DIFFUSION: + import diffusers + import torch + import PIL.Image + import PIL.ImageOps + + class GeneratorPipeline(diffusers.StableDiffusionControlNetPipeline): + @torch.no_grad() + def __call__( + self, + prompt: Union[str, List[str]] = None, + image: Union[torch.FloatTensor, PIL.Image.Image, List[torch.FloatTensor], List[PIL.Image.Image]] = None, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 50, + guidance_scale: float = 7.5, + negative_prompt: Optional[Union[str, List[str]]] = None, + num_images_per_prompt: Optional[int] = 1, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, + callback_steps: int = 1, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + controlnet_conditioning_scale: float = 1.0, + + **kwargs + ): + # 0. Default height and width to unet + height, width = self._default_height_width(height, width, image) + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt, image, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds + ) + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self._execution_device + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + do_classifier_free_guidance = guidance_scale > 1.0 + + # 3. Encode input prompt + prompt_embeds = self._encode_prompt( + prompt, + device, + num_images_per_prompt, + do_classifier_free_guidance, + negative_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + ) + + # 4. Prepare image + image = self.prepare_image( + image, + width, + height, + batch_size * num_images_per_prompt, + num_images_per_prompt, + device, + self.controlnet.dtype, + ) + + if do_classifier_free_guidance: + image = torch.cat([image] * 2) + + # 5. Prepare timesteps + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = self.scheduler.timesteps + + # 6. Prepare latent variables + num_channels_latents = self.unet.in_channels + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + prompt_embeds.dtype, + device, + generator, + latents, + ) + + # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + + # 8. Denoising loop + num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + + down_block_res_samples, mid_block_res_sample = self.controlnet( + latent_model_input, + t, + encoder_hidden_states=prompt_embeds, + controlnet_cond=image, + return_dict=False, + ) + + down_block_res_samples = [ + down_block_res_sample * controlnet_conditioning_scale + for down_block_res_sample in down_block_res_samples + ] + mid_block_res_sample *= controlnet_conditioning_scale + + # predict the noise residual + noise_pred = self.unet( + latent_model_input, + t, + encoder_hidden_states=prompt_embeds, + cross_attention_kwargs=cross_attention_kwargs, + down_block_additional_residuals=down_block_res_samples, + mid_block_additional_residual=mid_block_res_sample, + ).sample + + # perform guidance + if do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample + + # NOTE: Modified to yield the latents instead of calling a callback. + yield ImageGenerationResult.step_preview(self, kwargs['step_preview_mode'], width, height, latents, generator, i) + + # If we do sequential model offloading, let's offload unet and controlnet + # manually for max memory savings + if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: + self.unet.to("cpu") + self.controlnet.to("cpu") + torch.cuda.empty_cache() + + if output_type == "latent": + image = latents + has_nsfw_concept = None + elif output_type == "pil": + # 8. Post-processing + image = self.decode_latents(latents) + + # NOTE: Add UI to enable this. + # 9. Run safety checker + # image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype) + + # 10. Convert to PIL + image = self.numpy_to_pil(image) + else: + # 8. Post-processing + image = self.decode_latents(latents) + + # NOTE: Add UI to enable this. + # 9. Run safety checker + # image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype) + + # Offload last model to CPU + if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: + self.final_offload_hook.offload() + + # NOTE: Modified to yield the decoded image as a numpy array. + yield ImageGenerationResult( + [np.asarray(PIL.ImageOps.flip(image).convert('RGBA'), dtype=np.float32) / 255. + for i, image in enumerate(image)], + [gen.initial_seed() for gen in generator] if isinstance(generator, list) else [generator.initial_seed()], + num_inference_steps, + True + ) + + if optimizations.cpu_only: + device = "cpu" + else: + device = self.choose_device() + + # Load the ControlNet model + controlnet = load_pipe(self, "control_net_model", diffusers.ControlNetModel, control_net, optimizations, None, device) + + # StableDiffusionPipeline w/ caching + pipe = load_pipe(self, "control_net", GeneratorPipeline, model, optimizations, scheduler, device, controlnet=controlnet) + + # Optimizations + pipe = optimizations.apply(pipe, device) + + # RNG + batch_size = len(prompt) if isinstance(prompt, list) else 1 + generator = [] + for _ in range(batch_size): + gen = torch.Generator(device="cpu" if device in ("mps", "privateuseone") else device) # MPS and DML do not support the `Generator` API + generator.append(gen.manual_seed(random.randrange(0, np.iinfo(np.uint32).max) if seed is None else seed)) + if batch_size == 1: + # Some schedulers don't handle a list of generators: https://github.com/huggingface/diffusers/issues/1909 + generator = generator[0] + + # Init Image + # FIXME: The `unet.config.sample_size` of the depth model is `32`, not `64`. For now, this will be hardcoded to `512`. + height = height or 512 + width = width or 512 + rounded_size = ( + int(8 * (width // 8)), + int(8 * (height // 8)), + ) + control_image = PIL.Image.fromarray(np.uint8(control * 255)).convert('RGB').resize(rounded_size) if control is not None else None + init_image = None if image is None else (PIL.Image.open(image) if isinstance(image, str) else PIL.Image.fromarray(image.astype(np.uint8))).convert('RGB').resize(rounded_size) + + # Seamless + if seamless_axes == SeamlessAxes.AUTO: + init_sa = None if init_image is None else self.detect_seamless(np.array(init_image) / 255) + control_sa = None if control_image is None else self.detect_seamless(np.array(control_image) / 255) + if init_sa is not None and control_sa is not None: + seamless_axes = SeamlessAxes((init_sa.x and control_sa.x, init_sa.y and control_sa.y)) + elif init_sa is not None: + seamless_axes = init_sa + elif control_sa is not None: + seamless_axes = control_sa + _configure_model_padding(pipe.unet, seamless_axes) + _configure_model_padding(pipe.vae, seamless_axes) + + # Inference + with (torch.inference_mode() if device not in ('mps', "privateuseone") else nullcontext()), \ + (torch.autocast(device) if optimizations.can_use("amp", device) else nullcontext()): + yield from pipe( + prompt=prompt, + image=control_image, + controlnet_conditioning_scale=controlnet_conditioning_scale, + # image=init_image, + strength=strength, + width=rounded_size[0], + height=rounded_size[1], + num_inference_steps=steps, + guidance_scale=cfg_scale, + negative_prompt=negative_prompt if use_negative_prompt else None, + num_images_per_prompt=1, + eta=0.0, + generator=generator, + latents=None, + output_type="pil", + return_dict=True, + callback=None, + callback_steps=1, + step_preview_mode=step_preview_mode + ) + case Pipeline.STABILITY_SDK: + import stability_sdk + raise NotImplementedError() + case _: + raise Exception(f"Unsupported pipeline {pipeline}.") \ No newline at end of file diff --git a/generator_process/actions/huggingface_hub.py b/generator_process/actions/huggingface_hub.py index 657085df..fa1683b5 100644 --- a/generator_process/actions/huggingface_hub.py +++ b/generator_process/actions/huggingface_hub.py @@ -25,6 +25,8 @@ class ModelType(enum.IntEnum): UPSCALING = 7 INPAINTING = 9 + CONTROL_NET = -1 + @classmethod def _missing_(cls, _): return cls.UNKNOWN @@ -68,7 +70,7 @@ def hf_list_models( api = HfApi() setattr(self, "huggingface_hub_api", api) - filter = ModelFilter(tags="diffusers", task="text-to-image") + filter = ModelFilter(tags="diffusers") models = api.list_models( filter=filter, search=query, @@ -87,9 +89,17 @@ def hf_list_installed_models(self) -> list[Model]: def detect_model_type(snapshot_folder): unet_config = os.path.join(snapshot_folder, 'unet', 'config.json') + config = os.path.join(snapshot_folder, 'config.json') if os.path.exists(unet_config): with open(unet_config, 'r') as f: return ModelType(json.load(f)['in_channels']) + elif os.path.exists(config): + with open(config, 'r') as f: + config_dict = json.load(f) + if '_class_name' in config_dict and config_dict['_class_name'] == 'ControlNetModel': + return ModelType.CONTROL_NET + else: + return ModelType.UNKNOWN else: return ModelType.UNKNOWN @@ -101,7 +111,10 @@ def _map_model(file): snapshot_folder = storage_folder model_type = detect_model_type(snapshot_folder) else: - for revision in os.listdir(os.path.join(storage_folder, "refs")): + refs_folder = os.path.join(storage_folder, "refs") + if not os.path.exists(refs_folder): + return None + for revision in os.listdir(refs_folder): ref_path = os.path.join(storage_folder, "refs", revision) with open(ref_path) as f: commit_hash = f.read() @@ -152,24 +165,30 @@ def hf_snapshot_download( from diffusers.utils import DIFFUSERS_CACHE, WEIGHTS_NAME, CONFIG_NAME, ONNX_WEIGHTS_NAME from diffusers.schedulers.scheduling_utils import SCHEDULER_CONFIG_NAME from diffusers.utils.hub_utils import http_user_agent - config_dict = StableDiffusionPipeline.load_config( - model, - cache_dir=DIFFUSERS_CACHE, - resume_download=True, - force_download=False, - use_auth_token=token - ) - # make sure we only download sub-folders and `diffusers` filenames - folder_names = [k for k in config_dict.keys() if not k.startswith("_")] - allow_patterns = [os.path.join(k, "*") for k in folder_names] - allow_patterns += [WEIGHTS_NAME, SCHEDULER_CONFIG_NAME, CONFIG_NAME, ONNX_WEIGHTS_NAME, StableDiffusionPipeline.config_name] - - # make sure we don't download flax, safetensors, or ckpt weights. - ignore_patterns = ["*.msgpack", "*.safetensors", "*.ckpt"] - requested_pipeline_class = config_dict.get("_class_name", StableDiffusionPipeline.__name__) - user_agent = {"pipeline_class": requested_pipeline_class} - user_agent = http_user_agent(user_agent) + try: + config_dict = StableDiffusionPipeline.load_config( + model, + cache_dir=DIFFUSERS_CACHE, + resume_download=True, + force_download=False, + use_auth_token=token + ) + # make sure we only download sub-folders and `diffusers` filenames + folder_names = [k for k in config_dict.keys() if not k.startswith("_")] + allow_patterns = [os.path.join(k, "*") for k in folder_names] + allow_patterns += [WEIGHTS_NAME, SCHEDULER_CONFIG_NAME, CONFIG_NAME, ONNX_WEIGHTS_NAME, StableDiffusionPipeline.config_name] + + # make sure we don't download flax, safetensors, or ckpt weights. + ignore_patterns = ["*.msgpack", "*.safetensors", "*.ckpt"] + + requested_pipeline_class = config_dict.get("_class_name", StableDiffusionPipeline.__name__) + user_agent = {"pipeline_class": requested_pipeline_class} + user_agent = http_user_agent(user_agent) + except: + allow_patterns = None + ignore_patterns = None + user_agent = None # download all allow_patterns diff --git a/generator_process/actions/prompt_to_image.py b/generator_process/actions/prompt_to_image.py index bce36d1b..3b443a56 100644 --- a/generator_process/actions/prompt_to_image.py +++ b/generator_process/actions/prompt_to_image.py @@ -34,7 +34,7 @@ def __init__(self, pipeline: Any, invalidation_properties: tuple, snapshot_folde def is_valid(self, properties: tuple): return properties == self.invalidation_properties -def load_pipe(self, action, generator_pipeline, model, optimizations, scheduler, device): +def load_pipe(self, action, generator_pipeline, model, optimizations, scheduler, device, **kwargs): """ Use a cached pipeline, or create the pipeline class and cache it. @@ -64,17 +64,19 @@ def load_pipe(self, action, generator_pipeline, model, optimizations, scheduler, snapshot_folder, revision=revision, torch_dtype=torch.float16 if optimizations.can_use_half(device) else torch.float32, + **kwargs ) pipe = pipe.to(device) setattr(self, "_cached_pipe", CachedPipeline(pipe, invalidation_properties, snapshot_folder)) cached_pipe = self._cached_pipe - if 'scheduler' in os.listdir(cached_pipe.snapshot_folder): - pipe.scheduler = scheduler.create(pipe, { - 'model_path': cached_pipe.snapshot_folder, - 'subfolder': 'scheduler', - }) - else: - pipe.scheduler = scheduler.create(pipe, None) + if scheduler is not None: + if 'scheduler' in os.listdir(cached_pipe.snapshot_folder): + pipe.scheduler = scheduler.create(pipe, { + 'model_path': cached_pipe.snapshot_folder, + 'subfolder': 'scheduler', + }) + else: + pipe.scheduler = scheduler.create(pipe, None) return pipe class Scheduler(enum.Enum): diff --git a/operators/dream_texture.py b/operators/dream_texture.py index 792f3d8c..8f482f7c 100644 --- a/operators/dream_texture.py +++ b/operators/dream_texture.py @@ -210,6 +210,12 @@ def require_depth(): depth=np.flipud(init_image.astype(np.float32) / 255.), **generated_args, ) + case 'control_net': + f = gen.control_net( + image=None, + control=init_image, + **generated_args + ) case 'inpaint': f = gen.inpaint( image=init_image, diff --git a/property_groups/dream_prompt.py b/property_groups/dream_prompt.py index 62261683..7fa9c9b6 100644 --- a/property_groups/dream_prompt.py +++ b/property_groups/dream_prompt.py @@ -62,6 +62,8 @@ def modify_action_source_type(self, context): ('depth_generated', 'Color and Generated Depth', 'Use MiDaS to infer the depth of the initial image and include it in the conditioning. Can give results that more closely match the composition of the source image', 2), ('depth_map', 'Color and Depth Map', 'Specify a secondary image to use as the depth map. Can give results that closely match the composition of the depth map', 3), ('depth', 'Depth', 'Treat the initial image as a depth map, and ignore any color. Matches the composition of the source image without any color influence', 4), + ('control_net', 'Control Net', 'Treat the initial image as the input to a ControlNet model', 5), + ('control_net_color', 'Color and Control Net', 'Specify a secondary image to use with a ControlNet model', 6), ] def model_options(self, context): @@ -76,6 +78,8 @@ def model_case(model, i): ) models = {} for i, model in enumerate(context.preferences.addons[StableDiffusionPreferences.bl_idname].preferences.installed_models): + if model.model_type in {ModelType.CONTROL_NET.name, ModelType.UNKNOWN.name}: + continue if model.model_type not in models: models[model.model_type] = [model_case(model, i)] else: @@ -104,7 +108,13 @@ def model_case(model, i): def pipeline_options(self, context): return [ (Pipeline.STABLE_DIFFUSION.name, 'Stable Diffusion', 'Stable Diffusion on your own hardware', 1), - (Pipeline.STABILITY_SDK.name, 'DreamStudio', 'Cloud compute via DreamStudio', 2) + (Pipeline.STABILITY_SDK.name, 'DreamStudio', 'Cloud compute via DreamStudio', 2), + ] + +def control_net_options(self, context): + return [ + (model.model_base, model.model_base.replace('models--', '').replace('--', '/'), '') for model in context.preferences.addons[StableDiffusionPreferences.bl_idname].preferences.installed_models + if model.model_type == ModelType.CONTROL_NET.name ] def seed_clamp(self, ctx): @@ -119,6 +129,9 @@ def seed_clamp(self, ctx): attributes = { "pipeline": EnumProperty(name="Pipeline", items=pipeline_options, default=1 if Pipeline.local_available() else 2, description="Specify which model and target should be used."), "model": EnumProperty(name="Model", items=model_options, description="Specify which model to use for inference"), + + "control_net": EnumProperty(name="Control Net", items=control_net_options, description="Specify which ControlNet to use"), + "controlnet_conditioning_scale": FloatProperty(name="ControlNet Conditioning Scale", default=1.0, description="Increases the strength of the ControlNet's effect"), # Prompt "prompt_structure": EnumProperty(name="Preset", items=prompt_structures_items, description="Fill in a few simple options to create interesting images quickly"), diff --git a/ui/panels/dream_texture.py b/ui/panels/dream_texture.py index bfd67725..496726cf 100644 --- a/ui/panels/dream_texture.py +++ b/ui/panels/dream_texture.py @@ -232,7 +232,10 @@ def _outpaint_warning_box(warning): layout.prop(prompt, "use_init_img_color") if prompt.init_img_action == 'modify': layout.prop(prompt, "modify_action_source_type") - if prompt.modify_action_source_type == 'depth_map': + if prompt.modify_action_source_type in {'control_net', 'control_net_color'}: + layout.prop(context.scene.dream_textures_prompt, 'control_net') + layout.prop(context.scene.dream_textures_prompt, 'controlnet_conditioning_scale') + if prompt.modify_action_source_type == 'depth_map' or prompt.modify_action_source_type == 'control_net_color': layout.template_ID(context.scene, "init_depth", open="image.open") yield InitImagePanel @@ -360,4 +363,4 @@ def draw(self, context): e.draw(context, error_box) except Exception as e: print(e) - return ActionsPanel + return ActionsPanel \ No newline at end of file From d180f9d51bb70c42fd0101a34eeba93e8063f54a Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Thu, 2 Mar 2023 22:10:32 -0500 Subject: [PATCH 02/28] Support ControlNet in project --- __init__.py | 1 + operators/project.py | 28 ++++++++++++++++++++-------- property_groups/dream_prompt.py | 6 +++--- 3 files changed, 24 insertions(+), 11 deletions(-) diff --git a/__init__.py b/__init__.py index 7f906cff..0f5f8f93 100644 --- a/__init__.py +++ b/__init__.py @@ -101,6 +101,7 @@ def get_selection_preview(self): bpy.types.Scene.dream_textures_project_prompt = PointerProperty(type=DreamPrompt) bpy.types.Scene.dream_textures_project_framebuffer_arguments = EnumProperty(name="Inputs", items=framebuffer_arguments) bpy.types.Scene.dream_textures_project_bake = BoolProperty(name="Bake", default=False, description="Re-maps the generated texture onto the specified UV map") + bpy.types.Scene.dream_textures_project_use_control_net = BoolProperty(name="Use ControlNet", default=False, description="Use a depth ControlNet instead of a depth model") for cls in CLASSES: bpy.utils.register_class(cls) diff --git a/operators/project.py b/operators/project.py index ce3e0c99..d1ef31d2 100644 --- a/operators/project.py +++ b/operators/project.py @@ -98,7 +98,6 @@ def draw(self, context): if Pipeline[context.scene.dream_textures_project_prompt.pipeline].model(): layout.prop(context.scene.dream_textures_project_prompt, 'model') - yield DREAM_PT_dream_panel_projection def get_prompt(context): @@ -125,6 +124,12 @@ def draw(self, context): layout.prop(prompt, "strength") col = layout.column() + + col.prop(context.scene, "dream_textures_project_use_control_net") + if context.scene.dream_textures_project_use_control_net: + col.prop(prompt, "control_net", text="Depth ControlNet") + col.prop(prompt, "controlnet_conditioning_scale") + col.prop(context.scene, "dream_textures_project_bake") if context.scene.dream_textures_project_bake: for obj in context.selected_objects: @@ -150,7 +155,7 @@ def draw(self, context): # Validation try: - prompt.validate(context, task=ModelType.DEPTH) + prompt.validate(context, task=None if context.scene.dream_textures_project_use_control_net else ModelType.DEPTH) _validate_projection(context) except FixItError as e: error_box = layout.box() @@ -258,7 +263,7 @@ class ProjectDreamTexture(bpy.types.Operator): @classmethod def poll(cls, context): try: - context.scene.dream_textures_project_prompt.validate(context, task=ModelType.DEPTH) + context.scene.dream_textures_project_prompt.validate(context, task=None if context.scene.dream_textures_project_use_control_net else ModelType.DEPTH) _validate_projection(context) except: return False @@ -433,11 +438,18 @@ def on_exception(_, exception): raise exception context.scene.dream_textures_info = "Starting..." - future = gen.depth_to_image( - depth=depth, - image=init_img_path, - **context.scene.dream_textures_project_prompt.generate_args() - ) + if context.scene.dream_textures_project_use_control_net: + future = gen.control_net( + control=np.flipud(depth), # the depth control needs to be flipped. + image=init_img_path, + **context.scene.dream_textures_project_prompt.generate_args() + ) + else: + future = gen.depth_to_image( + depth=depth, + image=init_img_path, + **context.scene.dream_textures_project_prompt.generate_args() + ) gen._active_generation_future = future future.call_done_on_exception = False future.add_response_callback(on_response) diff --git a/property_groups/dream_prompt.py b/property_groups/dream_prompt.py index 7fa9c9b6..90e86047 100644 --- a/property_groups/dream_prompt.py +++ b/property_groups/dream_prompt.py @@ -62,8 +62,8 @@ def modify_action_source_type(self, context): ('depth_generated', 'Color and Generated Depth', 'Use MiDaS to infer the depth of the initial image and include it in the conditioning. Can give results that more closely match the composition of the source image', 2), ('depth_map', 'Color and Depth Map', 'Specify a secondary image to use as the depth map. Can give results that closely match the composition of the depth map', 3), ('depth', 'Depth', 'Treat the initial image as a depth map, and ignore any color. Matches the composition of the source image without any color influence', 4), - ('control_net', 'Control Net', 'Treat the initial image as the input to a ControlNet model', 5), - ('control_net_color', 'Color and Control Net', 'Specify a secondary image to use with a ControlNet model', 6), + ('control_net', 'ControlNet', 'Treat the initial image as the input to a ControlNet model', 5), + ('control_net_color', 'Color and ControlNet', 'Specify a secondary image to use with a ControlNet model', 6), ] def model_options(self, context): @@ -130,7 +130,7 @@ def seed_clamp(self, ctx): "pipeline": EnumProperty(name="Pipeline", items=pipeline_options, default=1 if Pipeline.local_available() else 2, description="Specify which model and target should be used."), "model": EnumProperty(name="Model", items=model_options, description="Specify which model to use for inference"), - "control_net": EnumProperty(name="Control Net", items=control_net_options, description="Specify which ControlNet to use"), + "control_net": EnumProperty(name="ControlNet", items=control_net_options, description="Specify which ControlNet to use"), "controlnet_conditioning_scale": FloatProperty(name="ControlNet Conditioning Scale", default=1.0, description="Increases the strength of the ControlNet's effect"), # Prompt From bd8d71b4e82ed872054c3411b0e1d87d75de7084 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Fri, 3 Mar 2023 20:07:02 -0500 Subject: [PATCH 03/28] Add custom node tree and render engine --- __init__.py | 13 +++ classes.py | 31 +++--- engine/__init__.py | 59 +++++++++++ engine/engine.py | 146 +++++++++++++++++++++++++++ engine/node.py | 7 ++ engine/node_executor.py | 31 ++++++ engine/node_tree.py | 10 ++ engine/nodes/input_nodes.py | 77 ++++++++++++++ engine/nodes/pipeline_nodes.py | 178 +++++++++++++++++++++++++++++++++ ui/panels/dream_texture.py | 4 +- 10 files changed, 542 insertions(+), 14 deletions(-) create mode 100644 engine/__init__.py create mode 100644 engine/engine.py create mode 100644 engine/node.py create mode 100644 engine/node_executor.py create mode 100644 engine/node_tree.py create mode 100644 engine/nodes/input_nodes.py create mode 100644 engine/nodes/pipeline_nodes.py diff --git a/__init__.py b/__init__.py index 0f5f8f93..25ab4a57 100644 --- a/__init__.py +++ b/__init__.py @@ -26,6 +26,7 @@ if current_process().name != "__actor__": import bpy from bpy.props import IntProperty, PointerProperty, EnumProperty, BoolProperty, CollectionProperty, FloatProperty + import nodeitems_utils import sys import os @@ -47,6 +48,8 @@ def clear_modules(): from .property_groups.seamless_result import SeamlessResult from .preferences import StableDiffusionPreferences from .ui.presets import register_default_presets + + from . import engine requirements_path_items = ( ('requirements/win-linux-cuda.txt', 'Linux/Windows (CUDA)', 'Linux or Windows with NVIDIA GPU'), @@ -102,12 +105,18 @@ def get_selection_preview(self): bpy.types.Scene.dream_textures_project_framebuffer_arguments = EnumProperty(name="Inputs", items=framebuffer_arguments) bpy.types.Scene.dream_textures_project_bake = BoolProperty(name="Bake", default=False, description="Re-maps the generated texture onto the specified UV map") bpy.types.Scene.dream_textures_project_use_control_net = BoolProperty(name="Use ControlNet", default=False, description="Use a depth ControlNet instead of a depth model") + + bpy.types.Scene.dream_textures_render_engine = PointerProperty(type=engine.DreamTexturesRenderEngineProperties) + + bpy.types.RENDER_PT_context.append(engine.draw_device) for cls in CLASSES: bpy.utils.register_class(cls) for tool in TOOLS: bpy.utils.register_tool(tool) + + engine.register() # Monkey patch cycles render passes register_render_pass() @@ -122,6 +131,10 @@ def unregister(): bpy.utils.unregister_class(cls) for tool in TOOLS: bpy.utils.unregister_tool(tool) + + bpy.types.RENDER_PT_context.remove(engine.draw_device) + + engine.unregister() unregister_render_pass() diff --git a/classes.py b/classes.py index ef980502..9eeb0a22 100644 --- a/classes.py +++ b/classes.py @@ -13,6 +13,8 @@ from .ui.presets import DREAM_PT_AdvancedPresets, DREAM_MT_AdvancedPresets, AddAdvancedPreset, RestoreDefaultPresets +from . import engine + CLASSES = ( *render_properties.render_properties_panels(), @@ -35,6 +37,9 @@ AddAdvancedPreset, NotifyResult, + + engine.DreamTexturesRenderEngine, + *engine.engine_panels(), # The order these are registered in matters *dream_texture.dream_texture_panels(), @@ -44,15 +49,17 @@ ) PREFERENCE_CLASSES = ( - PREFERENCES_UL_ModelList, - ModelSearch, - InstallModel, - Model, - DreamPrompt, - SeamlessResult, - UninstallDependencies, - InstallDependencies, - OpenURL, - ImportWeights, - RestoreDefaultPresets, - StableDiffusionPreferences) \ No newline at end of file + PREFERENCES_UL_ModelList, + ModelSearch, + InstallModel, + Model, + DreamPrompt, + SeamlessResult, + UninstallDependencies, + InstallDependencies, + OpenURL, + ImportWeights, + RestoreDefaultPresets, + StableDiffusionPreferences, + engine.DreamTexturesRenderEngineProperties, +) \ No newline at end of file diff --git a/engine/__init__.py b/engine/__init__.py new file mode 100644 index 00000000..e0ad2d6b --- /dev/null +++ b/engine/__init__.py @@ -0,0 +1,59 @@ +from .engine import * +from .node_tree import * +from .node_executor import * +from .node import * +from .nodes.input_nodes import * +from .nodes.pipeline_nodes import * + +import bpy +import nodeitems_utils + +class DreamTexturesNodeCategory(nodeitems_utils.NodeCategory): + @classmethod + def poll(cls, context): + return context.space_data.tree_type == DreamTexturesNodeTree.__name__ + +categories = [ + DreamTexturesNodeCategory("DREAM_TEXTURES_PIPELINE", "Pipeline", items = [ + nodeitems_utils.NodeItem(NodeStableDiffusion.bl_idname), + ]), + DreamTexturesNodeCategory("DREAM_TEXTURES_INPUT", "Input", items = [ + nodeitems_utils.NodeItem(NodeInteger.bl_idname), + nodeitems_utils.NodeItem(NodeString.bl_idname), + nodeitems_utils.NodeItem(NodeCollection.bl_idname), + nodeitems_utils.NodeItem(NodeRandomValue.bl_idname), + ]), + DreamTexturesNodeCategory("DREAM_TEXTURES_GROUP", "Group", items = [ + nodeitems_utils.NodeItem(bpy.types.NodeGroupOutput.__name__), + ]), +] + +def register(): + bpy.utils.register_class(DreamTexturesNodeTree) + + # Nodes + bpy.utils.register_class(NodeSocketControlNet) + bpy.utils.register_class(NodeStableDiffusion) + bpy.utils.register_class(NodeControlNet) + + bpy.utils.register_class(NodeInteger) + bpy.utils.register_class(NodeString) + bpy.utils.register_class(NodeCollection) + bpy.utils.register_class(NodeRandomValue) + + nodeitems_utils.register_node_categories("DREAM_TEXTURES_CATEGORIES", categories) + +def unregister(): + bpy.utils.unregister_class(DreamTexturesNodeTree) + + # Nodes + bpy.utils.unregister_class(NodeSocketControlNet) + bpy.utils.unregister_class(NodeStableDiffusion) + bpy.utils.unregister_class(NodeControlNet) + + bpy.utils.unregister_class(NodeInteger) + bpy.utils.unregister_class(NodeString) + bpy.utils.unregister_class(NodeCollection) + bpy.utils.unregister_class(NodeRandomValue) + + nodeitems_utils.unregister_node_categories("DREAM_TEXTURES_CATEGORIES") \ No newline at end of file diff --git a/engine/engine.py b/engine/engine.py new file mode 100644 index 00000000..028e6bc4 --- /dev/null +++ b/engine/engine.py @@ -0,0 +1,146 @@ +import bpy +import gpu +from bl_ui.properties_render import RenderButtonsPanel +from bl_ui.properties_output import RenderOutputButtonsPanel +from ..ui.panels.dream_texture import create_panel, prompt_panel, advanced_panel, size_panel +from ..property_groups.dream_prompt import control_net_options + +class DreamTexturesRenderEngine(bpy.types.RenderEngine): + """A custom Dream Textures render engine, that uses Stable Diffusion and scene data to render images, instead of as a pass on top of Cycles.""" + + bl_idname = "DREAM_TEXTURES" + bl_label = "Dream Textures" + bl_use_preview = False + + def __init__(self): + pass + + def __del__(self): + pass + + def render(self, depsgraph): + scene = depsgraph.scene + scale = scene.render.resolution_percentage / 100.0 + self.size_x = int(scene.render.resolution_x * scale) + self.size_y = int(scene.render.resolution_y * scale) + + # Fill the render result with a flat color. The framebuffer is + # defined as a list of pixels, each pixel itself being a list of + # R,G,B,A values. + if self.is_preview: + color = [0.1, 0.2, 0.1, 1.0] + else: + color = [0.2, 0.1, 0.1, 1.0] + + pixel_count = self.size_x * self.size_y + rect = [color] * pixel_count + + # Here we write the pixel values to the RenderResult + result = self.begin_result(0, 0, self.size_x, self.size_y) + layer = result.layers[0].passes["Combined"] + layer.rect = rect + self.end_result(result) + + def view_update(self, context, depsgraph): + region = context.region + view3d = context.space_data + scene = depsgraph.scene + + # Get viewport dimensions + dimensions = region.width, region.height + + if not self.scene_data: + # First time initialization + self.scene_data = [] + first_time = True + + # Loop over all datablocks used in the scene. + for datablock in depsgraph.ids: + pass + else: + first_time = False + + # Test which datablocks changed + for update in depsgraph.updates: + print("Datablock updated: ", update.id.name) + + # Test if any material was added, removed or changed. + if depsgraph.id_type_updated('MATERIAL'): + print("Materials updated") + + # Loop over all object instances in the scene. + if first_time or depsgraph.id_type_updated('OBJECT'): + for instance in depsgraph.object_instances: + pass + + # For viewport renders, this method is called whenever Blender redraws + # the 3D viewport. The renderer is expected to quickly draw the render + # with OpenGL, and not perform other expensive work. + # Blender will draw overlays for selection and editing on top of the + # rendered image automatically. + def view_draw(self, context, depsgraph): + region = context.region + scene = depsgraph.scene + + # Get viewport dimensions + dimensions = region.width, region.height + + # Bind shader that converts from scene linear to display space, + gpu.state.blend_set('ALPHA_PREMULT') + self.bind_display_space_shader(scene) + + if not self.draw_data or self.draw_data.dimensions != dimensions: + self.draw_data = CustomDrawData(dimensions) + + self.draw_data.draw() + + self.unbind_display_space_shader() + gpu.state.blend_set('NONE') + +def draw_device(self, context): + scene = context.scene + layout = self.layout + layout.use_property_split = True + layout.use_property_decorate = False + + if context.engine == DreamTexturesRenderEngine.bl_idname: + layout.prop(scene.dream_textures_prompt, "pipeline") + layout.prop(scene.dream_textures_prompt, "model") + +class DreamTexturesRenderEngineProperties(bpy.types.PropertyGroup): + depth_controlnet: bpy.props.EnumProperty(name="Depth ControlNet", items=control_net_options, description="Select a depth map ControlNet") + armature_controlnet: bpy.props.EnumProperty(name="Armature ControlNet", items=control_net_options, description="Select an OpenPose ControlNet") + normals_controlnet: bpy.props.EnumProperty(name="Normals ControlNet", items=control_net_options, description="Select a normal map ControlNet") + +class ControlsPanel(bpy.types.Panel, RenderButtonsPanel): + COMPAT_ENGINES = {DreamTexturesRenderEngine.bl_idname} + bl_label = "Controls" + bl_idname = f"DREAM_PT_render_engine_controls" + + def draw(self, context): + self.layout.use_property_split = True + self.layout.prop(context.scene.dream_textures_render_engine, "depth_controlnet", icon="OBJECT_DATA") + self.layout.prop(context.scene.dream_textures_render_engine, "armature_controlnet", icon="ARMATURE_DATA") + self.layout.prop(context.scene.dream_textures_render_engine, "normals_controlnet", icon="NORMALS_FACE") + +def engine_panels(): + bpy.types.RENDER_PT_output.COMPAT_ENGINES.add(DreamTexturesRenderEngine.bl_idname) + def get_prompt(context): + return context.scene.dream_textures_prompt + class RenderPanel(bpy.types.Panel, RenderButtonsPanel): + COMPAT_ENGINES = {DreamTexturesRenderEngine.bl_idname} + def draw(self, context): + self.layout.use_property_decorate = True + class OutputPanel(bpy.types.Panel, RenderOutputButtonsPanel): + COMPAT_ENGINES = {DreamTexturesRenderEngine.bl_idname} + + def draw(self, context): + self.layout.use_property_decorate = True + + # Render Properties + yield from prompt_panel(RenderPanel, 'engine', get_prompt) + yield ControlsPanel + yield from advanced_panel(RenderPanel, 'engine', get_prompt) + + # Output Properties + yield size_panel(OutputPanel, 'engine', get_prompt) \ No newline at end of file diff --git a/engine/node.py b/engine/node.py new file mode 100644 index 00000000..18113f09 --- /dev/null +++ b/engine/node.py @@ -0,0 +1,7 @@ +import bpy +from .node_tree import DreamTexturesNodeTree + +class DreamTexturesNode(bpy.types.Node): + @classmethod + def poll(cls, tree): + return tree.bl_idname == DreamTexturesNodeTree.__name__ \ No newline at end of file diff --git a/engine/node_executor.py b/engine/node_executor.py new file mode 100644 index 00000000..b11607a1 --- /dev/null +++ b/engine/node_executor.py @@ -0,0 +1,31 @@ +import bpy +import numpy as np +# from dream_textures.engine import node_executor +# node_executor.execute(bpy.data.node_groups["NodeTree"], bpy.context) + +def execute_node(node, context, cache): + if node in cache: + return cache[node] + kwargs = { + input.name.lower().replace(' ', '_'): ([ + execute_node(link.from_socket.node, context, cache)[link.from_socket.name] + for link in input.links + ] if len(input.links) > 1 else execute_node(input.links[0].from_socket.node, context, cache)[input.links[0].from_socket.name]) + if input.is_linked else getattr(input, 'default_value', None) + for input in node.inputs + } + if node.type == 'GROUP_OUTPUT': + return list(kwargs.values())[0] + result = node.execute(context, **kwargs) + print(node.name, result) + cache[node] = result + return result + +def execute(node_tree, context): + output = next(n for n in node_tree.nodes if n.type == 'GROUP_OUTPUT') + cache = {} + result = execute_node(output, context, cache) + print(result) + image = bpy.data.images.new("test", width=result.shape[0], height=result.shape[1]) + image.pixels.foreach_set(result.ravel()) + return image \ No newline at end of file diff --git a/engine/node_tree.py b/engine/node_tree.py new file mode 100644 index 00000000..e5daea02 --- /dev/null +++ b/engine/node_tree.py @@ -0,0 +1,10 @@ +import bpy +from .engine import DreamTexturesRenderEngine + +class DreamTexturesNodeTree(bpy.types.NodeTree): + bl_label = "Dream Textures Node Editor" + bl_icon = 'NODETREE' + + @classmethod + def poll(cls, context): + return context.scene.render.engine == DreamTexturesRenderEngine.bl_idname \ No newline at end of file diff --git a/engine/nodes/input_nodes.py b/engine/nodes/input_nodes.py new file mode 100644 index 00000000..9f1b6ea9 --- /dev/null +++ b/engine/nodes/input_nodes.py @@ -0,0 +1,77 @@ +import bpy +import nodeitems_utils +from ..node import DreamTexturesNode +import random + +class NodeString(DreamTexturesNode): + bl_idname = "dream_textures.node_string" + bl_label = "String" + + value: bpy.props.StringProperty(name="") + + def init(self, context): + self.outputs.new("NodeSocketString", "String") + + def draw_buttons(self, context, layout): + layout.prop(self, "value") + + def execute(self, context): + return { + 'String': self.value + } + +class NodeInteger(DreamTexturesNode): + bl_idname = "dream_textures.node_integer" + bl_label = "Integer" + + value: bpy.props.IntProperty(name="") + + def init(self, context): + self.outputs.new("NodeSocketInt", "Integer") + + def draw_buttons(self, context, layout): + layout.prop(self, "value") + + def execute(self, context): + return { + 'Integer': self.value + } + +class NodeCollection(DreamTexturesNode): + bl_idname = "dream_textures.node_collection" + bl_label = "Collection" + + value: bpy.props.PointerProperty(type=bpy.types.Collection, name="") + + def init(self, context): + self.outputs.new("NodeSocketCollection", "Collection") + + def draw_buttons(self, context, layout): + layout.prop(self, "value") + + def execute(self, context): + return { + 'Collection': self.value + } + +class NodeRandomValue(DreamTexturesNode): + bl_idname = "dream_textures.node_random_value" + bl_label = "Random Value" + + data_type: bpy.props.EnumProperty(name="", items=( + ('integer', 'Integer', '', 1), + )) + + def init(self, context): + self.inputs.new("NodeSocketInt", "Min") + self.inputs.new("NodeSocketInt", "Max") + + self.outputs.new("NodeSocketInt", "Value") + + def draw_buttons(self, context, layout): + layout.prop(self, "data_type") + + def execute(self, context, min, max): + return { + 'Value': random.randrange(min, max) + } \ No newline at end of file diff --git a/engine/nodes/pipeline_nodes.py b/engine/nodes/pipeline_nodes.py new file mode 100644 index 00000000..9aef1fef --- /dev/null +++ b/engine/nodes/pipeline_nodes.py @@ -0,0 +1,178 @@ +import bpy +from ..node import DreamTexturesNode +from ...generator_process import Generator +from ...generator_process.actions.prompt_to_image import StepPreviewMode +from ...property_groups.dream_prompt import DreamPrompt, control_net_options +import numpy as np +from dataclasses import dataclass +from typing import Any +import enum + +class NodeSocketControlNet(bpy.types.NodeSocket): + bl_idname = "NodeSocketControlNet" + bl_label = "ControlNet Socket" + + def __init__(self): + self.link_limit = 0 + + def draw(self, context, layout, node, text): + layout.label(text=text) + + def draw_color(self, context, node): + return (0.63, 0.63, 0.63, 1) + +class ControlType(enum.IntEnum): + DEPTH = 1 + OPENPOSE = 2 + NORMAL = 3 + +@dataclass +class ControlNet: + model: str + image: Any + collection: Any + control_type: ControlType + conditioning_scale: float + +def _update_stable_diffusion_sockets(self, context): + self.inputs['Source Image'].enabled = self.task in {'image_to_image', 'depth_to_image'} + self.inputs['Noise Strength'].enabled = self.task in {'image_to_image', 'depth_to_image'} + self.inputs['Depth Map'].enabled = self.task == 'depth_to_image' + self.inputs['ControlNets'].enabled = self.task != 'depth_to_image' +class NodeStableDiffusion(DreamTexturesNode): + bl_idname = "dream_textures.node_stable_diffusion" + bl_label = "Stable Diffusion" + + prompt: bpy.props.PointerProperty(type=DreamPrompt) + task: bpy.props.EnumProperty(name="", items=( + ('prompt_to_image', 'Prompt to Image', '', 1), + ('image_to_image', 'Image to Image', '', 2), + ('depth_to_image', 'Depth to Image', '', 3), + ), update=_update_stable_diffusion_sockets) + + def init(self, context): + self.inputs.new("NodeSocketColor", "Depth Map") + self.inputs.new("NodeSocketColor", "Source Image") + self.inputs.new("NodeSocketFloat", "Noise Strength").default_value = 0.75 + + self.inputs.new("NodeSocketString", "Prompt") + self.inputs.new("NodeSocketString", "Negative Prompt") + + self.inputs.new("NodeSocketInt", "Width").default_value = 512 + self.inputs.new("NodeSocketInt", "Height").default_value = 512 + + self.inputs.new("NodeSocketInt", "Steps").default_value = 25 + self.inputs.new("NodeSocketInt", "Seed") + self.inputs.new("NodeSocketFloat", "CFG Scale").default_value = 7.50 + + self.inputs.new("NodeSocketControlNet", "ControlNets") + + self.outputs.new("NodeSocketColor", "Image") + self.outputs.new("NodeSocketInt", "Seed") + + _update_stable_diffusion_sockets(self, context) + + def draw_buttons(self, context, layout): + layout.prop(self, "task") + prompt = self.prompt + layout.prop(prompt, "pipeline", text="") + layout.prop(prompt, "model", text="") + layout.prop(prompt, "scheduler", text="") + + def execute(self, context, prompt, negative_prompt, width, height, steps, seed, cfg_scale, controlnets, depth_map, source_image, noise_strength): + self.prompt.use_negative_prompt = True + self.prompt.negative_prompt = negative_prompt + self.prompt.steps = steps + self.prompt.seed = str(seed) + self.prompt.cfg_scale = cfg_scale + args = self.prompt.generate_args() + + match self.task: + case 'prompt_to_image': + result = Generator.shared().prompt_to_image( + pipeline=args['pipeline'], + model=args['model'], + scheduler=args['scheduler'], + optimizations=args['optimizations'], + seamless_axes=args['seamless_axes'], + iterations=args['iterations'], + prompt=prompt, + steps=steps, + seed=seed, + width=width, + height=height, + cfg_scale=cfg_scale, + use_negative_prompt=True, + negative_prompt=negative_prompt, + step_preview_mode=StepPreviewMode.NONE + ).result() + case 'image_to_image': + result = Generator.shared().image_to_image( + pipeline=args['pipeline'], + model=args['model'], + scheduler=args['scheduler'], + optimizations=args['optimizations'], + seamless_axes=args['seamless_axes'], + iterations=args['iterations'], + + image=np.uint8(source_image * 255), + strength=noise_strength, + fit=True, + + prompt=prompt, + steps=steps, + seed=seed, + width=width, + height=height, + cfg_scale=cfg_scale, + use_negative_prompt=True, + negative_prompt=negative_prompt, + step_preview_mode=StepPreviewMode.NONE + ).result() + return { + 'Image': result[-1].images[-1], + 'Seed': result[-1].seeds[-1] + } + +def _update_control_net_sockets(self, context): + self.inputs['Collection'].enabled = self.input_type == 'collection' + self.inputs['Image'].enabled = self.input_type == 'image' +class NodeControlNet(DreamTexturesNode): + bl_idname = "dream_textures.node_control_net" + bl_label = "ControlNet" + + control_net: bpy.props.EnumProperty(name="", items=control_net_options) + input_type: bpy.props.EnumProperty(name="", items=( + ('collection', 'Collection', '', 1), + ('image', 'Image', '', 2), + ), update=_update_control_net_sockets) + control_type: bpy.props.EnumProperty(name="", items=( + ('DEPTH', 'Depth', '', 1), + ('OPENPOSE', 'OpenPose', '', 2), + ('NORMAL', 'Normal Map', '', 3), + )) + + def init(self, context): + self.inputs.new("NodeSocketCollection", "Collection") + self.inputs.new("NodeSocketColor", "Image") + self.inputs.new("NodeSocketFloat", "Conditioning Scale").default_value = 1 + + self.outputs.new(NodeSocketControlNet.bl_idname, "Control") + + _update_control_net_sockets(self, context) + + def draw_buttons(self, context, layout): + layout.prop(self, "control_net") + layout.prop(self, "input_type") + layout.prop(self, "control_type") + + def execute(self, context, collection, image, conditioning_scale): + return { + 'Control': ControlNet( + self.control_net, + image if self.input_type == 'image' else None, + collection if self.input_type == 'collection' else None, + ControlType[self.control_type], + conditioning_scale + ) + } \ No newline at end of file diff --git a/ui/panels/dream_texture.py b/ui/panels/dream_texture.py index 496726cf..8dbf63a9 100644 --- a/ui/panels/dream_texture.py +++ b/ui/panels/dream_texture.py @@ -95,8 +95,8 @@ class SubPanel(BasePanel): def draw(self, context): self.layout.use_property_decorate = use_property_decorate - - return ctor(SubPanel, space_type, get_prompt, **kwargs) + + return ctor(kwargs.pop('base_panel', SubPanel), space_type, get_prompt, **kwargs) def prompt_panel(sub_panel, space_type, get_prompt, get_seamless_result=None): class PromptPanel(sub_panel): From 89ed7e50711037d117075e2d29eb6ec99b62f1ba Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Fri, 3 Mar 2023 20:37:23 -0500 Subject: [PATCH 04/28] Add node_tree property to render engine --- __init__.py | 10 +++++----- classes.py | 2 +- engine/engine.py | 23 +++++------------------ engine/node_tree.py | 4 ++-- 4 files changed, 13 insertions(+), 26 deletions(-) diff --git a/__init__.py b/__init__.py index 25ab4a57..ec2ca87a 100644 --- a/__init__.py +++ b/__init__.py @@ -105,18 +105,18 @@ def get_selection_preview(self): bpy.types.Scene.dream_textures_project_framebuffer_arguments = EnumProperty(name="Inputs", items=framebuffer_arguments) bpy.types.Scene.dream_textures_project_bake = BoolProperty(name="Bake", default=False, description="Re-maps the generated texture onto the specified UV map") bpy.types.Scene.dream_textures_project_use_control_net = BoolProperty(name="Use ControlNet", default=False, description="Use a depth ControlNet instead of a depth model") - - bpy.types.Scene.dream_textures_render_engine = PointerProperty(type=engine.DreamTexturesRenderEngineProperties) - bpy.types.RENDER_PT_context.append(engine.draw_device) + engine.register() for cls in CLASSES: bpy.utils.register_class(cls) for tool in TOOLS: bpy.utils.register_tool(tool) - - engine.register() + + bpy.types.Scene.dream_textures_render_engine = PointerProperty(type=engine.DreamTexturesRenderEngineProperties) + + bpy.types.RENDER_PT_context.append(engine.draw_device) # Monkey patch cycles render passes register_render_pass() diff --git a/classes.py b/classes.py index 9eeb0a22..aa3cfd79 100644 --- a/classes.py +++ b/classes.py @@ -38,6 +38,7 @@ NotifyResult, + engine.DreamTexturesRenderEngineProperties, engine.DreamTexturesRenderEngine, *engine.engine_panels(), @@ -61,5 +62,4 @@ ImportWeights, RestoreDefaultPresets, StableDiffusionPreferences, - engine.DreamTexturesRenderEngineProperties, ) \ No newline at end of file diff --git a/engine/engine.py b/engine/engine.py index 028e6bc4..ad6dfc24 100644 --- a/engine/engine.py +++ b/engine/engine.py @@ -4,6 +4,7 @@ from bl_ui.properties_output import RenderOutputButtonsPanel from ..ui.panels.dream_texture import create_panel, prompt_panel, advanced_panel, size_panel from ..property_groups.dream_prompt import control_net_options +from .node_tree import DreamTexturesNodeTree class DreamTexturesRenderEngine(bpy.types.RenderEngine): """A custom Dream Textures render engine, that uses Stable Diffusion and scene data to render images, instead of as a pass on top of Cycles.""" @@ -104,24 +105,12 @@ def draw_device(self, context): layout.use_property_decorate = False if context.engine == DreamTexturesRenderEngine.bl_idname: - layout.prop(scene.dream_textures_prompt, "pipeline") - layout.prop(scene.dream_textures_prompt, "model") + layout.template_ID(scene.dream_textures_render_engine, "node_tree", text="Node Tree") +def _poll_node_tree(self, value): + return value.bl_idname == "dream_textures.node_tree" class DreamTexturesRenderEngineProperties(bpy.types.PropertyGroup): - depth_controlnet: bpy.props.EnumProperty(name="Depth ControlNet", items=control_net_options, description="Select a depth map ControlNet") - armature_controlnet: bpy.props.EnumProperty(name="Armature ControlNet", items=control_net_options, description="Select an OpenPose ControlNet") - normals_controlnet: bpy.props.EnumProperty(name="Normals ControlNet", items=control_net_options, description="Select a normal map ControlNet") - -class ControlsPanel(bpy.types.Panel, RenderButtonsPanel): - COMPAT_ENGINES = {DreamTexturesRenderEngine.bl_idname} - bl_label = "Controls" - bl_idname = f"DREAM_PT_render_engine_controls" - - def draw(self, context): - self.layout.use_property_split = True - self.layout.prop(context.scene.dream_textures_render_engine, "depth_controlnet", icon="OBJECT_DATA") - self.layout.prop(context.scene.dream_textures_render_engine, "armature_controlnet", icon="ARMATURE_DATA") - self.layout.prop(context.scene.dream_textures_render_engine, "normals_controlnet", icon="NORMALS_FACE") + node_tree: bpy.props.PointerProperty(type=DreamTexturesNodeTree, name="Node Tree", poll=_poll_node_tree) def engine_panels(): bpy.types.RENDER_PT_output.COMPAT_ENGINES.add(DreamTexturesRenderEngine.bl_idname) @@ -138,8 +127,6 @@ def draw(self, context): self.layout.use_property_decorate = True # Render Properties - yield from prompt_panel(RenderPanel, 'engine', get_prompt) - yield ControlsPanel yield from advanced_panel(RenderPanel, 'engine', get_prompt) # Output Properties diff --git a/engine/node_tree.py b/engine/node_tree.py index e5daea02..4fafa951 100644 --- a/engine/node_tree.py +++ b/engine/node_tree.py @@ -1,10 +1,10 @@ import bpy -from .engine import DreamTexturesRenderEngine class DreamTexturesNodeTree(bpy.types.NodeTree): + bl_idname = "dream_textures.node_tree" bl_label = "Dream Textures Node Editor" bl_icon = 'NODETREE' @classmethod def poll(cls, context): - return context.scene.render.engine == DreamTexturesRenderEngine.bl_idname \ No newline at end of file + return context.scene.render.engine == 'DREAM_TEXTURES' \ No newline at end of file From d6c1089ca038531b918b8e9ae09b8d97166ae7b1 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Sun, 12 Mar 2023 22:12:17 -0400 Subject: [PATCH 05/28] Node tree improvements --- __init__.py | 4 +- classes.py | 1 + engine/__init__.py | 17 ++ engine/engine.py | 86 +++++++--- engine/node.py | 2 +- engine/node_executor.py | 13 +- engine/node_tree.py | 2 +- engine/nodes/input_nodes.py | 88 ++++++++-- engine/nodes/pipeline_nodes.py | 161 +++++++++++++------ engine/nodes/utility_nodes.py | 65 ++++++++ generator_process/actions/control_net.py | 1 + generator_process/actions/huggingface_hub.py | 115 +++++++------ generator_process/actions/prompt_to_image.py | 2 + ui/panels/dream_texture.py | 7 +- 14 files changed, 410 insertions(+), 154 deletions(-) create mode 100644 engine/nodes/utility_nodes.py diff --git a/__init__.py b/__init__.py index ec2ca87a..caa39279 100644 --- a/__init__.py +++ b/__init__.py @@ -25,8 +25,7 @@ if current_process().name != "__actor__": import bpy - from bpy.props import IntProperty, PointerProperty, EnumProperty, BoolProperty, CollectionProperty, FloatProperty - import nodeitems_utils + from bpy.props import IntProperty, PointerProperty, EnumProperty, BoolProperty, CollectionProperty import sys import os @@ -46,7 +45,6 @@ def clear_modules(): from .operators.dream_texture import DreamTexture, kill_generator from .property_groups.dream_prompt import DreamPrompt from .property_groups.seamless_result import SeamlessResult - from .preferences import StableDiffusionPreferences from .ui.presets import register_default_presets from . import engine diff --git a/classes.py b/classes.py index aa3cfd79..83ea7af3 100644 --- a/classes.py +++ b/classes.py @@ -40,6 +40,7 @@ engine.DreamTexturesRenderEngineProperties, engine.DreamTexturesRenderEngine, + engine.NewEngineNodeTree, *engine.engine_panels(), # The order these are registered in matters diff --git a/engine/__init__.py b/engine/__init__.py index e0ad2d6b..9fdb32ff 100644 --- a/engine/__init__.py +++ b/engine/__init__.py @@ -4,6 +4,7 @@ from .node import * from .nodes.input_nodes import * from .nodes.pipeline_nodes import * +from .nodes.utility_nodes import * import bpy import nodeitems_utils @@ -20,7 +21,12 @@ def poll(cls, context): DreamTexturesNodeCategory("DREAM_TEXTURES_INPUT", "Input", items = [ nodeitems_utils.NodeItem(NodeInteger.bl_idname), nodeitems_utils.NodeItem(NodeString.bl_idname), + nodeitems_utils.NodeItem(NodeImage.bl_idname), nodeitems_utils.NodeItem(NodeCollection.bl_idname), + nodeitems_utils.NodeItem(NodeSceneInfo.bl_idname), + ]), + DreamTexturesNodeCategory("DREAM_TEXTURES_UTILITY", "Utilities", items = [ + nodeitems_utils.NodeItem(NodeMath.bl_idname), nodeitems_utils.NodeItem(NodeRandomValue.bl_idname), ]), DreamTexturesNodeCategory("DREAM_TEXTURES_GROUP", "Group", items = [ @@ -29,6 +35,8 @@ def poll(cls, context): ] def register(): + bpy.types.Scene.dream_textures_engine_prompt = bpy.props.PointerProperty(type=DreamPrompt) + bpy.utils.register_class(DreamTexturesNodeTree) # Nodes @@ -39,6 +47,10 @@ def register(): bpy.utils.register_class(NodeInteger) bpy.utils.register_class(NodeString) bpy.utils.register_class(NodeCollection) + bpy.utils.register_class(NodeSceneInfo) + bpy.utils.register_class(NodeImage) + + bpy.utils.register_class(NodeMath) bpy.utils.register_class(NodeRandomValue) nodeitems_utils.register_node_categories("DREAM_TEXTURES_CATEGORIES", categories) @@ -54,6 +66,11 @@ def unregister(): bpy.utils.unregister_class(NodeInteger) bpy.utils.unregister_class(NodeString) bpy.utils.unregister_class(NodeCollection) + bpy.utils.unregister_class(NodeSceneInfo) + bpy.utils.unregister_class(NodeImage) + + bpy.utils.unregister_class(NodeMath) bpy.utils.unregister_class(NodeRandomValue) + nodeitems_utils.unregister_node_categories("DREAM_TEXTURES_CATEGORIES") \ No newline at end of file diff --git a/engine/engine.py b/engine/engine.py index ad6dfc24..9af89f20 100644 --- a/engine/engine.py +++ b/engine/engine.py @@ -2,9 +2,10 @@ import gpu from bl_ui.properties_render import RenderButtonsPanel from bl_ui.properties_output import RenderOutputButtonsPanel -from ..ui.panels.dream_texture import create_panel, prompt_panel, advanced_panel, size_panel -from ..property_groups.dream_prompt import control_net_options +import numpy as np +from ..ui.panels.dream_texture import optimization_panels from .node_tree import DreamTexturesNodeTree +from ..engine import node_executor class DreamTexturesRenderEngine(bpy.types.RenderEngine): """A custom Dream Textures render engine, that uses Stable Diffusion and scene data to render images, instead of as a pass on top of Cycles.""" @@ -12,6 +13,7 @@ class DreamTexturesRenderEngine(bpy.types.RenderEngine): bl_idname = "DREAM_TEXTURES" bl_label = "Dream Textures" bl_use_preview = False + bl_use_gpu_context = True def __init__(self): pass @@ -21,25 +23,40 @@ def __del__(self): def render(self, depsgraph): scene = depsgraph.scene - scale = scene.render.resolution_percentage / 100.0 - self.size_x = int(scene.render.resolution_x * scale) - self.size_y = int(scene.render.resolution_y * scale) - - # Fill the render result with a flat color. The framebuffer is - # defined as a list of pixels, each pixel itself being a list of - # R,G,B,A values. - if self.is_preview: - color = [0.1, 0.2, 0.1, 1.0] - else: - color = [0.2, 0.1, 0.1, 1.0] - - pixel_count = self.size_x * self.size_y - rect = [color] * pixel_count - # Here we write the pixel values to the RenderResult - result = self.begin_result(0, 0, self.size_x, self.size_y) + def prepare_result(result): + if len(result.shape) == 2: + return np.concatenate( + ( + np.stack((result,)*3, axis=-1), + np.ones((*result.shape, 1)) + ), + axis=-1 + ) + else: + return result + + result = self.begin_result(0, 0, scene.render.resolution_x, scene.render.resolution_y) layer = result.layers[0].passes["Combined"] - layer.rect = rect + + try: + progress = 0 + def update_result(node, result): + nonlocal progress + progress += 1 + if isinstance(result, np.ndarray): + node_result = prepare_result(result) + layer.rect = node_result.reshape(-1, node_result.shape[-1]) + self.update_result(result) + self.update_stats("Node", node.name) + self.update_progress(progress / len(scene.dream_textures_render_engine.node_tree.nodes)) + node_result = node_executor.execute(scene.dream_textures_render_engine.node_tree, depsgraph, on_execute=update_result) + node_result = prepare_result(node_result) + except Exception as error: + self.report({'ERROR'}, str(error)) + raise error + + layer.rect = node_result.reshape(-1, node_result.shape[-1]) self.end_result(result) def view_update(self, context, depsgraph): @@ -98,6 +115,14 @@ def view_draw(self, context, depsgraph): self.unbind_display_space_shader() gpu.state.blend_set('NONE') +class NewEngineNodeTree(bpy.types.Operator): + bl_idname = "dream_textures.new_engine_node_tree" + bl_label = "New Node Tree" + + def execute(self, context): + bpy.ops.node.new_node_tree(type="DreamTexturesNodeTree") + return {'FINISHED'} + def draw_device(self, context): scene = context.scene layout = self.layout @@ -105,17 +130,17 @@ def draw_device(self, context): layout.use_property_decorate = False if context.engine == DreamTexturesRenderEngine.bl_idname: - layout.template_ID(scene.dream_textures_render_engine, "node_tree", text="Node Tree") + layout.template_ID(scene.dream_textures_render_engine, "node_tree", text="Node Tree", new=NewEngineNodeTree.bl_idname) def _poll_node_tree(self, value): - return value.bl_idname == "dream_textures.node_tree" + return value.bl_idname == "DreamTexturesNodeTree" class DreamTexturesRenderEngineProperties(bpy.types.PropertyGroup): node_tree: bpy.props.PointerProperty(type=DreamTexturesNodeTree, name="Node Tree", poll=_poll_node_tree) def engine_panels(): bpy.types.RENDER_PT_output.COMPAT_ENGINES.add(DreamTexturesRenderEngine.bl_idname) def get_prompt(context): - return context.scene.dream_textures_prompt + return context.scene.dream_textures_engine_prompt class RenderPanel(bpy.types.Panel, RenderButtonsPanel): COMPAT_ENGINES = {DreamTexturesRenderEngine.bl_idname} def draw(self, context): @@ -127,7 +152,20 @@ def draw(self, context): self.layout.use_property_decorate = True # Render Properties - yield from advanced_panel(RenderPanel, 'engine', get_prompt) + yield from optimization_panels(RenderPanel, 'engine', get_prompt, "") # Output Properties - yield size_panel(OutputPanel, 'engine', get_prompt) \ No newline at end of file + class FormatPanel(OutputPanel): + """Create a subpanel for format options""" + bl_idname = f"DREAM_PT_dream_panel_format_engine" + bl_label = "Format" + + def draw(self, context): + super().draw(context) + layout = self.layout + layout.use_property_split = True + + col = layout.column(align=True) + col.prop(context.scene.render, "resolution_x") + col.prop(context.scene.render, "resolution_y", text="Y") + yield FormatPanel \ No newline at end of file diff --git a/engine/node.py b/engine/node.py index 18113f09..b17afb93 100644 --- a/engine/node.py +++ b/engine/node.py @@ -4,4 +4,4 @@ class DreamTexturesNode(bpy.types.Node): @classmethod def poll(cls, tree): - return tree.bl_idname == DreamTexturesNodeTree.__name__ \ No newline at end of file + return tree.bl_idname == DreamTexturesNodeTree.bl_idname \ No newline at end of file diff --git a/engine/node_executor.py b/engine/node_executor.py index b11607a1..cb45feaa 100644 --- a/engine/node_executor.py +++ b/engine/node_executor.py @@ -3,7 +3,7 @@ # from dream_textures.engine import node_executor # node_executor.execute(bpy.data.node_groups["NodeTree"], bpy.context) -def execute_node(node, context, cache): +def execute_node(node, context, cache, on_execute=lambda _, __: None): if node in cache: return cache[node] kwargs = { @@ -17,15 +17,12 @@ def execute_node(node, context, cache): if node.type == 'GROUP_OUTPUT': return list(kwargs.values())[0] result = node.execute(context, **kwargs) - print(node.name, result) + on_execute(node, result) cache[node] = result return result -def execute(node_tree, context): +def execute(node_tree, context, on_execute=lambda _, __: None): output = next(n for n in node_tree.nodes if n.type == 'GROUP_OUTPUT') cache = {} - result = execute_node(output, context, cache) - print(result) - image = bpy.data.images.new("test", width=result.shape[0], height=result.shape[1]) - image.pixels.foreach_set(result.ravel()) - return image \ No newline at end of file + result = execute_node(output, context, cache, on_execute) + return result \ No newline at end of file diff --git a/engine/node_tree.py b/engine/node_tree.py index 4fafa951..812ba51d 100644 --- a/engine/node_tree.py +++ b/engine/node_tree.py @@ -1,7 +1,7 @@ import bpy class DreamTexturesNodeTree(bpy.types.NodeTree): - bl_idname = "dream_textures.node_tree" + bl_idname = "DreamTexturesNodeTree" bl_label = "Dream Textures Node Editor" bl_icon = 'NODETREE' diff --git a/engine/nodes/input_nodes.py b/engine/nodes/input_nodes.py index 9f1b6ea9..a54d1404 100644 --- a/engine/nodes/input_nodes.py +++ b/engine/nodes/input_nodes.py @@ -1,7 +1,9 @@ import bpy -import nodeitems_utils -from ..node import DreamTexturesNode +import gpu +from gpu_extras.batch import batch_for_shader import random +import numpy as np +from ..node import DreamTexturesNode class NodeString(DreamTexturesNode): bl_idname = "dream_textures.node_string" @@ -54,24 +56,80 @@ def execute(self, context): 'Collection': self.value } -class NodeRandomValue(DreamTexturesNode): - bl_idname = "dream_textures.node_random_value" - bl_label = "Random Value" +class NodeImage(DreamTexturesNode): + bl_idname = "dream_textures.node_image" + bl_label = "Image" + + value: bpy.props.PointerProperty(type=bpy.types.Image) + + def init(self, context): + self.outputs.new("NodeSocketImage", "Image") + + def draw_buttons(self, context, layout): + layout.prop(self, "value", text="") + + def execute(self, context): + return { + 'Image': np.array(self.value.pixels).reshape((*self.value.size, self.value.channels)) + } - data_type: bpy.props.EnumProperty(name="", items=( - ('integer', 'Integer', '', 1), - )) +class NodeSceneInfo(DreamTexturesNode): + bl_idname = "dream_textures.node_scene" + bl_label = "Scene Info" def init(self, context): - self.inputs.new("NodeSocketInt", "Min") - self.inputs.new("NodeSocketInt", "Max") - - self.outputs.new("NodeSocketInt", "Value") + self.outputs.new("NodeSocketImage", "Depth Map") def draw_buttons(self, context, layout): - layout.prop(self, "data_type") + pass + + @classmethod + def render_depth_map(cls, context, collection=None): + width, height = context.scene.render.resolution_x, context.scene.render.resolution_y + matrix = context.scene.camera.matrix_world.inverted() + projection_matrix = context.scene.camera.calc_matrix_camera( + context, + x=width, + y=height + ) + offscreen = gpu.types.GPUOffScreen(width, height) + + with offscreen.bind(): + fb = gpu.state.active_framebuffer_get() + fb.clear(color=(0.0, 0.0, 0.0, 0.0)) + gpu.state.depth_test_set('LESS_EQUAL') + gpu.state.depth_mask_set(True) + with gpu.matrix.push_pop(): + gpu.matrix.load_matrix(matrix) + gpu.matrix.load_projection_matrix(projection_matrix) + + shader = gpu.shader.from_builtin('3D_UNIFORM_COLOR') + + for object in (context.scene.objects if collection is None else collection.objects): + try: + mesh = object.to_mesh(depsgraph=context) + except: + continue + if mesh is None: + continue + vertices = np.empty((len(mesh.vertices), 3), 'f') + indices = np.empty((len(mesh.loop_triangles), 3), 'i') + + mesh.vertices.foreach_get("co", np.reshape(vertices, len(mesh.vertices) * 3)) + mesh.loop_triangles.foreach_get("vertices", np.reshape(indices, len(mesh.loop_triangles) * 3)) + + batch = batch_for_shader( + shader, 'TRIS', + {"pos": vertices}, + indices=indices, + ) + batch.draw(shader) + depth = np.array(fb.read_depth(0, 0, width, height).to_list()) + depth = np.interp(depth, [np.ma.masked_equal(depth, 0, copy=False).min(), depth.max()], [0, 1]).clip(0, 1) + offscreen.free() + return depth - def execute(self, context, min, max): + def execute(self, context): return { - 'Value': random.randrange(min, max) + 'Depth Map': NodeSceneInfo.render_depth_map(context) } \ No newline at end of file diff --git a/engine/nodes/pipeline_nodes.py b/engine/nodes/pipeline_nodes.py index 9aef1fef..6db7176e 100644 --- a/engine/nodes/pipeline_nodes.py +++ b/engine/nodes/pipeline_nodes.py @@ -3,6 +3,7 @@ from ...generator_process import Generator from ...generator_process.actions.prompt_to_image import StepPreviewMode from ...property_groups.dream_prompt import DreamPrompt, control_net_options +from .input_nodes import NodeSceneInfo import numpy as np from dataclasses import dataclass from typing import Any @@ -34,9 +35,23 @@ class ControlNet: control_type: ControlType conditioning_scale: float + def control(self, context): + if self.image is not None: + return np.flipud(self.image) + else: + match self.control_type: + case ControlType.DEPTH: + return np.flipud(NodeSceneInfo.render_depth_map(context, collection=self.collection)) + case ControlType.OPENPOSE: + pass + case ControlType.NORMAL: + pass + def _update_stable_diffusion_sockets(self, context): self.inputs['Source Image'].enabled = self.task in {'image_to_image', 'depth_to_image'} self.inputs['Noise Strength'].enabled = self.task in {'image_to_image', 'depth_to_image'} + if self.task == 'depth_to_image': + self.inputs['Noise Strength'].default_value = 1.0 self.inputs['Depth Map'].enabled = self.task == 'depth_to_image' self.inputs['ControlNets'].enabled = self.task != 'depth_to_image' class NodeStableDiffusion(DreamTexturesNode): @@ -51,8 +66,8 @@ class NodeStableDiffusion(DreamTexturesNode): ), update=_update_stable_diffusion_sockets) def init(self, context): - self.inputs.new("NodeSocketColor", "Depth Map") - self.inputs.new("NodeSocketColor", "Source Image") + self.inputs.new("NodeSocketImage", "Depth Map") + self.inputs.new("NodeSocketImage", "Source Image") self.inputs.new("NodeSocketFloat", "Noise Strength").default_value = 0.75 self.inputs.new("NodeSocketString", "Prompt") @@ -68,7 +83,6 @@ def init(self, context): self.inputs.new("NodeSocketControlNet", "ControlNets") self.outputs.new("NodeSocketColor", "Image") - self.outputs.new("NodeSocketInt", "Seed") _update_stable_diffusion_sockets(self, context) @@ -87,51 +101,104 @@ def execute(self, context, prompt, negative_prompt, width, height, steps, seed, self.prompt.cfg_scale = cfg_scale args = self.prompt.generate_args() - match self.task: - case 'prompt_to_image': - result = Generator.shared().prompt_to_image( - pipeline=args['pipeline'], - model=args['model'], - scheduler=args['scheduler'], - optimizations=args['optimizations'], - seamless_axes=args['seamless_axes'], - iterations=args['iterations'], - prompt=prompt, - steps=steps, - seed=seed, - width=width, - height=height, - cfg_scale=cfg_scale, - use_negative_prompt=True, - negative_prompt=negative_prompt, - step_preview_mode=StepPreviewMode.NONE - ).result() - case 'image_to_image': - result = Generator.shared().image_to_image( - pipeline=args['pipeline'], - model=args['model'], - scheduler=args['scheduler'], - optimizations=args['optimizations'], - seamless_axes=args['seamless_axes'], - iterations=args['iterations'], - - image=np.uint8(source_image * 255), - strength=noise_strength, - fit=True, - - prompt=prompt, - steps=steps, - seed=seed, - width=width, - height=height, - cfg_scale=cfg_scale, - use_negative_prompt=True, - negative_prompt=negative_prompt, - step_preview_mode=StepPreviewMode.NONE - ).result() + shared_args = context.scene.dream_textures_engine_prompt.generate_args() + + if controlnets is not None: + if not isinstance(controlnets, ControlNet): + controlnets = controlnets[0] + result = Generator.shared().control_net( + pipeline=args['pipeline'], + model=args['model'], + scheduler=args['scheduler'], + optimizations=shared_args['optimizations'], + seamless_axes=args['seamless_axes'], + iterations=args['iterations'], + + control_net=controlnets.model, + control=controlnets.control(context), + controlnet_conditioning_scale=controlnets.conditioning_scale, + + image=np.uint8(source_image * 255) if self.task == 'image_to_image' else None, + strength=noise_strength, + + prompt=prompt, + steps=steps, + seed=seed, + width=width, + height=height, + cfg_scale=cfg_scale, + use_negative_prompt=True, + negative_prompt=negative_prompt, + step_preview_mode=StepPreviewMode.NONE + ).result() + else: + match self.task: + case 'prompt_to_image': + result = Generator.shared().prompt_to_image( + pipeline=args['pipeline'], + model=args['model'], + scheduler=args['scheduler'], + optimizations=shared_args['optimizations'], + seamless_axes=args['seamless_axes'], + iterations=args['iterations'], + prompt=prompt, + steps=steps, + seed=seed, + width=width, + height=height, + cfg_scale=cfg_scale, + use_negative_prompt=True, + negative_prompt=negative_prompt, + step_preview_mode=StepPreviewMode.NONE + ).result() + case 'image_to_image': + result = Generator.shared().image_to_image( + pipeline=args['pipeline'], + model=args['model'], + scheduler=args['scheduler'], + optimizations=shared_args['optimizations'], + seamless_axes=args['seamless_axes'], + iterations=args['iterations'], + + image=np.uint8(source_image * 255), + strength=noise_strength, + fit=True, + + prompt=prompt, + steps=steps, + seed=seed, + width=width, + height=height, + cfg_scale=cfg_scale, + use_negative_prompt=True, + negative_prompt=negative_prompt, + step_preview_mode=StepPreviewMode.NONE + ).result() + case 'depth_to_image': + result = Generator.shared().depth_to_image( + pipeline=args['pipeline'], + model=args['model'], + scheduler=args['scheduler'], + optimizations=shared_args['optimizations'], + seamless_axes=args['seamless_axes'], + iterations=args['iterations'], + + depth=depth_map, + image=np.uint8(source_image * 255) if source_image is not None else None, + strength=noise_strength, + + prompt=prompt, + steps=steps, + seed=seed, + width=width, + height=height, + cfg_scale=cfg_scale, + use_negative_prompt=True, + negative_prompt=negative_prompt, + step_preview_mode=StepPreviewMode.NONE + ).result() return { - 'Image': result[-1].images[-1], - 'Seed': result[-1].seeds[-1] + 'Image': result[-1].images[-1] } def _update_control_net_sockets(self, context): diff --git a/engine/nodes/utility_nodes.py b/engine/nodes/utility_nodes.py new file mode 100644 index 00000000..7e6f711a --- /dev/null +++ b/engine/nodes/utility_nodes.py @@ -0,0 +1,65 @@ +import bpy +import numpy as np +import random +from ..node import DreamTexturesNode + +class NodeMath(DreamTexturesNode): + bl_idname = "dream_textures.node_math" + bl_label = "Math" + + operation: bpy.props.EnumProperty( + name="Operation", + items=( + ("add", "Add", ""), + ("subtract", "Subtract", ""), + ("multiply", "Multiply", ""), + ("divide", "Divide", ""), + ) + ) + + def init(self, context): + self.inputs.new("NodeSocketFloat", "A") + self.inputs.new("NodeSocketFloat", "B") + + self.outputs.new("NodeSocketFloat", "Value") + + def draw_buttons(self, context, layout): + layout.prop(self, "operation", text="") + + def perform(self, a, b): + match self.operation: + case 'add': + return a + b + case 'subtract': + return a - b + case 'multiply': + return a * b + case 'divide': + return a / b + + def execute(self, context, a, b): + return { + 'Value': self.perform(a, b) + } + +class NodeRandomValue(DreamTexturesNode): + bl_idname = "dream_textures.node_random_value" + bl_label = "Random Value" + + data_type: bpy.props.EnumProperty(name="", items=( + ('integer', 'Integer', '', 1), + )) + + def init(self, context): + self.inputs.new("NodeSocketInt", "Min") + self.inputs.new("NodeSocketInt", "Max").default_value = np.iinfo(np.int32).max + + self.outputs.new("NodeSocketInt", "Value") + + def draw_buttons(self, context, layout): + layout.prop(self, "data_type") + + def execute(self, context, min, max): + return { + 'Value': random.randrange(min, max) + } \ No newline at end of file diff --git a/generator_process/actions/control_net.py b/generator_process/actions/control_net.py index 09756773..1031db54 100644 --- a/generator_process/actions/control_net.py +++ b/generator_process/actions/control_net.py @@ -256,6 +256,7 @@ def __call__( int(8 * (width // 8)), int(8 * (height // 8)), ) + print(control) control_image = PIL.Image.fromarray(np.uint8(control * 255)).convert('RGB').resize(rounded_size) if control is not None else None init_image = None if image is None else (PIL.Image.open(image) if isinstance(image, str) else PIL.Image.fromarray(image.astype(np.uint8))).convert('RGB').resize(rounded_size) diff --git a/generator_process/actions/huggingface_hub.py b/generator_process/actions/huggingface_hub.py index fa1683b5..804a334b 100644 --- a/generator_process/actions/huggingface_hub.py +++ b/generator_process/actions/huggingface_hub.py @@ -84,59 +84,68 @@ def hf_list_models( def hf_list_installed_models(self) -> list[Model]: from diffusers.utils import DIFFUSERS_CACHE - if not os.path.exists(DIFFUSERS_CACHE): - return [] - - def detect_model_type(snapshot_folder): - unet_config = os.path.join(snapshot_folder, 'unet', 'config.json') - config = os.path.join(snapshot_folder, 'config.json') - if os.path.exists(unet_config): - with open(unet_config, 'r') as f: - return ModelType(json.load(f)['in_channels']) - elif os.path.exists(config): - with open(config, 'r') as f: - config_dict = json.load(f) - if '_class_name' in config_dict and config_dict['_class_name'] == 'ControlNetModel': - return ModelType.CONTROL_NET - else: - return ModelType.UNKNOWN - else: - return ModelType.UNKNOWN - - def _map_model(file): - storage_folder = os.path.join(DIFFUSERS_CACHE, file) - model_type = ModelType.UNKNOWN - - if os.path.exists(os.path.join(storage_folder, 'model_index.json')): - snapshot_folder = storage_folder - model_type = detect_model_type(snapshot_folder) - else: - refs_folder = os.path.join(storage_folder, "refs") - if not os.path.exists(refs_folder): - return None - for revision in os.listdir(refs_folder): - ref_path = os.path.join(storage_folder, "refs", revision) - with open(ref_path) as f: - commit_hash = f.read() - snapshot_folder = os.path.join(storage_folder, "snapshots", commit_hash) - if (detected_type := detect_model_type(snapshot_folder)) != ModelType.UNKNOWN: - model_type = detected_type - break - - return Model( - storage_folder, - "", - [], - -1, - -1, - model_type - ) - return [ - model for model in ( - _map_model(file) for file in os.listdir(DIFFUSERS_CACHE) if os.path.isdir(os.path.join(DIFFUSERS_CACHE, file)) - ) - if model is not None - ] + from diffusers.utils.hub_utils import old_diffusers_cache + + def list_dir(cache_dir): + if not os.path.exists(cache_dir): + return [] + + def detect_model_type(snapshot_folder): + unet_config = os.path.join(snapshot_folder, 'unet', 'config.json') + config = os.path.join(snapshot_folder, 'config.json') + if os.path.exists(unet_config): + with open(unet_config, 'r') as f: + return ModelType(json.load(f)['in_channels']) + elif os.path.exists(config): + with open(config, 'r') as f: + config_dict = json.load(f) + if '_class_name' in config_dict and config_dict['_class_name'] == 'ControlNetModel': + return ModelType.CONTROL_NET + else: + return ModelType.UNKNOWN + else: + return ModelType.UNKNOWN + + def _map_model(file): + storage_folder = os.path.join(cache_dir, file) + model_type = ModelType.UNKNOWN + + if os.path.exists(os.path.join(storage_folder, 'model_index.json')): + snapshot_folder = storage_folder + model_type = detect_model_type(snapshot_folder) + else: + refs_folder = os.path.join(storage_folder, "refs") + if not os.path.exists(refs_folder): + return None + for revision in os.listdir(refs_folder): + ref_path = os.path.join(storage_folder, "refs", revision) + with open(ref_path) as f: + commit_hash = f.read() + snapshot_folder = os.path.join(storage_folder, "snapshots", commit_hash) + if (detected_type := detect_model_type(snapshot_folder)) != ModelType.UNKNOWN: + model_type = detected_type + break + + return Model( + storage_folder, + "", + [], + -1, + -1, + model_type + ) + return [ + model for model in ( + _map_model(file) for file in os.listdir(cache_dir) if os.path.isdir(os.path.join(cache_dir, file)) + ) + if model is not None + ] + new_cache_list = list_dir(DIFFUSERS_CACHE) + model_ids = [os.path.basename(m.id) for m in new_cache_list] + for model in list_dir(old_diffusers_cache): + if os.path.basename(model.id) not in model_ids: + new_cache_list.append(model) + return new_cache_list @dataclass class DownloadStatus: diff --git a/generator_process/actions/prompt_to_image.py b/generator_process/actions/prompt_to_image.py index 3b443a56..3927c419 100644 --- a/generator_process/actions/prompt_to_image.py +++ b/generator_process/actions/prompt_to_image.py @@ -375,6 +375,8 @@ def model_snapshot_folder(model, preferred_revision: str | None = None): """ Try to find the preferred revision, but fallback to another revision if necessary. """ import diffusers storage_folder = os.path.join(diffusers.utils.DIFFUSERS_CACHE, model) + if not os.path.exists(os.path.join(storage_folder, "refs")): + storage_folder = os.path.join(diffusers.utils.hub_utils.old_diffusers_cache, model) if os.path.exists(os.path.join(storage_folder, 'model_index.json')): # converted model snapshot_folder = storage_folder else: # hub model diff --git a/ui/panels/dream_texture.py b/ui/panels/dream_texture.py index 8dbf63a9..6e543ae6 100644 --- a/ui/panels/dream_texture.py +++ b/ui/panels/dream_texture.py @@ -265,11 +265,14 @@ def draw(self, context): yield AdvancedPanel + yield from optimization_panels(sub_panel, space_type, get_prompt, AdvancedPanel.bl_idname) + +def optimization_panels(sub_panel, space_type, get_prompt, parent_id=""): class SpeedOptimizationPanel(sub_panel): """Create a subpanel for speed optimizations""" bl_idname = f"DREAM_PT_dream_panel_speed_optimizations_{space_type}" bl_label = "Speed Optimizations" - bl_parent_id = AdvancedPanel.bl_idname + bl_parent_id = parent_id def draw(self, context): super().draw(context) @@ -295,7 +298,7 @@ class MemoryOptimizationPanel(sub_panel): """Create a subpanel for memory optimizations""" bl_idname = f"DREAM_PT_dream_panel_memory_optimizations_{space_type}" bl_label = "Memory Optimizations" - bl_parent_id = AdvancedPanel.bl_idname + bl_parent_id = parent_id def draw(self, context): super().draw(context) From 0380ce5587c3ec0dc76d3be9a6bc819b78880eb2 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Mon, 13 Mar 2023 23:29:12 -0400 Subject: [PATCH 06/28] OpenPose annotation support --- engine/engine.py | 1 + engine/node_executor.py | 2 +- engine/nodes/input_nodes.py | 209 ++++++++++++++++++++++++++++++++- engine/nodes/pipeline_nodes.py | 4 +- 4 files changed, 212 insertions(+), 4 deletions(-) diff --git a/engine/engine.py b/engine/engine.py index 9af89f20..93450235 100644 --- a/engine/engine.py +++ b/engine/engine.py @@ -139,6 +139,7 @@ class DreamTexturesRenderEngineProperties(bpy.types.PropertyGroup): def engine_panels(): bpy.types.RENDER_PT_output.COMPAT_ENGINES.add(DreamTexturesRenderEngine.bl_idname) + bpy.types.RENDER_PT_color_management.COMPAT_ENGINES.add(DreamTexturesRenderEngine.bl_idname) def get_prompt(context): return context.scene.dream_textures_engine_prompt class RenderPanel(bpy.types.Panel, RenderButtonsPanel): diff --git a/engine/node_executor.py b/engine/node_executor.py index cb45feaa..815ae244 100644 --- a/engine/node_executor.py +++ b/engine/node_executor.py @@ -1,7 +1,7 @@ import bpy import numpy as np # from dream_textures.engine import node_executor -# node_executor.execute(bpy.data.node_groups["NodeTree"], bpy.context) +# node_executor.execute(bpy.data.node_groups["NodeTree"], bpy.context.evaluated_depsgraph_get()) def execute_node(node, context, cache, on_execute=lambda _, __: None): if node in cache: diff --git a/engine/nodes/input_nodes.py b/engine/nodes/input_nodes.py index a54d1404..907523f3 100644 --- a/engine/nodes/input_nodes.py +++ b/engine/nodes/input_nodes.py @@ -1,10 +1,53 @@ import bpy +import bpy_extras import gpu from gpu_extras.batch import batch_for_shader -import random +from gpu_extras.presets import draw_circle_2d +import mathutils +import math import numpy as np +import enum from ..node import DreamTexturesNode +def draw_circle_2d(center, radius, segments, color): + m = (1.0 / (segments - 1)) * (math.pi * 2) + + coords = [ + ( + center[0] + math.cos(m * p) * radius, + center[1] + math.sin(m * p) * radius + ) + for p in range(segments) + ] + + shader = gpu.shader.from_builtin('3D_UNIFORM_COLOR') + batch = batch_for_shader(shader, 'TRI_FAN', {"pos": coords}) + shader.uniform_float("color", color) + batch.draw(shader) + +def draw_ellipse_2d(start, end, thickness, segments, color): + length = math.sqrt((end[0] - start[0]) ** 2 + (end[1] - start[1]) ** 2) + theta = math.atan2(end[1] - start[1], end[0] - start[0]) + center = ( + (start[0] + end[0]) / 2, + (start[1] + end[1]) / 2 + ) + major, minor = length / 2, thickness + m = (1.0 / (segments - 1)) * (math.pi * 2) + + coords = [ + ( + center[0] + major * math.cos(m * p) * math.cos(theta) - minor * math.sin(m * p) * math.sin(theta), + center[1] + major * math.cos(m * p) * math.sin(theta) + minor * math.sin(m * p) * math.cos(theta) + ) + for p in range(segments) + ] + + shader = gpu.shader.from_builtin('3D_UNIFORM_COLOR') + batch = batch_for_shader(shader, 'TRI_FAN', {"pos": coords}) + shader.uniform_float("color", color) + batch.draw(shader) + class NodeString(DreamTexturesNode): bl_idname = "dream_textures.node_string" bl_label = "String" @@ -79,6 +122,7 @@ class NodeSceneInfo(DreamTexturesNode): def init(self, context): self.outputs.new("NodeSocketImage", "Depth Map") + self.outputs.new("NodeSocketImage", "OpenPose Map") def draw_buttons(self, context, layout): pass @@ -129,7 +173,168 @@ def render_depth_map(cls, context, collection=None): offscreen.free() return depth + @classmethod + def render_openpose_map(cls, context, collection=None): + width, height = context.scene.render.resolution_x, context.scene.render.resolution_y + offscreen = gpu.types.GPUOffScreen(width, height) + + with offscreen.bind(): + fb = gpu.state.active_framebuffer_get() + fb.clear(color=(0.0, 0.0, 0.0, 0.0)) + gpu.state.depth_test_set('LESS_EQUAL') + gpu.state.depth_mask_set(True) + + class Side(enum.IntEnum): + HEAD = 0 + TAIL = 1 + + class Bone(enum.IntEnum): + NOSE = 0 + CHEST = 1 + + SHOULDER_L = 2 + SHOULDER_R = 3 + ELBOW_L = 4 + ELBOW_R = 5 + HAND_L = 6 + HAND_R = 7 + + HIP_L = 8 + HIP_R = 9 + KNEE_L = 10 + KNEE_R = 11 + FOOT_L = 12 + FOOT_R = 13 + + EYE_L = 14 + EYE_R = 15 + + EAR_L = 16 + EAR_R = 17 + + def identify(self, pose): + options = self.name_detection_options() + for option in options: + if (result := pose.bones.get(option[0], None)) is not None: + return result, option[1] + return None, None + + def name_detection_options(self): + match self: + case Bone.NOSE: + return [('nose_ik.001', Side.TAIL), ('nose.001', Side.TAIL)] + case Bone.CHEST: + return [('spine_fk.003', Side.TAIL), ('spine.003', Side.TAIL)] + case Bone.SHOULDER_L: + return [('shoulder_ik.L', Side.TAIL), ('shoulder.L', Side.TAIL)] + case Bone.SHOULDER_R: + return [('shoulder_ik.R', Side.TAIL), ('shoulder.R', Side.TAIL)] + case Bone.ELBOW_L: + return [('upper_arm_ik.L', Side.TAIL), ('upper_arm.L', Side.TAIL)] + case Bone.ELBOW_R: + return [('upper_arm_ik.R', Side.TAIL), ('upper_arm.R', Side.TAIL)] + case Bone.HAND_L: + return [('hand_ik.L', Side.TAIL), ('forearm.L', Side.TAIL)] + case Bone.HAND_R: + return [('hand_ik.R', Side.TAIL), ('forearm.R', Side.TAIL)] + case Bone.HIP_L: + return [('thigh_ik.L', Side.HEAD), ('thigh.L', Side.HEAD)] + case Bone.HIP_R: + return [('thigh_ik.R', Side.HEAD), ('thigh.R', Side.HEAD)] + case Bone.KNEE_L: + return [('thigh_ik.L', Side.TAIL), ('thigh.L', Side.TAIL)] + case Bone.KNEE_R: + return [('thigh_ik.R', Side.TAIL), ('thigh.R', Side.TAIL)] + case Bone.FOOT_L: + return [('foot_ik.L', Side.TAIL), ('shin.L', Side.TAIL)] + case Bone.FOOT_R: + return [('foot_ik.R', Side.TAIL), ('shin.R', Side.TAIL)] + case Bone.EYE_L: + return [('master_eye.L', Side.TAIL), ('eye.L', Side.TAIL)] + case Bone.EYE_R: + return [('master_eye.R', Side.TAIL), ('eye.R', Side.TAIL)] + case Bone.EAR_L: + return [('ear.L', Side.TAIL), ('ear.L.001', Side.TAIL)] + case Bone.EAR_R: + return [('ear.R', Side.TAIL), ('ear.R.001', Side.TAIL)] + + def color(self): + match self: + case Bone.NOSE: return (255, 0, 0) + case Bone.CHEST: return (255, 85, 0) + case Bone.SHOULDER_L: return (85, 255, 0) + case Bone.SHOULDER_R: return (255, 170, 0) + case Bone.ELBOW_L: return (0, 255, 0) + case Bone.ELBOW_R: return (255, 255, 0) + case Bone.HAND_L: return (0, 255, 85) + case Bone.HAND_R: return (170, 255, 0) + case Bone.HIP_L: return (0, 85, 255) + case Bone.HIP_R: return (0, 255, 170) + case Bone.KNEE_L: return (0, 0, 255) + case Bone.KNEE_R: return (0, 255, 255) + case Bone.FOOT_L: return (85, 0, 255) + case Bone.FOOT_R: return (0, 170, 255) + case Bone.EYE_L: return (255, 0, 255) + case Bone.EYE_R: return (170, 0, 255) + case Bone.EAR_L: return (255, 0, 85) + case Bone.EAR_R: return (255, 0, 170) + + lines = { + (Bone.NOSE, Bone.CHEST): (0, 0, 255), + (Bone.CHEST, Bone.SHOULDER_L): (255, 85, 0), + (Bone.CHEST, Bone.SHOULDER_R): (255, 0, 0), + (Bone.SHOULDER_L, Bone.ELBOW_L): (170, 255, 0), + (Bone.SHOULDER_R, Bone.ELBOW_R): (255, 170, 0), + (Bone.ELBOW_L, Bone.HAND_L): (85, 255, 0), + (Bone.ELBOW_R, Bone.HAND_R): (255, 255, 0), + (Bone.CHEST, Bone.HIP_L): (0, 255, 255), + (Bone.CHEST, Bone.HIP_R): (0, 255, 0), + (Bone.HIP_L, Bone.KNEE_L): (0, 170, 255), + (Bone.HIP_R, Bone.KNEE_R): (0, 255, 85), + (Bone.KNEE_L, Bone.FOOT_L): (0, 85, 255), + (Bone.KNEE_R, Bone.FOOT_R): (0, 255, 170), + (Bone.NOSE, Bone.EYE_L): (255, 0, 255), + (Bone.NOSE, Bone.EYE_R): (85, 0, 255), + (Bone.EYE_L, Bone.EAR_L): (255, 0, 170), + (Bone.EYE_R, Bone.EAR_R): (170, 0, 255), + } + + with gpu.matrix.push_pop(): + gpu.matrix.load_matrix(mathutils.Matrix.Identity(4)) + gpu.matrix.load_projection_matrix(mathutils.Matrix.Identity(4)) + gpu.state.blend_set('ALPHA') + + shader = gpu.shader.from_builtin('3D_UNIFORM_COLOR') + batch = batch_for_shader(shader, 'TRI_STRIP', {"pos": [(-1, -1, 0), (-1, 1, 0), (1, -1, 0), (1, 1, 0)]}) + shader.uniform_float("color", (0, 0, 0, 1)) + batch.draw(shader) + + for object in (context.scene.objects if collection is None else collection.objects): + if object.hide_render: + continue + if object.pose is None: + continue + for connection, color in lines.items(): + a, a_side = connection[0].identify(object.pose) + b, b_side = connection[1].identify(object.pose) + if a is None or b is None: + continue + a = bpy_extras.object_utils.world_to_camera_view(context.scene, context.scene.camera, object.matrix_world @ (a.tail if a_side == Side.TAIL else a.head)) + b = bpy_extras.object_utils.world_to_camera_view(context.scene, context.scene.camera, object.matrix_world @ (b.tail if b_side == Side.TAIL else b.head)) + draw_ellipse_2d(((a[0] - 0.5) * 2, (a[1] - 0.5) * 2), ((b[0] - 0.5) * 2, (b[1] - 0.5) * 2), .015, 32, (color[0] / 255.0, color[1] / 255.0, color[2] / 255.0, 0.5)) + for b in Bone: + bone, side = b.identify(object.pose) + color = b.color() + if bone is None: continue + tail = bpy_extras.object_utils.world_to_camera_view(context.scene, context.scene.camera, object.matrix_world @ (bone.tail if side == Side.TAIL else bone.head)) + draw_circle_2d(((tail[0] - 0.5) * 2, (tail[1] - 0.5) * 2), .015, 16, (color[0] / 255.0, color[1] / 255.0, color[2] / 255.0, 0.5)) + + depth = np.array(fb.read_color(0, 0, width, height, 4, 0, 'FLOAT').to_list()) + offscreen.free() + return depth + def execute(self, context): return { - 'Depth Map': NodeSceneInfo.render_depth_map(context) + 'Depth Map': NodeSceneInfo.render_depth_map(context), + 'OpenPose Map': NodeSceneInfo.render_openpose_map(context) } \ No newline at end of file diff --git a/engine/nodes/pipeline_nodes.py b/engine/nodes/pipeline_nodes.py index 6db7176e..d91e3655 100644 --- a/engine/nodes/pipeline_nodes.py +++ b/engine/nodes/pipeline_nodes.py @@ -8,6 +8,8 @@ from dataclasses import dataclass from typing import Any import enum +import gpu +from gpu_extras.batch import batch_for_shader class NodeSocketControlNet(bpy.types.NodeSocket): bl_idname = "NodeSocketControlNet" @@ -43,7 +45,7 @@ def control(self, context): case ControlType.DEPTH: return np.flipud(NodeSceneInfo.render_depth_map(context, collection=self.collection)) case ControlType.OPENPOSE: - pass + return np.flipud(NodeSceneInfo.render_openpose_map(context, collection=self.collection)) case ControlType.NORMAL: pass From 11196b92e69b31c7616498d01a81c6b6c5b6d275 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Tue, 14 Mar 2023 00:40:11 -0400 Subject: [PATCH 07/28] Improve bone detection --- engine/engine.py | 1 + engine/nodes/input_nodes.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/engine/engine.py b/engine/engine.py index 93450235..dab4b2b4 100644 --- a/engine/engine.py +++ b/engine/engine.py @@ -140,6 +140,7 @@ class DreamTexturesRenderEngineProperties(bpy.types.PropertyGroup): def engine_panels(): bpy.types.RENDER_PT_output.COMPAT_ENGINES.add(DreamTexturesRenderEngine.bl_idname) bpy.types.RENDER_PT_color_management.COMPAT_ENGINES.add(DreamTexturesRenderEngine.bl_idname) + bpy.types.DATA_PT_lens.COMPAT_ENGINES.add(DreamTexturesRenderEngine.bl_idname) def get_prompt(context): return context.scene.dream_textures_engine_prompt class RenderPanel(bpy.types.Panel, RenderButtonsPanel): diff --git a/engine/nodes/input_nodes.py b/engine/nodes/input_nodes.py index 907523f3..f2609379 100644 --- a/engine/nodes/input_nodes.py +++ b/engine/nodes/input_nodes.py @@ -222,7 +222,7 @@ def identify(self, pose): def name_detection_options(self): match self: case Bone.NOSE: - return [('nose_ik.001', Side.TAIL), ('nose.001', Side.TAIL)] + return [('nose_master', Side.TAIL), ('nose.001', Side.TAIL)] case Bone.CHEST: return [('spine_fk.003', Side.TAIL), ('spine.003', Side.TAIL)] case Bone.SHOULDER_L: From 867671dc0ed5cf9a4d6f6ba370473011402d6a43 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Tue, 14 Mar 2023 23:10:47 -0400 Subject: [PATCH 08/28] Add manual bone configuration --- engine/__init__.py | 50 +++++--- engine/annotations/openpose.py | 210 +++++++++++++++++++++++++++++++++ engine/engine.py | 32 ++++- engine/nodes/input_nodes.py | 206 +------------------------------- engine/nodes/pipeline_nodes.py | 13 +- 5 files changed, 282 insertions(+), 229 deletions(-) create mode 100644 engine/annotations/openpose.py diff --git a/engine/__init__.py b/engine/__init__.py index 9fdb32ff..214db80a 100644 --- a/engine/__init__.py +++ b/engine/__init__.py @@ -5,6 +5,7 @@ from .nodes.input_nodes import * from .nodes.pipeline_nodes import * from .nodes.utility_nodes import * +from .annotations import openpose import bpy import nodeitems_utils @@ -15,27 +16,42 @@ def poll(cls, context): return context.space_data.tree_type == DreamTexturesNodeTree.__name__ categories = [ - DreamTexturesNodeCategory("DREAM_TEXTURES_PIPELINE", "Pipeline", items = [ - nodeitems_utils.NodeItem(NodeStableDiffusion.bl_idname), - ]), - DreamTexturesNodeCategory("DREAM_TEXTURES_INPUT", "Input", items = [ - nodeitems_utils.NodeItem(NodeInteger.bl_idname), - nodeitems_utils.NodeItem(NodeString.bl_idname), - nodeitems_utils.NodeItem(NodeImage.bl_idname), - nodeitems_utils.NodeItem(NodeCollection.bl_idname), - nodeitems_utils.NodeItem(NodeSceneInfo.bl_idname), - ]), - DreamTexturesNodeCategory("DREAM_TEXTURES_UTILITY", "Utilities", items = [ - nodeitems_utils.NodeItem(NodeMath.bl_idname), - nodeitems_utils.NodeItem(NodeRandomValue.bl_idname), - ]), - DreamTexturesNodeCategory("DREAM_TEXTURES_GROUP", "Group", items = [ - nodeitems_utils.NodeItem(bpy.types.NodeGroupOutput.__name__), - ]), + DreamTexturesNodeCategory("DREAM_TEXTURES_PIPELINE", "Pipeline", items = [ + nodeitems_utils.NodeItem(NodeStableDiffusion.bl_idname), + ]), + DreamTexturesNodeCategory("DREAM_TEXTURES_INPUT", "Input", items = [ + nodeitems_utils.NodeItem(NodeInteger.bl_idname), + nodeitems_utils.NodeItem(NodeString.bl_idname), + nodeitems_utils.NodeItem(NodeImage.bl_idname), + nodeitems_utils.NodeItem(NodeCollection.bl_idname), + nodeitems_utils.NodeItem(NodeSceneInfo.bl_idname), + ]), + DreamTexturesNodeCategory("DREAM_TEXTURES_UTILITY", "Utilities", items = [ + nodeitems_utils.NodeItem(NodeMath.bl_idname), + nodeitems_utils.NodeItem(NodeRandomValue.bl_idname), + ]), + DreamTexturesNodeCategory("DREAM_TEXTURES_GROUP", "Group", items = [ + nodeitems_utils.NodeItem(bpy.types.NodeGroupOutput.__name__), + ]), ] def register(): + # Prompt bpy.types.Scene.dream_textures_engine_prompt = bpy.props.PointerProperty(type=DreamPrompt) + + # Bone + bpy.types.Bone.dream_textures_openpose = bpy.props.BoolProperty( + name="Use OpenPose", + default=False + ) + bpy.types.Bone.dream_textures_openpose_bone = bpy.props.EnumProperty( + name="OpenPose Bone", + items=((str(b.value), b.name.title(), '') for b in openpose.Bone) + ) + bpy.types.Bone.dream_textures_openpose_bone_side = bpy.props.EnumProperty( + name="Endpoint Side", + items=((str(s.value), s.name.title(), '') for s in openpose.Side) + ) bpy.utils.register_class(DreamTexturesNodeTree) diff --git a/engine/annotations/openpose.py b/engine/annotations/openpose.py new file mode 100644 index 00000000..38682809 --- /dev/null +++ b/engine/annotations/openpose.py @@ -0,0 +1,210 @@ +import bpy_extras +import gpu +from gpu_extras.batch import batch_for_shader +from gpu_extras.presets import draw_circle_2d +import mathutils +import numpy as np +import enum +import math + +class Side(enum.IntEnum): + HEAD = 0 + TAIL = 1 + +class Bone(enum.IntEnum): + NOSE = 0 + CHEST = 1 + + SHOULDER_L = 2 + SHOULDER_R = 3 + ELBOW_L = 4 + ELBOW_R = 5 + HAND_L = 6 + HAND_R = 7 + + HIP_L = 8 + HIP_R = 9 + KNEE_L = 10 + KNEE_R = 11 + FOOT_L = 12 + FOOT_R = 13 + + EYE_L = 14 + EYE_R = 15 + + EAR_L = 16 + EAR_R = 17 + + def identify(self, pose): + for bone in pose.bones: + if bone.bone.dream_textures_openpose: + if bone.bone.dream_textures_openpose_bone == str(self.value): + return bone, Side(int(bone.bone.dream_textures_openpose_bone_side)) + options = self.name_detection_options() + for option in options: + if (result := pose.bones.get(option[0], None)) is not None: + return result, option[1] + return None, None + + def name_detection_options(self): + match self: + case Bone.NOSE: + return [('nose_master', Side.TAIL), ('nose.001', Side.TAIL)] + case Bone.CHEST: + return [('spine_fk.003', Side.TAIL), ('spine.003', Side.TAIL)] + case Bone.SHOULDER_L: + return [('shoulder_ik.L', Side.TAIL), ('shoulder.L', Side.TAIL)] + case Bone.SHOULDER_R: + return [('shoulder_ik.R', Side.TAIL), ('shoulder.R', Side.TAIL)] + case Bone.ELBOW_L: + return [('upper_arm_ik.L', Side.TAIL), ('upper_arm.L', Side.TAIL)] + case Bone.ELBOW_R: + return [('upper_arm_ik.R', Side.TAIL), ('upper_arm.R', Side.TAIL)] + case Bone.HAND_L: + return [('hand_ik.L', Side.TAIL), ('forearm.L', Side.TAIL)] + case Bone.HAND_R: + return [('hand_ik.R', Side.TAIL), ('forearm.R', Side.TAIL)] + case Bone.HIP_L: + return [('thigh_ik.L', Side.HEAD), ('thigh.L', Side.HEAD)] + case Bone.HIP_R: + return [('thigh_ik.R', Side.HEAD), ('thigh.R', Side.HEAD)] + case Bone.KNEE_L: + return [('thigh_ik.L', Side.TAIL), ('thigh.L', Side.TAIL)] + case Bone.KNEE_R: + return [('thigh_ik.R', Side.TAIL), ('thigh.R', Side.TAIL)] + case Bone.FOOT_L: + return [('foot_ik.L', Side.TAIL), ('shin.L', Side.TAIL)] + case Bone.FOOT_R: + return [('foot_ik.R', Side.TAIL), ('shin.R', Side.TAIL)] + case Bone.EYE_L: + return [('master_eye.L', Side.TAIL), ('eye.L', Side.TAIL)] + case Bone.EYE_R: + return [('master_eye.R', Side.TAIL), ('eye.R', Side.TAIL)] + case Bone.EAR_L: + return [('ear.L', Side.TAIL), ('ear.L.001', Side.TAIL)] + case Bone.EAR_R: + return [('ear.R', Side.TAIL), ('ear.R.001', Side.TAIL)] + + def color(self): + match self: + case Bone.NOSE: return (255, 0, 0) + case Bone.CHEST: return (255, 85, 0) + case Bone.SHOULDER_L: return (85, 255, 0) + case Bone.SHOULDER_R: return (255, 170, 0) + case Bone.ELBOW_L: return (0, 255, 0) + case Bone.ELBOW_R: return (255, 255, 0) + case Bone.HAND_L: return (0, 255, 85) + case Bone.HAND_R: return (170, 255, 0) + case Bone.HIP_L: return (0, 85, 255) + case Bone.HIP_R: return (0, 255, 170) + case Bone.KNEE_L: return (0, 0, 255) + case Bone.KNEE_R: return (0, 255, 255) + case Bone.FOOT_L: return (85, 0, 255) + case Bone.FOOT_R: return (0, 170, 255) + case Bone.EYE_L: return (255, 0, 255) + case Bone.EYE_R: return (170, 0, 255) + case Bone.EAR_L: return (255, 0, 85) + case Bone.EAR_R: return (255, 0, 170) + +def render_openpose_map(context, collection=None): + width, height = context.scene.render.resolution_x, context.scene.render.resolution_y + offscreen = gpu.types.GPUOffScreen(width, height) + + with offscreen.bind(): + fb = gpu.state.active_framebuffer_get() + fb.clear(color=(0.0, 0.0, 0.0, 0.0)) + gpu.state.depth_test_set('LESS_EQUAL') + gpu.state.depth_mask_set(True) + + lines = { + (Bone.NOSE, Bone.CHEST): (0, 0, 255), + (Bone.CHEST, Bone.SHOULDER_L): (255, 85, 0), + (Bone.CHEST, Bone.SHOULDER_R): (255, 0, 0), + (Bone.SHOULDER_L, Bone.ELBOW_L): (170, 255, 0), + (Bone.SHOULDER_R, Bone.ELBOW_R): (255, 170, 0), + (Bone.ELBOW_L, Bone.HAND_L): (85, 255, 0), + (Bone.ELBOW_R, Bone.HAND_R): (255, 255, 0), + (Bone.CHEST, Bone.HIP_L): (0, 255, 255), + (Bone.CHEST, Bone.HIP_R): (0, 255, 0), + (Bone.HIP_L, Bone.KNEE_L): (0, 170, 255), + (Bone.HIP_R, Bone.KNEE_R): (0, 255, 85), + (Bone.KNEE_L, Bone.FOOT_L): (0, 85, 255), + (Bone.KNEE_R, Bone.FOOT_R): (0, 255, 170), + (Bone.NOSE, Bone.EYE_L): (255, 0, 255), + (Bone.NOSE, Bone.EYE_R): (85, 0, 255), + (Bone.EYE_L, Bone.EAR_L): (255, 0, 170), + (Bone.EYE_R, Bone.EAR_R): (170, 0, 255), + } + + with gpu.matrix.push_pop(): + gpu.matrix.load_matrix(mathutils.Matrix.Identity(4)) + gpu.matrix.load_projection_matrix(mathutils.Matrix.Identity(4)) + gpu.state.blend_set('ALPHA') + + shader = gpu.shader.from_builtin('3D_UNIFORM_COLOR') + batch = batch_for_shader(shader, 'TRI_STRIP', {"pos": [(-1, -1, 0), (-1, 1, 0), (1, -1, 0), (1, 1, 0)]}) + shader.uniform_float("color", (0, 0, 0, 1)) + batch.draw(shader) + + for object in (context.scene.objects if collection is None else collection.objects): + if object.hide_render: + continue + if object.pose is None: + continue + for connection, color in lines.items(): + a, a_side = connection[0].identify(object.pose) + b, b_side = connection[1].identify(object.pose) + if a is None or b is None: + continue + a = bpy_extras.object_utils.world_to_camera_view(context.scene, context.scene.camera, object.matrix_world @ (a.tail if a_side == Side.TAIL else a.head)) + b = bpy_extras.object_utils.world_to_camera_view(context.scene, context.scene.camera, object.matrix_world @ (b.tail if b_side == Side.TAIL else b.head)) + draw_ellipse_2d(((a[0] - 0.5) * 2, (a[1] - 0.5) * 2), ((b[0] - 0.5) * 2, (b[1] - 0.5) * 2), .015, 32, (color[0] / 255.0, color[1] / 255.0, color[2] / 255.0, 0.5)) + for b in Bone: + bone, side = b.identify(object.pose) + color = b.color() + if bone is None: continue + tail = bpy_extras.object_utils.world_to_camera_view(context.scene, context.scene.camera, object.matrix_world @ (bone.tail if side == Side.TAIL else bone.head)) + draw_circle_2d(((tail[0] - 0.5) * 2, (tail[1] - 0.5) * 2), .015, 16, (color[0] / 255.0, color[1] / 255.0, color[2] / 255.0, 0.5)) + + depth = np.array(fb.read_color(0, 0, width, height, 4, 0, 'FLOAT').to_list()) + offscreen.free() + return depth + +def draw_circle_2d(center, radius, segments, color): + m = (1.0 / (segments - 1)) * (math.pi * 2) + + coords = [ + ( + center[0] + math.cos(m * p) * radius, + center[1] + math.sin(m * p) * radius + ) + for p in range(segments) + ] + + shader = gpu.shader.from_builtin('3D_UNIFORM_COLOR') + batch = batch_for_shader(shader, 'TRI_FAN', {"pos": coords}) + shader.uniform_float("color", color) + batch.draw(shader) + +def draw_ellipse_2d(start, end, thickness, segments, color): + length = math.sqrt((end[0] - start[0]) ** 2 + (end[1] - start[1]) ** 2) + theta = math.atan2(end[1] - start[1], end[0] - start[0]) + center = ( + (start[0] + end[0]) / 2, + (start[1] + end[1]) / 2 + ) + major, minor = length / 2, thickness + m = (1.0 / (segments - 1)) * (math.pi * 2) + + coords = [ + ( + center[0] + major * math.cos(m * p) * math.cos(theta) - minor * math.sin(m * p) * math.sin(theta), + center[1] + major * math.cos(m * p) * math.sin(theta) + minor * math.sin(m * p) * math.cos(theta) + ) + for p in range(segments) + ] + + shader = gpu.shader.from_builtin('3D_UNIFORM_COLOR') + batch = batch_for_shader(shader, 'TRI_FAN', {"pos": coords}) + shader.uniform_float("color", color) + batch.draw(shader) \ No newline at end of file diff --git a/engine/engine.py b/engine/engine.py index dab4b2b4..0ee89eff 100644 --- a/engine/engine.py +++ b/engine/engine.py @@ -170,4 +170,34 @@ def draw(self, context): col = layout.column(align=True) col.prop(context.scene.render, "resolution_x") col.prop(context.scene.render, "resolution_y", text="Y") - yield FormatPanel \ No newline at end of file + yield FormatPanel + + # Bone properties + class OpenPoseBonePanel(bpy.types.Panel): + bl_idname = "DREAM_PT_dream_textures_engine_bone_properties" + bl_label = "OpenPose" + bl_space_type = 'PROPERTIES' + bl_region_type = 'WINDOW' + bl_context = "bone" + + @classmethod + def poll(cls, context): + return (context.bone or context.edit_bone) and context.scene.render.engine == 'DREAM_TEXTURES' + + def draw_header(self, context): + bone = context.bone or context.edit_bone + if bone: + self.layout.prop(bone, "dream_textures_openpose", text="") + + def draw(self, context): + layout = self.layout + layout.use_property_split = True + + bone = context.bone or context.edit_bone + + if bone: + layout.enabled = bone.dream_textures_openpose + layout.prop(bone, "dream_textures_openpose_bone") + layout.prop(bone, "dream_textures_openpose_bone_side") + + yield OpenPoseBonePanel \ No newline at end of file diff --git a/engine/nodes/input_nodes.py b/engine/nodes/input_nodes.py index f2609379..acf11133 100644 --- a/engine/nodes/input_nodes.py +++ b/engine/nodes/input_nodes.py @@ -1,52 +1,10 @@ import bpy -import bpy_extras import gpu from gpu_extras.batch import batch_for_shader -from gpu_extras.presets import draw_circle_2d -import mathutils import math import numpy as np -import enum from ..node import DreamTexturesNode - -def draw_circle_2d(center, radius, segments, color): - m = (1.0 / (segments - 1)) * (math.pi * 2) - - coords = [ - ( - center[0] + math.cos(m * p) * radius, - center[1] + math.sin(m * p) * radius - ) - for p in range(segments) - ] - - shader = gpu.shader.from_builtin('3D_UNIFORM_COLOR') - batch = batch_for_shader(shader, 'TRI_FAN', {"pos": coords}) - shader.uniform_float("color", color) - batch.draw(shader) - -def draw_ellipse_2d(start, end, thickness, segments, color): - length = math.sqrt((end[0] - start[0]) ** 2 + (end[1] - start[1]) ** 2) - theta = math.atan2(end[1] - start[1], end[0] - start[0]) - center = ( - (start[0] + end[0]) / 2, - (start[1] + end[1]) / 2 - ) - major, minor = length / 2, thickness - m = (1.0 / (segments - 1)) * (math.pi * 2) - - coords = [ - ( - center[0] + major * math.cos(m * p) * math.cos(theta) - minor * math.sin(m * p) * math.sin(theta), - center[1] + major * math.cos(m * p) * math.sin(theta) + minor * math.sin(m * p) * math.cos(theta) - ) - for p in range(segments) - ] - - shader = gpu.shader.from_builtin('3D_UNIFORM_COLOR') - batch = batch_for_shader(shader, 'TRI_FAN', {"pos": coords}) - shader.uniform_float("color", color) - batch.draw(shader) +from ..annotations import openpose class NodeString(DreamTexturesNode): bl_idname = "dream_textures.node_string" @@ -173,168 +131,8 @@ def render_depth_map(cls, context, collection=None): offscreen.free() return depth - @classmethod - def render_openpose_map(cls, context, collection=None): - width, height = context.scene.render.resolution_x, context.scene.render.resolution_y - offscreen = gpu.types.GPUOffScreen(width, height) - - with offscreen.bind(): - fb = gpu.state.active_framebuffer_get() - fb.clear(color=(0.0, 0.0, 0.0, 0.0)) - gpu.state.depth_test_set('LESS_EQUAL') - gpu.state.depth_mask_set(True) - - class Side(enum.IntEnum): - HEAD = 0 - TAIL = 1 - - class Bone(enum.IntEnum): - NOSE = 0 - CHEST = 1 - - SHOULDER_L = 2 - SHOULDER_R = 3 - ELBOW_L = 4 - ELBOW_R = 5 - HAND_L = 6 - HAND_R = 7 - - HIP_L = 8 - HIP_R = 9 - KNEE_L = 10 - KNEE_R = 11 - FOOT_L = 12 - FOOT_R = 13 - - EYE_L = 14 - EYE_R = 15 - - EAR_L = 16 - EAR_R = 17 - - def identify(self, pose): - options = self.name_detection_options() - for option in options: - if (result := pose.bones.get(option[0], None)) is not None: - return result, option[1] - return None, None - - def name_detection_options(self): - match self: - case Bone.NOSE: - return [('nose_master', Side.TAIL), ('nose.001', Side.TAIL)] - case Bone.CHEST: - return [('spine_fk.003', Side.TAIL), ('spine.003', Side.TAIL)] - case Bone.SHOULDER_L: - return [('shoulder_ik.L', Side.TAIL), ('shoulder.L', Side.TAIL)] - case Bone.SHOULDER_R: - return [('shoulder_ik.R', Side.TAIL), ('shoulder.R', Side.TAIL)] - case Bone.ELBOW_L: - return [('upper_arm_ik.L', Side.TAIL), ('upper_arm.L', Side.TAIL)] - case Bone.ELBOW_R: - return [('upper_arm_ik.R', Side.TAIL), ('upper_arm.R', Side.TAIL)] - case Bone.HAND_L: - return [('hand_ik.L', Side.TAIL), ('forearm.L', Side.TAIL)] - case Bone.HAND_R: - return [('hand_ik.R', Side.TAIL), ('forearm.R', Side.TAIL)] - case Bone.HIP_L: - return [('thigh_ik.L', Side.HEAD), ('thigh.L', Side.HEAD)] - case Bone.HIP_R: - return [('thigh_ik.R', Side.HEAD), ('thigh.R', Side.HEAD)] - case Bone.KNEE_L: - return [('thigh_ik.L', Side.TAIL), ('thigh.L', Side.TAIL)] - case Bone.KNEE_R: - return [('thigh_ik.R', Side.TAIL), ('thigh.R', Side.TAIL)] - case Bone.FOOT_L: - return [('foot_ik.L', Side.TAIL), ('shin.L', Side.TAIL)] - case Bone.FOOT_R: - return [('foot_ik.R', Side.TAIL), ('shin.R', Side.TAIL)] - case Bone.EYE_L: - return [('master_eye.L', Side.TAIL), ('eye.L', Side.TAIL)] - case Bone.EYE_R: - return [('master_eye.R', Side.TAIL), ('eye.R', Side.TAIL)] - case Bone.EAR_L: - return [('ear.L', Side.TAIL), ('ear.L.001', Side.TAIL)] - case Bone.EAR_R: - return [('ear.R', Side.TAIL), ('ear.R.001', Side.TAIL)] - - def color(self): - match self: - case Bone.NOSE: return (255, 0, 0) - case Bone.CHEST: return (255, 85, 0) - case Bone.SHOULDER_L: return (85, 255, 0) - case Bone.SHOULDER_R: return (255, 170, 0) - case Bone.ELBOW_L: return (0, 255, 0) - case Bone.ELBOW_R: return (255, 255, 0) - case Bone.HAND_L: return (0, 255, 85) - case Bone.HAND_R: return (170, 255, 0) - case Bone.HIP_L: return (0, 85, 255) - case Bone.HIP_R: return (0, 255, 170) - case Bone.KNEE_L: return (0, 0, 255) - case Bone.KNEE_R: return (0, 255, 255) - case Bone.FOOT_L: return (85, 0, 255) - case Bone.FOOT_R: return (0, 170, 255) - case Bone.EYE_L: return (255, 0, 255) - case Bone.EYE_R: return (170, 0, 255) - case Bone.EAR_L: return (255, 0, 85) - case Bone.EAR_R: return (255, 0, 170) - - lines = { - (Bone.NOSE, Bone.CHEST): (0, 0, 255), - (Bone.CHEST, Bone.SHOULDER_L): (255, 85, 0), - (Bone.CHEST, Bone.SHOULDER_R): (255, 0, 0), - (Bone.SHOULDER_L, Bone.ELBOW_L): (170, 255, 0), - (Bone.SHOULDER_R, Bone.ELBOW_R): (255, 170, 0), - (Bone.ELBOW_L, Bone.HAND_L): (85, 255, 0), - (Bone.ELBOW_R, Bone.HAND_R): (255, 255, 0), - (Bone.CHEST, Bone.HIP_L): (0, 255, 255), - (Bone.CHEST, Bone.HIP_R): (0, 255, 0), - (Bone.HIP_L, Bone.KNEE_L): (0, 170, 255), - (Bone.HIP_R, Bone.KNEE_R): (0, 255, 85), - (Bone.KNEE_L, Bone.FOOT_L): (0, 85, 255), - (Bone.KNEE_R, Bone.FOOT_R): (0, 255, 170), - (Bone.NOSE, Bone.EYE_L): (255, 0, 255), - (Bone.NOSE, Bone.EYE_R): (85, 0, 255), - (Bone.EYE_L, Bone.EAR_L): (255, 0, 170), - (Bone.EYE_R, Bone.EAR_R): (170, 0, 255), - } - - with gpu.matrix.push_pop(): - gpu.matrix.load_matrix(mathutils.Matrix.Identity(4)) - gpu.matrix.load_projection_matrix(mathutils.Matrix.Identity(4)) - gpu.state.blend_set('ALPHA') - - shader = gpu.shader.from_builtin('3D_UNIFORM_COLOR') - batch = batch_for_shader(shader, 'TRI_STRIP', {"pos": [(-1, -1, 0), (-1, 1, 0), (1, -1, 0), (1, 1, 0)]}) - shader.uniform_float("color", (0, 0, 0, 1)) - batch.draw(shader) - - for object in (context.scene.objects if collection is None else collection.objects): - if object.hide_render: - continue - if object.pose is None: - continue - for connection, color in lines.items(): - a, a_side = connection[0].identify(object.pose) - b, b_side = connection[1].identify(object.pose) - if a is None or b is None: - continue - a = bpy_extras.object_utils.world_to_camera_view(context.scene, context.scene.camera, object.matrix_world @ (a.tail if a_side == Side.TAIL else a.head)) - b = bpy_extras.object_utils.world_to_camera_view(context.scene, context.scene.camera, object.matrix_world @ (b.tail if b_side == Side.TAIL else b.head)) - draw_ellipse_2d(((a[0] - 0.5) * 2, (a[1] - 0.5) * 2), ((b[0] - 0.5) * 2, (b[1] - 0.5) * 2), .015, 32, (color[0] / 255.0, color[1] / 255.0, color[2] / 255.0, 0.5)) - for b in Bone: - bone, side = b.identify(object.pose) - color = b.color() - if bone is None: continue - tail = bpy_extras.object_utils.world_to_camera_view(context.scene, context.scene.camera, object.matrix_world @ (bone.tail if side == Side.TAIL else bone.head)) - draw_circle_2d(((tail[0] - 0.5) * 2, (tail[1] - 0.5) * 2), .015, 16, (color[0] / 255.0, color[1] / 255.0, color[2] / 255.0, 0.5)) - - depth = np.array(fb.read_color(0, 0, width, height, 4, 0, 'FLOAT').to_list()) - offscreen.free() - return depth - def execute(self, context): return { 'Depth Map': NodeSceneInfo.render_depth_map(context), - 'OpenPose Map': NodeSceneInfo.render_openpose_map(context) + 'OpenPose Map': openpose.render_openpose_map(context) } \ No newline at end of file diff --git a/engine/nodes/pipeline_nodes.py b/engine/nodes/pipeline_nodes.py index d91e3655..f1684528 100644 --- a/engine/nodes/pipeline_nodes.py +++ b/engine/nodes/pipeline_nodes.py @@ -1,15 +1,14 @@ import bpy +import numpy as np +from dataclasses import dataclass +from typing import Any +import enum from ..node import DreamTexturesNode from ...generator_process import Generator from ...generator_process.actions.prompt_to_image import StepPreviewMode from ...property_groups.dream_prompt import DreamPrompt, control_net_options from .input_nodes import NodeSceneInfo -import numpy as np -from dataclasses import dataclass -from typing import Any -import enum -import gpu -from gpu_extras.batch import batch_for_shader +from ..annotations import openpose class NodeSocketControlNet(bpy.types.NodeSocket): bl_idname = "NodeSocketControlNet" @@ -45,7 +44,7 @@ def control(self, context): case ControlType.DEPTH: return np.flipud(NodeSceneInfo.render_depth_map(context, collection=self.collection)) case ControlType.OPENPOSE: - return np.flipud(NodeSceneInfo.render_openpose_map(context, collection=self.collection)) + return np.flipud(openpose.render_openpose_map(context, collection=self.collection)) case ControlType.NORMAL: pass From 1f34b10aa6bb2bc7a573b21e36d530ecc734bdbe Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Tue, 14 Mar 2023 23:44:14 -0400 Subject: [PATCH 09/28] Allow disabling individual bones --- engine/__init__.py | 24 +++++++------ engine/annotations/openpose.py | 38 +++++++++++++++----- engine/engine.py | 64 ++++++++++++++++++++++++++++++---- 3 files changed, 101 insertions(+), 25 deletions(-) diff --git a/engine/__init__.py b/engine/__init__.py index 214db80a..58a073a4 100644 --- a/engine/__init__.py +++ b/engine/__init__.py @@ -39,18 +39,14 @@ def register(): # Prompt bpy.types.Scene.dream_textures_engine_prompt = bpy.props.PointerProperty(type=DreamPrompt) - # Bone - bpy.types.Bone.dream_textures_openpose = bpy.props.BoolProperty( - name="Use OpenPose", - default=False + # OpenPose + bpy.utils.register_class(openpose.ArmatureOpenPoseData) + bpy.types.Armature.dream_textures_openpose = bpy.props.PointerProperty( + type=openpose.ArmatureOpenPoseData ) - bpy.types.Bone.dream_textures_openpose_bone = bpy.props.EnumProperty( - name="OpenPose Bone", - items=((str(b.value), b.name.title(), '') for b in openpose.Bone) - ) - bpy.types.Bone.dream_textures_openpose_bone_side = bpy.props.EnumProperty( - name="Endpoint Side", - items=((str(s.value), s.name.title(), '') for s in openpose.Side) + bpy.utils.register_class(openpose.BoneOpenPoseData) + bpy.types.Bone.dream_textures_openpose = bpy.props.PointerProperty( + type=openpose.BoneOpenPoseData ) bpy.utils.register_class(DreamTexturesNodeTree) @@ -72,6 +68,12 @@ def register(): nodeitems_utils.register_node_categories("DREAM_TEXTURES_CATEGORIES", categories) def unregister(): + # OpenPose + del bpy.types.Armature.dream_textures_openpose + bpy.utils.unregister_class(openpose.ArmatureOpenPoseData) + del bpy.types.Bone.dream_textures_openpose + bpy.utils.unregister_class(openpose.BoneOpenPoseData) + bpy.utils.unregister_class(DreamTexturesNodeTree) # Nodes diff --git a/engine/annotations/openpose.py b/engine/annotations/openpose.py index 38682809..23146d21 100644 --- a/engine/annotations/openpose.py +++ b/engine/annotations/openpose.py @@ -1,7 +1,7 @@ +import bpy import bpy_extras import gpu from gpu_extras.batch import batch_for_shader -from gpu_extras.presets import draw_circle_2d import mathutils import numpy as np import enum @@ -35,11 +35,13 @@ class Bone(enum.IntEnum): EAR_L = 16 EAR_R = 17 - def identify(self, pose): + def identify(self, armature, pose): + if not getattr(armature.dream_textures_openpose, self.name): + return None, None for bone in pose.bones: - if bone.bone.dream_textures_openpose: - if bone.bone.dream_textures_openpose_bone == str(self.value): - return bone, Side(int(bone.bone.dream_textures_openpose_bone_side)) + if bone.bone.dream_textures_openpose.enabled: + if bone.bone.dream_textures_openpose.bone == str(self.value): + return bone, Side(int(bone.bone.dream_textures_openpose.side)) options = self.name_detection_options() for option in options: if (result := pose.bones.get(option[0], None)) is not None: @@ -106,6 +108,26 @@ def color(self): case Bone.EAR_L: return (255, 0, 85) case Bone.EAR_R: return (255, 0, 170) +class BoneOpenPoseData(bpy.types.PropertyGroup): + bl_label = "OpenPose" + bl_idname = "dream_textures.BoneOpenPoseData" + + enabled: bpy.props.BoolProperty(name="Enabled", default=False) + bone: bpy.props.EnumProperty( + name="OpenPose Bone", + items=((str(b.value), b.name.title(), '') for b in Bone) + ) + side: bpy.props.EnumProperty( + name="Bone Side", + items=((str(s.value), s.name.title(), '') for s in Side) + ) + +ArmatureOpenPoseData = type('ArmatureOpenPoseData', (bpy.types.PropertyGroup,), { + "bl_label": "OpenPose", + "bl_idname": "dream_textures.ArmatureOpenPoseData", + "__annotations__": { b.name: bpy.props.BoolProperty(name=b.name.title(), default=True) for b in Bone }, +}) + def render_openpose_map(context, collection=None): width, height = context.scene.render.resolution_x, context.scene.render.resolution_y offscreen = gpu.types.GPUOffScreen(width, height) @@ -152,15 +174,15 @@ def render_openpose_map(context, collection=None): if object.pose is None: continue for connection, color in lines.items(): - a, a_side = connection[0].identify(object.pose) - b, b_side = connection[1].identify(object.pose) + a, a_side = connection[0].identify(object.data, object.pose) + b, b_side = connection[1].identify(object.data, object.pose) if a is None or b is None: continue a = bpy_extras.object_utils.world_to_camera_view(context.scene, context.scene.camera, object.matrix_world @ (a.tail if a_side == Side.TAIL else a.head)) b = bpy_extras.object_utils.world_to_camera_view(context.scene, context.scene.camera, object.matrix_world @ (b.tail if b_side == Side.TAIL else b.head)) draw_ellipse_2d(((a[0] - 0.5) * 2, (a[1] - 0.5) * 2), ((b[0] - 0.5) * 2, (b[1] - 0.5) * 2), .015, 32, (color[0] / 255.0, color[1] / 255.0, color[2] / 255.0, 0.5)) for b in Bone: - bone, side = b.identify(object.pose) + bone, side = b.identify(object.data, object.pose) color = b.color() if bone is None: continue tail = bpy_extras.object_utils.world_to_camera_view(context.scene, context.scene.camera, object.matrix_world @ (bone.tail if side == Side.TAIL else bone.head)) diff --git a/engine/engine.py b/engine/engine.py index 0ee89eff..7bff258a 100644 --- a/engine/engine.py +++ b/engine/engine.py @@ -6,6 +6,7 @@ from ..ui.panels.dream_texture import optimization_panels from .node_tree import DreamTexturesNodeTree from ..engine import node_executor +from .annotations import openpose class DreamTexturesRenderEngine(bpy.types.RenderEngine): """A custom Dream Textures render engine, that uses Stable Diffusion and scene data to render images, instead of as a pass on top of Cycles.""" @@ -173,8 +174,60 @@ def draw(self, context): yield FormatPanel # Bone properties + class OpenPoseArmaturePanel(bpy.types.Panel): + bl_idname = "DREAM_PT_dream_textures_armature_openpose" + bl_label = "OpenPose" + bl_space_type = 'PROPERTIES' + bl_region_type = 'WINDOW' + bl_context = "data" + + @classmethod + def poll(cls, context): + return context.armature + + def draw_header(self, context): + bone = context.bone or context.edit_bone + if bone: + self.layout.prop(bone.dream_textures_openpose, "enabled", text="") + + def draw(self, context): + layout = self.layout + + armature = context.armature + + p = armature.dream_textures_openpose + + row = layout.row() + row.prop(p, "EAR_L", toggle=True) + row.prop(p, "EYE_L", toggle=True) + row.prop(p, "EYE_R", toggle=True) + row.prop(p, "EAR_R", toggle=True) + layout.prop(p, "NOSE", toggle=True) + row = layout.row() + row.prop(p, "SHOULDER_L", toggle=True) + row.prop(p, "CHEST", toggle=True) + row.prop(p, "SHOULDER_R", toggle=True) + row = layout.row() + row.prop(p, "ELBOW_L", toggle=True) + row.separator() + row.prop(p, "HIP_L", toggle=True) + row.prop(p, "HIP_R", toggle=True) + row.separator() + row.prop(p, "ELBOW_R", toggle=True) + row = layout.row() + row.prop(p, "HAND_L", toggle=True) + row.separator() + row.prop(p, "KNEE_L", toggle=True) + row.prop(p, "KNEE_R", toggle=True) + row.separator() + row.prop(p, "HAND_R", toggle=True) + row = layout.row() + row.prop(p, "FOOT_L", toggle=True) + row.prop(p, "FOOT_R", toggle=True) + + yield OpenPoseArmaturePanel class OpenPoseBonePanel(bpy.types.Panel): - bl_idname = "DREAM_PT_dream_textures_engine_bone_properties" + bl_idname = "DREAM_PT_dream_textures_bone_openpose" bl_label = "OpenPose" bl_space_type = 'PROPERTIES' bl_region_type = 'WINDOW' @@ -187,7 +240,7 @@ def poll(cls, context): def draw_header(self, context): bone = context.bone or context.edit_bone if bone: - self.layout.prop(bone, "dream_textures_openpose", text="") + self.layout.prop(bone.dream_textures_openpose, "enabled", text="") def draw(self, context): layout = self.layout @@ -195,9 +248,8 @@ def draw(self, context): bone = context.bone or context.edit_bone - if bone: - layout.enabled = bone.dream_textures_openpose - layout.prop(bone, "dream_textures_openpose_bone") - layout.prop(bone, "dream_textures_openpose_bone_side") + layout.enabled = bone.dream_textures_openpose.enabled + layout.prop(bone.dream_textures_openpose, "bone") + layout.prop(bone.dream_textures_openpose, "side") yield OpenPoseBonePanel \ No newline at end of file From f26040198cda43e22f10e8b0578fe3c5e9cba539 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Wed, 15 Mar 2023 00:23:57 -0400 Subject: [PATCH 10/28] Proper aspect ratio handling --- engine/annotations/openpose.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/engine/annotations/openpose.py b/engine/annotations/openpose.py index 23146d21..048b1abc 100644 --- a/engine/annotations/openpose.py +++ b/engine/annotations/openpose.py @@ -159,12 +159,25 @@ def render_openpose_map(context, collection=None): } with gpu.matrix.push_pop(): + ratio = width / height + projection_matrix = mathutils.Matrix(( + (1 / ratio, 0, 0, 0), + (0, 1, 0, 0), + (0, 0, -1, 0), + (0, 0, 0, 1) + )) gpu.matrix.load_matrix(mathutils.Matrix.Identity(4)) - gpu.matrix.load_projection_matrix(mathutils.Matrix.Identity(4)) + gpu.matrix.load_projection_matrix(projection_matrix) gpu.state.blend_set('ALPHA') + def transform(x, y): + return ( + (x - 0.5) * 2 * ratio, + (y - 0.5) * 2 + ) + shader = gpu.shader.from_builtin('3D_UNIFORM_COLOR') - batch = batch_for_shader(shader, 'TRI_STRIP', {"pos": [(-1, -1, 0), (-1, 1, 0), (1, -1, 0), (1, 1, 0)]}) + batch = batch_for_shader(shader, 'TRI_STRIP', {"pos": [(-ratio, -1, 0), (-ratio, 1, 0), (ratio, -1, 0), (ratio, 1, 0)]}) shader.uniform_float("color", (0, 0, 0, 1)) batch.draw(shader) @@ -180,13 +193,13 @@ def render_openpose_map(context, collection=None): continue a = bpy_extras.object_utils.world_to_camera_view(context.scene, context.scene.camera, object.matrix_world @ (a.tail if a_side == Side.TAIL else a.head)) b = bpy_extras.object_utils.world_to_camera_view(context.scene, context.scene.camera, object.matrix_world @ (b.tail if b_side == Side.TAIL else b.head)) - draw_ellipse_2d(((a[0] - 0.5) * 2, (a[1] - 0.5) * 2), ((b[0] - 0.5) * 2, (b[1] - 0.5) * 2), .015, 32, (color[0] / 255.0, color[1] / 255.0, color[2] / 255.0, 0.5)) + draw_ellipse_2d(transform(a[0], a[1]), transform(b[0], b[1]), .015, 32, (color[0] / 255.0, color[1] / 255.0, color[2] / 255.0, 0.5)) for b in Bone: bone, side = b.identify(object.data, object.pose) color = b.color() if bone is None: continue tail = bpy_extras.object_utils.world_to_camera_view(context.scene, context.scene.camera, object.matrix_world @ (bone.tail if side == Side.TAIL else bone.head)) - draw_circle_2d(((tail[0] - 0.5) * 2, (tail[1] - 0.5) * 2), .015, 16, (color[0] / 255.0, color[1] / 255.0, color[2] / 255.0, 0.5)) + draw_circle_2d(transform(tail[0], tail[1]), .015, 16, (color[0] / 255.0, color[1] / 255.0, color[2] / 255.0, 0.5)) depth = np.array(fb.read_color(0, 0, width, height, 4, 0, 'FLOAT').to_list()) offscreen.free() From d67fd3de4efb121229df3929f7d1e05882d9752c Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Wed, 15 Mar 2023 00:38:47 -0400 Subject: [PATCH 11/28] Improve bone panels --- engine/engine.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/engine/engine.py b/engine/engine.py index 7bff258a..28cbfa82 100644 --- a/engine/engine.py +++ b/engine/engine.py @@ -235,10 +235,10 @@ class OpenPoseBonePanel(bpy.types.Panel): @classmethod def poll(cls, context): - return (context.bone or context.edit_bone) and context.scene.render.engine == 'DREAM_TEXTURES' + return context.bone and context.scene.render.engine == 'DREAM_TEXTURES' def draw_header(self, context): - bone = context.bone or context.edit_bone + bone = context.bone if bone: self.layout.prop(bone.dream_textures_openpose, "enabled", text="") @@ -246,7 +246,7 @@ def draw(self, context): layout = self.layout layout.use_property_split = True - bone = context.bone or context.edit_bone + bone = context.bone layout.enabled = bone.dream_textures_openpose.enabled layout.prop(bone.dream_textures_openpose, "bone") From 6df01fe8f834b989b23d199a2cf05278353b4b4a Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Wed, 15 Mar 2023 21:11:58 -0400 Subject: [PATCH 12/28] Fix depth transformation --- engine/nodes/input_nodes.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/engine/nodes/input_nodes.py b/engine/nodes/input_nodes.py index acf11133..3c88f4c0 100644 --- a/engine/nodes/input_nodes.py +++ b/engine/nodes/input_nodes.py @@ -1,7 +1,6 @@ import bpy import gpu from gpu_extras.batch import batch_for_shader -import math import numpy as np from ..node import DreamTexturesNode from ..annotations import openpose @@ -109,7 +108,7 @@ def render_depth_map(cls, context, collection=None): for object in (context.scene.objects if collection is None else collection.objects): try: - mesh = object.to_mesh(depsgraph=context) + mesh = object.to_mesh(depsgraph=context).copy() except: continue if mesh is None: @@ -117,6 +116,7 @@ def render_depth_map(cls, context, collection=None): vertices = np.empty((len(mesh.vertices), 3), 'f') indices = np.empty((len(mesh.loop_triangles), 3), 'i') + mesh.transform(object.matrix_world) mesh.vertices.foreach_get("co", np.reshape(vertices, len(mesh.vertices) * 3)) mesh.loop_triangles.foreach_get("vertices", np.reshape(indices, len(mesh.loop_triangles) * 3)) From ee6d2c663f39cba3f0c821c1413945b5ad6037e9 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Wed, 15 Mar 2023 22:38:56 -0400 Subject: [PATCH 13/28] Add multi-controlnet support --- engine/nodes/pipeline_nodes.py | 10 +-- generator_process/actions/control_net.py | 80 ++++++++++++++++-------- operators/dream_texture.py | 2 +- operators/project.py | 2 +- property_groups/dream_prompt.py | 3 + 5 files changed, 65 insertions(+), 32 deletions(-) diff --git a/engine/nodes/pipeline_nodes.py b/engine/nodes/pipeline_nodes.py index f1684528..5382781e 100644 --- a/engine/nodes/pipeline_nodes.py +++ b/engine/nodes/pipeline_nodes.py @@ -105,8 +105,8 @@ def execute(self, context, prompt, negative_prompt, width, height, steps, seed, shared_args = context.scene.dream_textures_engine_prompt.generate_args() if controlnets is not None: - if not isinstance(controlnets, ControlNet): - controlnets = controlnets[0] + if not isinstance(controlnets, list): + controlnets = [controlnets] result = Generator.shared().control_net( pipeline=args['pipeline'], model=args['model'], @@ -115,9 +115,9 @@ def execute(self, context, prompt, negative_prompt, width, height, steps, seed, seamless_axes=args['seamless_axes'], iterations=args['iterations'], - control_net=controlnets.model, - control=controlnets.control(context), - controlnet_conditioning_scale=controlnets.conditioning_scale, + control_net=[c.model for c in controlnets], + control=[c.control(context) for c in controlnets], + controlnet_conditioning_scale=[c.conditioning_scale for c in controlnets], image=np.uint8(source_image * 255) if self.task == 'image_to_image' else None, strength=noise_strength, diff --git a/generator_process/actions/control_net.py b/generator_process/actions/control_net.py index 1031db54..7cf15971 100644 --- a/generator_process/actions/control_net.py +++ b/generator_process/actions/control_net.py @@ -18,9 +18,9 @@ def control_net( optimizations: Optimizations, - control_net: str, - control: NDArray | None, - controlnet_conditioning_scale: float, + control_net: list[str], + control: list[NDArray] | None, + controlnet_conditioning_scale: list[float], image: NDArray | str | None, strength: float, prompt: str | list[str], @@ -43,6 +43,7 @@ def control_net( match pipeline: case Pipeline.STABLE_DIFFUSION: import diffusers + from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import MultiControlNetModel, ControlNetModel import torch import PIL.Image import PIL.ImageOps @@ -69,7 +70,7 @@ def __call__( callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, callback_steps: int = 1, cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: float = 1.0, + controlnet_conditioning_scale: Union[float, List[float]] = 1.0, **kwargs ): @@ -78,7 +79,15 @@ def __call__( # 1. Check inputs. Raise error if not correct self.check_inputs( - prompt, image, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds + prompt, + image, + height, + width, + callback_steps, + negative_prompt, + prompt_embeds, + negative_prompt_embeds, + controlnet_conditioning_scale ) # 2. Define call parameters @@ -95,6 +104,9 @@ def __call__( # corresponds to doing no classifier free guidance. do_classifier_free_guidance = guidance_scale > 1.0 + if isinstance(self.controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float): + controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(self.controlnet.nets) + # 3. Encode input prompt prompt_embeds = self._encode_prompt( prompt, @@ -107,18 +119,37 @@ def __call__( ) # 4. Prepare image - image = self.prepare_image( - image, - width, - height, - batch_size * num_images_per_prompt, - num_images_per_prompt, - device, - self.controlnet.dtype, - ) + if isinstance(self.controlnet, ControlNetModel): + image = self.prepare_image( + image=image, + width=width, + height=height, + batch_size=batch_size * num_images_per_prompt, + num_images_per_prompt=num_images_per_prompt, + device=device, + dtype=self.controlnet.dtype, + do_classifier_free_guidance=do_classifier_free_guidance, + ) + elif isinstance(self.controlnet, MultiControlNetModel): + images = [] + + for image_ in image: + image_ = self.prepare_image( + image=image_, + width=width, + height=height, + batch_size=batch_size * num_images_per_prompt, + num_images_per_prompt=num_images_per_prompt, + device=device, + dtype=self.controlnet.dtype, + do_classifier_free_guidance=do_classifier_free_guidance, + ) + + images.append(image_) - if do_classifier_free_guidance: - image = torch.cat([image] * 2) + image = images + else: + assert False # 5. Prepare timesteps self.scheduler.set_timesteps(num_inference_steps, device=device) @@ -148,20 +179,16 @@ def __call__( latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + # controlnet(s) inference down_block_res_samples, mid_block_res_sample = self.controlnet( latent_model_input, t, encoder_hidden_states=prompt_embeds, controlnet_cond=image, + conditioning_scale=controlnet_conditioning_scale, return_dict=False, ) - down_block_res_samples = [ - down_block_res_sample * controlnet_conditioning_scale - for down_block_res_sample in down_block_res_samples - ] - mid_block_res_sample *= controlnet_conditioning_scale - # predict the noise residual noise_pred = self.unet( latent_model_input, @@ -230,7 +257,10 @@ def __call__( device = self.choose_device() # Load the ControlNet model - controlnet = load_pipe(self, "control_net_model", diffusers.ControlNetModel, control_net, optimizations, None, device) + controlnet = [] + for controlnet_name in control_net: + controlnet.append(load_pipe(self, f"control_net_model-{controlnet_name}", diffusers.ControlNetModel, controlnet_name, optimizations, None, device)) + controlnet = MultiControlNetModel(controlnet) # StableDiffusionPipeline w/ caching pipe = load_pipe(self, "control_net", GeneratorPipeline, model, optimizations, scheduler, device, controlnet=controlnet) @@ -257,13 +287,13 @@ def __call__( int(8 * (height // 8)), ) print(control) - control_image = PIL.Image.fromarray(np.uint8(control * 255)).convert('RGB').resize(rounded_size) if control is not None else None + control_image = [PIL.Image.fromarray(np.uint8(c * 255)).convert('RGB').resize(rounded_size) for c in control] if control is not None else None init_image = None if image is None else (PIL.Image.open(image) if isinstance(image, str) else PIL.Image.fromarray(image.astype(np.uint8))).convert('RGB').resize(rounded_size) # Seamless if seamless_axes == SeamlessAxes.AUTO: init_sa = None if init_image is None else self.detect_seamless(np.array(init_image) / 255) - control_sa = None if control_image is None else self.detect_seamless(np.array(control_image) / 255) + control_sa = None if control_image is None else self.detect_seamless(np.array(control_image[0]) / 255) if init_sa is not None and control_sa is not None: seamless_axes = SeamlessAxes((init_sa.x and control_sa.x, init_sa.y and control_sa.y)) elif init_sa is not None: diff --git a/operators/dream_texture.py b/operators/dream_texture.py index 8f482f7c..c8b9e4c2 100644 --- a/operators/dream_texture.py +++ b/operators/dream_texture.py @@ -213,7 +213,7 @@ def require_depth(): case 'control_net': f = gen.control_net( image=None, - control=init_image, + control=[init_image], **generated_args ) case 'inpaint': diff --git a/operators/project.py b/operators/project.py index d1ef31d2..e6cb483d 100644 --- a/operators/project.py +++ b/operators/project.py @@ -440,7 +440,7 @@ def on_exception(_, exception): context.scene.dream_textures_info = "Starting..." if context.scene.dream_textures_project_use_control_net: future = gen.control_net( - control=np.flipud(depth), # the depth control needs to be flipped. + control=[np.flipud(depth)], # the depth control needs to be flipped. image=init_img_path, **context.scene.dream_textures_project_prompt.generate_args() ) diff --git a/property_groups/dream_prompt.py b/property_groups/dream_prompt.py index 90e86047..c64001d7 100644 --- a/property_groups/dream_prompt.py +++ b/property_groups/dream_prompt.py @@ -278,6 +278,9 @@ def generate_args(self): args['seamless_axes'] = SeamlessAxes(args['seamless_axes']) args['width'] = args['width'] if args['use_size'] else None args['height'] = args['height'] if args['use_size'] else None + + args['control_net'] = [args['control_net']] + args['controlnet_conditioning_scale'] = [args['controlnet_conditioning_scale']] return args DreamPrompt.generate_prompt = generate_prompt From 2a0386babedec6519d3508d38ac57f4ffc092643 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Wed, 15 Mar 2023 22:50:26 -0400 Subject: [PATCH 14/28] Evaluate object modifiers --- engine/annotations/openpose.py | 1 + engine/nodes/input_nodes.py | 1 + 2 files changed, 2 insertions(+) diff --git a/engine/annotations/openpose.py b/engine/annotations/openpose.py index 048b1abc..5f7f435e 100644 --- a/engine/annotations/openpose.py +++ b/engine/annotations/openpose.py @@ -182,6 +182,7 @@ def transform(x, y): batch.draw(shader) for object in (context.scene.objects if collection is None else collection.objects): + object = object.evaluated_get(context) if object.hide_render: continue if object.pose is None: diff --git a/engine/nodes/input_nodes.py b/engine/nodes/input_nodes.py index 3c88f4c0..6e3c2390 100644 --- a/engine/nodes/input_nodes.py +++ b/engine/nodes/input_nodes.py @@ -107,6 +107,7 @@ def render_depth_map(cls, context, collection=None): shader = gpu.shader.from_builtin('3D_UNIFORM_COLOR') for object in (context.scene.objects if collection is None else collection.objects): + object = object.evaluated_get(context) try: mesh = object.to_mesh(depsgraph=context).copy() except: From e06b39680cba72b34ac8e6ffaf841a1d67d9aa1a Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Thu, 16 Mar 2023 23:07:27 -0400 Subject: [PATCH 15/28] Refactor into annotation nodes --- engine/__init__.py | 18 +++++++-- engine/annotations/depth.py | 55 ++++++++++++++++++++++++++ engine/nodes/annotation_nodes.py | 39 ++++++++++++++++++ engine/nodes/input_nodes.py | 68 +------------------------------- engine/nodes/pipeline_nodes.py | 8 ++-- engine/nodes/utility_nodes.py | 19 +++++++++ 6 files changed, 133 insertions(+), 74 deletions(-) create mode 100644 engine/annotations/depth.py create mode 100644 engine/nodes/annotation_nodes.py diff --git a/engine/__init__.py b/engine/__init__.py index 58a073a4..dbbdc72c 100644 --- a/engine/__init__.py +++ b/engine/__init__.py @@ -5,6 +5,7 @@ from .nodes.input_nodes import * from .nodes.pipeline_nodes import * from .nodes.utility_nodes import * +from .nodes.annotation_nodes import * from .annotations import openpose import bpy @@ -24,11 +25,15 @@ def poll(cls, context): nodeitems_utils.NodeItem(NodeString.bl_idname), nodeitems_utils.NodeItem(NodeImage.bl_idname), nodeitems_utils.NodeItem(NodeCollection.bl_idname), - nodeitems_utils.NodeItem(NodeSceneInfo.bl_idname), ]), DreamTexturesNodeCategory("DREAM_TEXTURES_UTILITY", "Utilities", items = [ nodeitems_utils.NodeItem(NodeMath.bl_idname), nodeitems_utils.NodeItem(NodeRandomValue.bl_idname), + nodeitems_utils.NodeItem(NodeClamp.bl_idname), + ]), + DreamTexturesNodeCategory("DREAM_TEXTURES_ANNOTATIONS", "Annotations", items = [ + nodeitems_utils.NodeItem(NodeAnnotationDepth.bl_idname), + nodeitems_utils.NodeItem(NodeAnnotationOpenPose.bl_idname), ]), DreamTexturesNodeCategory("DREAM_TEXTURES_GROUP", "Group", items = [ nodeitems_utils.NodeItem(bpy.types.NodeGroupOutput.__name__), @@ -59,11 +64,14 @@ def register(): bpy.utils.register_class(NodeInteger) bpy.utils.register_class(NodeString) bpy.utils.register_class(NodeCollection) - bpy.utils.register_class(NodeSceneInfo) bpy.utils.register_class(NodeImage) + + bpy.utils.register_class(NodeAnnotationDepth) + bpy.utils.register_class(NodeAnnotationOpenPose) bpy.utils.register_class(NodeMath) bpy.utils.register_class(NodeRandomValue) + bpy.utils.register_class(NodeClamp) nodeitems_utils.register_node_categories("DREAM_TEXTURES_CATEGORIES", categories) @@ -84,11 +92,13 @@ def unregister(): bpy.utils.unregister_class(NodeInteger) bpy.utils.unregister_class(NodeString) bpy.utils.unregister_class(NodeCollection) - bpy.utils.unregister_class(NodeSceneInfo) bpy.utils.unregister_class(NodeImage) + + bpy.utils.unregister_class(NodeAnnotationDepth) + bpy.utils.unregister_class(NodeAnnotationOpenPose) bpy.utils.unregister_class(NodeMath) bpy.utils.unregister_class(NodeRandomValue) - + bpy.utils.unregister_class(NodeClamp) nodeitems_utils.unregister_node_categories("DREAM_TEXTURES_CATEGORIES") \ No newline at end of file diff --git a/engine/annotations/depth.py b/engine/annotations/depth.py new file mode 100644 index 00000000..b9e38690 --- /dev/null +++ b/engine/annotations/depth.py @@ -0,0 +1,55 @@ +import bpy +import gpu +from gpu_extras.batch import batch_for_shader +import numpy as np + +def render_depth_map(context, collection=None, invert=True): + width, height = context.scene.render.resolution_x, context.scene.render.resolution_y + matrix = context.scene.camera.matrix_world.inverted() + projection_matrix = context.scene.camera.calc_matrix_camera( + context, + x=width, + y=height + ) + offscreen = gpu.types.GPUOffScreen(width, height) + + with offscreen.bind(): + fb = gpu.state.active_framebuffer_get() + fb.clear(color=(0.0, 0.0, 0.0, 0.0)) + gpu.state.depth_test_set('LESS_EQUAL') + gpu.state.depth_mask_set(True) + with gpu.matrix.push_pop(): + gpu.matrix.load_matrix(matrix) + gpu.matrix.load_projection_matrix(projection_matrix) + + shader = gpu.shader.from_builtin('3D_UNIFORM_COLOR') + + for object in (context.scene.objects if collection is None else collection.objects): + object = object.evaluated_get(context) + try: + mesh = object.to_mesh(depsgraph=context).copy() + except: + continue + if mesh is None: + continue + vertices = np.empty((len(mesh.vertices), 3), 'f') + indices = np.empty((len(mesh.loop_triangles), 3), 'i') + + mesh.transform(object.matrix_world) + mesh.vertices.foreach_get("co", np.reshape(vertices, len(mesh.vertices) * 3)) + mesh.loop_triangles.foreach_get("vertices", np.reshape(indices, len(mesh.loop_triangles) * 3)) + + batch = batch_for_shader( + shader, 'TRIS', + {"pos": vertices}, + indices=indices, + ) + batch.draw(shader) + depth = np.array(fb.read_depth(0, 0, width, height).to_list()) + if invert: + depth = 1 - depth + mask = np.array(fb.read_color(0, 0, width, height, 4, 0, 'UBYTE').to_list())[:, :, 3] + depth *= mask + depth = np.interp(depth, [np.ma.masked_equal(depth, 0, copy=False).min(), depth.max()], [0, 1]).clip(0, 1) + offscreen.free() + return depth \ No newline at end of file diff --git a/engine/nodes/annotation_nodes.py b/engine/nodes/annotation_nodes.py new file mode 100644 index 00000000..fcd03e0f --- /dev/null +++ b/engine/nodes/annotation_nodes.py @@ -0,0 +1,39 @@ +import bpy +from ..node import DreamTexturesNode +from ..annotations import depth +from ..annotations import openpose + +class NodeAnnotationDepth(DreamTexturesNode): + bl_idname = "dream_textures.node_annotation_depth" + bl_label = "Depth Map" + + def init(self, context): + self.inputs.new("NodeSocketCollection", "Collection") + self.inputs.new("NodeSocketBool", "Invert") + + self.outputs.new("NodeSocketColor", "Depth Map") + + def draw_buttons(self, context, layout): + pass + + def execute(self, context, collection, invert): + return { + 'Depth Map': depth.render_depth_map(context, collection=collection, invert=invert), + } + +class NodeAnnotationOpenPose(DreamTexturesNode): + bl_idname = "dream_textures.node_annotation_openpose" + bl_label = "OpenPose Map" + + def init(self, context): + self.inputs.new("NodeSocketCollection", "Collection") + + self.outputs.new("NodeSocketColor", "OpenPose Map") + + def draw_buttons(self, context, layout): + pass + + def execute(self, context, collection): + return { + 'OpenPose Map': openpose.render_openpose_map(context, collection=collection) + } \ No newline at end of file diff --git a/engine/nodes/input_nodes.py b/engine/nodes/input_nodes.py index 6e3c2390..19282011 100644 --- a/engine/nodes/input_nodes.py +++ b/engine/nodes/input_nodes.py @@ -4,6 +4,7 @@ import numpy as np from ..node import DreamTexturesNode from ..annotations import openpose +from ..annotations import depth class NodeString(DreamTexturesNode): bl_idname = "dream_textures.node_string" @@ -63,7 +64,7 @@ class NodeImage(DreamTexturesNode): value: bpy.props.PointerProperty(type=bpy.types.Image) def init(self, context): - self.outputs.new("NodeSocketImage", "Image") + self.outputs.new("NodeSocketColor", "Image") def draw_buttons(self, context, layout): layout.prop(self, "value", text="") @@ -71,69 +72,4 @@ def draw_buttons(self, context, layout): def execute(self, context): return { 'Image': np.array(self.value.pixels).reshape((*self.value.size, self.value.channels)) - } - -class NodeSceneInfo(DreamTexturesNode): - bl_idname = "dream_textures.node_scene" - bl_label = "Scene Info" - - def init(self, context): - self.outputs.new("NodeSocketImage", "Depth Map") - self.outputs.new("NodeSocketImage", "OpenPose Map") - - def draw_buttons(self, context, layout): - pass - - @classmethod - def render_depth_map(cls, context, collection=None): - width, height = context.scene.render.resolution_x, context.scene.render.resolution_y - matrix = context.scene.camera.matrix_world.inverted() - projection_matrix = context.scene.camera.calc_matrix_camera( - context, - x=width, - y=height - ) - offscreen = gpu.types.GPUOffScreen(width, height) - - with offscreen.bind(): - fb = gpu.state.active_framebuffer_get() - fb.clear(color=(0.0, 0.0, 0.0, 0.0)) - gpu.state.depth_test_set('LESS_EQUAL') - gpu.state.depth_mask_set(True) - with gpu.matrix.push_pop(): - gpu.matrix.load_matrix(matrix) - gpu.matrix.load_projection_matrix(projection_matrix) - - shader = gpu.shader.from_builtin('3D_UNIFORM_COLOR') - - for object in (context.scene.objects if collection is None else collection.objects): - object = object.evaluated_get(context) - try: - mesh = object.to_mesh(depsgraph=context).copy() - except: - continue - if mesh is None: - continue - vertices = np.empty((len(mesh.vertices), 3), 'f') - indices = np.empty((len(mesh.loop_triangles), 3), 'i') - - mesh.transform(object.matrix_world) - mesh.vertices.foreach_get("co", np.reshape(vertices, len(mesh.vertices) * 3)) - mesh.loop_triangles.foreach_get("vertices", np.reshape(indices, len(mesh.loop_triangles) * 3)) - - batch = batch_for_shader( - shader, 'TRIS', - {"pos": vertices}, - indices=indices, - ) - batch.draw(shader) - depth = np.array(fb.read_depth(0, 0, width, height).to_list()) - depth = np.interp(depth, [np.ma.masked_equal(depth, 0, copy=False).min(), depth.max()], [0, 1]).clip(0, 1) - offscreen.free() - return depth - - def execute(self, context): - return { - 'Depth Map': NodeSceneInfo.render_depth_map(context), - 'OpenPose Map': openpose.render_openpose_map(context) } \ No newline at end of file diff --git a/engine/nodes/pipeline_nodes.py b/engine/nodes/pipeline_nodes.py index 5382781e..c2e127f4 100644 --- a/engine/nodes/pipeline_nodes.py +++ b/engine/nodes/pipeline_nodes.py @@ -7,8 +7,8 @@ from ...generator_process import Generator from ...generator_process.actions.prompt_to_image import StepPreviewMode from ...property_groups.dream_prompt import DreamPrompt, control_net_options -from .input_nodes import NodeSceneInfo from ..annotations import openpose +from ..annotations import depth class NodeSocketControlNet(bpy.types.NodeSocket): bl_idname = "NodeSocketControlNet" @@ -42,7 +42,7 @@ def control(self, context): else: match self.control_type: case ControlType.DEPTH: - return np.flipud(NodeSceneInfo.render_depth_map(context, collection=self.collection)) + return np.flipud(depth.render_depth_map(context, collection=self.collection)) case ControlType.OPENPOSE: return np.flipud(openpose.render_openpose_map(context, collection=self.collection)) case ControlType.NORMAL: @@ -67,8 +67,8 @@ class NodeStableDiffusion(DreamTexturesNode): ), update=_update_stable_diffusion_sockets) def init(self, context): - self.inputs.new("NodeSocketImage", "Depth Map") - self.inputs.new("NodeSocketImage", "Source Image") + self.inputs.new("NodeSocketColor", "Depth Map") + self.inputs.new("NodeSocketColor", "Source Image") self.inputs.new("NodeSocketFloat", "Noise Strength").default_value = 0.75 self.inputs.new("NodeSocketString", "Prompt") diff --git a/engine/nodes/utility_nodes.py b/engine/nodes/utility_nodes.py index 7e6f711a..af787862 100644 --- a/engine/nodes/utility_nodes.py +++ b/engine/nodes/utility_nodes.py @@ -62,4 +62,23 @@ def draw_buttons(self, context, layout): def execute(self, context, min, max): return { 'Value': random.randrange(min, max) + } + +class NodeClamp(DreamTexturesNode): + bl_idname = "dream_textures.node_clamp" + bl_label = "Clamp" + + def init(self, context): + self.inputs.new("NodeSocketFloat", "Value") + self.inputs.new("NodeSocketFloat", "Min") + self.inputs.new("NodeSocketFloat", "Max") + + self.outputs.new("NodeSocketFloat", "Result") + + def draw_buttons(self, context, layout): + pass + + def execute(self, context, value, min, max): + return { + 'Result': np.clip(value, min, max) } \ No newline at end of file From 7583e38014d9cb209e19fb39b97a340fe453332c Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Fri, 17 Mar 2023 00:38:26 -0400 Subject: [PATCH 16/28] Improve responsiveness --- engine/annotations/depth.py | 101 +++++++------- engine/annotations/openpose.py | 159 ++++++++++++----------- engine/engine.py | 21 +-- engine/node_executor.py | 61 ++++++--- engine/nodes/annotation_nodes.py | 8 +- engine/nodes/pipeline_nodes.py | 50 ++++--- generator_process/actions/control_net.py | 1 - 7 files changed, 231 insertions(+), 170 deletions(-) diff --git a/engine/annotations/depth.py b/engine/annotations/depth.py index b9e38690..c7da7416 100644 --- a/engine/annotations/depth.py +++ b/engine/annotations/depth.py @@ -2,54 +2,63 @@ import gpu from gpu_extras.batch import batch_for_shader import numpy as np +import threading def render_depth_map(context, collection=None, invert=True): - width, height = context.scene.render.resolution_x, context.scene.render.resolution_y - matrix = context.scene.camera.matrix_world.inverted() - projection_matrix = context.scene.camera.calc_matrix_camera( - context, - x=width, - y=height - ) - offscreen = gpu.types.GPUOffScreen(width, height) + e = threading.Event() + result = None + def _execute(): + nonlocal result + width, height = context.scene.render.resolution_x, context.scene.render.resolution_y + matrix = context.scene.camera.matrix_world.inverted() + projection_matrix = context.scene.camera.calc_matrix_camera( + context, + x=width, + y=height + ) + offscreen = gpu.types.GPUOffScreen(width, height) - with offscreen.bind(): - fb = gpu.state.active_framebuffer_get() - fb.clear(color=(0.0, 0.0, 0.0, 0.0)) - gpu.state.depth_test_set('LESS_EQUAL') - gpu.state.depth_mask_set(True) - with gpu.matrix.push_pop(): - gpu.matrix.load_matrix(matrix) - gpu.matrix.load_projection_matrix(projection_matrix) - - shader = gpu.shader.from_builtin('3D_UNIFORM_COLOR') + with offscreen.bind(): + fb = gpu.state.active_framebuffer_get() + fb.clear(color=(0.0, 0.0, 0.0, 0.0)) + gpu.state.depth_test_set('LESS_EQUAL') + gpu.state.depth_mask_set(True) + with gpu.matrix.push_pop(): + gpu.matrix.load_matrix(matrix) + gpu.matrix.load_projection_matrix(projection_matrix) + + shader = gpu.shader.from_builtin('3D_UNIFORM_COLOR') - for object in (context.scene.objects if collection is None else collection.objects): - object = object.evaluated_get(context) - try: - mesh = object.to_mesh(depsgraph=context).copy() - except: - continue - if mesh is None: - continue - vertices = np.empty((len(mesh.vertices), 3), 'f') - indices = np.empty((len(mesh.loop_triangles), 3), 'i') + for object in (context.scene.objects if collection is None else collection.objects): + object = object.evaluated_get(context) + try: + mesh = object.to_mesh(depsgraph=context).copy() + except: + continue + if mesh is None: + continue + vertices = np.empty((len(mesh.vertices), 3), 'f') + indices = np.empty((len(mesh.loop_triangles), 3), 'i') - mesh.transform(object.matrix_world) - mesh.vertices.foreach_get("co", np.reshape(vertices, len(mesh.vertices) * 3)) - mesh.loop_triangles.foreach_get("vertices", np.reshape(indices, len(mesh.loop_triangles) * 3)) - - batch = batch_for_shader( - shader, 'TRIS', - {"pos": vertices}, - indices=indices, - ) - batch.draw(shader) - depth = np.array(fb.read_depth(0, 0, width, height).to_list()) - if invert: - depth = 1 - depth - mask = np.array(fb.read_color(0, 0, width, height, 4, 0, 'UBYTE').to_list())[:, :, 3] - depth *= mask - depth = np.interp(depth, [np.ma.masked_equal(depth, 0, copy=False).min(), depth.max()], [0, 1]).clip(0, 1) - offscreen.free() - return depth \ No newline at end of file + mesh.transform(object.matrix_world) + mesh.vertices.foreach_get("co", np.reshape(vertices, len(mesh.vertices) * 3)) + mesh.loop_triangles.foreach_get("vertices", np.reshape(indices, len(mesh.loop_triangles) * 3)) + + batch = batch_for_shader( + shader, 'TRIS', + {"pos": vertices}, + indices=indices, + ) + batch.draw(shader) + depth = np.array(fb.read_depth(0, 0, width, height).to_list()) + if invert: + depth = 1 - depth + mask = np.array(fb.read_color(0, 0, width, height, 4, 0, 'UBYTE').to_list())[:, :, 3] + depth *= mask + depth = np.interp(depth, [np.ma.masked_equal(depth, 0, copy=False).min(), depth.max()], [0, 1]).clip(0, 1) + offscreen.free() + result = depth + e.set() + bpy.app.timers.register(_execute, first_interval=0) + e.wait() + return result \ No newline at end of file diff --git a/engine/annotations/openpose.py b/engine/annotations/openpose.py index 5f7f435e..7de2483d 100644 --- a/engine/annotations/openpose.py +++ b/engine/annotations/openpose.py @@ -6,6 +6,7 @@ import numpy as np import enum import math +import threading class Side(enum.IntEnum): HEAD = 0 @@ -129,82 +130,90 @@ class BoneOpenPoseData(bpy.types.PropertyGroup): }) def render_openpose_map(context, collection=None): - width, height = context.scene.render.resolution_x, context.scene.render.resolution_y - offscreen = gpu.types.GPUOffScreen(width, height) - - with offscreen.bind(): - fb = gpu.state.active_framebuffer_get() - fb.clear(color=(0.0, 0.0, 0.0, 0.0)) - gpu.state.depth_test_set('LESS_EQUAL') - gpu.state.depth_mask_set(True) - - lines = { - (Bone.NOSE, Bone.CHEST): (0, 0, 255), - (Bone.CHEST, Bone.SHOULDER_L): (255, 85, 0), - (Bone.CHEST, Bone.SHOULDER_R): (255, 0, 0), - (Bone.SHOULDER_L, Bone.ELBOW_L): (170, 255, 0), - (Bone.SHOULDER_R, Bone.ELBOW_R): (255, 170, 0), - (Bone.ELBOW_L, Bone.HAND_L): (85, 255, 0), - (Bone.ELBOW_R, Bone.HAND_R): (255, 255, 0), - (Bone.CHEST, Bone.HIP_L): (0, 255, 255), - (Bone.CHEST, Bone.HIP_R): (0, 255, 0), - (Bone.HIP_L, Bone.KNEE_L): (0, 170, 255), - (Bone.HIP_R, Bone.KNEE_R): (0, 255, 85), - (Bone.KNEE_L, Bone.FOOT_L): (0, 85, 255), - (Bone.KNEE_R, Bone.FOOT_R): (0, 255, 170), - (Bone.NOSE, Bone.EYE_L): (255, 0, 255), - (Bone.NOSE, Bone.EYE_R): (85, 0, 255), - (Bone.EYE_L, Bone.EAR_L): (255, 0, 170), - (Bone.EYE_R, Bone.EAR_R): (170, 0, 255), - } - - with gpu.matrix.push_pop(): - ratio = width / height - projection_matrix = mathutils.Matrix(( - (1 / ratio, 0, 0, 0), - (0, 1, 0, 0), - (0, 0, -1, 0), - (0, 0, 0, 1) - )) - gpu.matrix.load_matrix(mathutils.Matrix.Identity(4)) - gpu.matrix.load_projection_matrix(projection_matrix) - gpu.state.blend_set('ALPHA') - - def transform(x, y): - return ( - (x - 0.5) * 2 * ratio, - (y - 0.5) * 2 - ) - - shader = gpu.shader.from_builtin('3D_UNIFORM_COLOR') - batch = batch_for_shader(shader, 'TRI_STRIP', {"pos": [(-ratio, -1, 0), (-ratio, 1, 0), (ratio, -1, 0), (ratio, 1, 0)]}) - shader.uniform_float("color", (0, 0, 0, 1)) - batch.draw(shader) - - for object in (context.scene.objects if collection is None else collection.objects): - object = object.evaluated_get(context) - if object.hide_render: - continue - if object.pose is None: - continue - for connection, color in lines.items(): - a, a_side = connection[0].identify(object.data, object.pose) - b, b_side = connection[1].identify(object.data, object.pose) - if a is None or b is None: + e = threading.Event() + result = None + def _execute(): + nonlocal result + width, height = context.scene.render.resolution_x, context.scene.render.resolution_y + offscreen = gpu.types.GPUOffScreen(width, height) + + with offscreen.bind(): + fb = gpu.state.active_framebuffer_get() + fb.clear(color=(0.0, 0.0, 0.0, 0.0)) + gpu.state.depth_test_set('LESS_EQUAL') + gpu.state.depth_mask_set(True) + + lines = { + (Bone.NOSE, Bone.CHEST): (0, 0, 255), + (Bone.CHEST, Bone.SHOULDER_L): (255, 85, 0), + (Bone.CHEST, Bone.SHOULDER_R): (255, 0, 0), + (Bone.SHOULDER_L, Bone.ELBOW_L): (170, 255, 0), + (Bone.SHOULDER_R, Bone.ELBOW_R): (255, 170, 0), + (Bone.ELBOW_L, Bone.HAND_L): (85, 255, 0), + (Bone.ELBOW_R, Bone.HAND_R): (255, 255, 0), + (Bone.CHEST, Bone.HIP_L): (0, 255, 255), + (Bone.CHEST, Bone.HIP_R): (0, 255, 0), + (Bone.HIP_L, Bone.KNEE_L): (0, 170, 255), + (Bone.HIP_R, Bone.KNEE_R): (0, 255, 85), + (Bone.KNEE_L, Bone.FOOT_L): (0, 85, 255), + (Bone.KNEE_R, Bone.FOOT_R): (0, 255, 170), + (Bone.NOSE, Bone.EYE_L): (255, 0, 255), + (Bone.NOSE, Bone.EYE_R): (85, 0, 255), + (Bone.EYE_L, Bone.EAR_L): (255, 0, 170), + (Bone.EYE_R, Bone.EAR_R): (170, 0, 255), + } + + with gpu.matrix.push_pop(): + ratio = width / height + projection_matrix = mathutils.Matrix(( + (1 / ratio, 0, 0, 0), + (0, 1, 0, 0), + (0, 0, -1, 0), + (0, 0, 0, 1) + )) + gpu.matrix.load_matrix(mathutils.Matrix.Identity(4)) + gpu.matrix.load_projection_matrix(projection_matrix) + gpu.state.blend_set('ALPHA') + + def transform(x, y): + return ( + (x - 0.5) * 2 * ratio, + (y - 0.5) * 2 + ) + + shader = gpu.shader.from_builtin('3D_UNIFORM_COLOR') + batch = batch_for_shader(shader, 'TRI_STRIP', {"pos": [(-ratio, -1, 0), (-ratio, 1, 0), (ratio, -1, 0), (ratio, 1, 0)]}) + shader.uniform_float("color", (0, 0, 0, 1)) + batch.draw(shader) + + for object in (context.scene.objects if collection is None else collection.objects): + object = object.evaluated_get(context) + if object.hide_render: continue - a = bpy_extras.object_utils.world_to_camera_view(context.scene, context.scene.camera, object.matrix_world @ (a.tail if a_side == Side.TAIL else a.head)) - b = bpy_extras.object_utils.world_to_camera_view(context.scene, context.scene.camera, object.matrix_world @ (b.tail if b_side == Side.TAIL else b.head)) - draw_ellipse_2d(transform(a[0], a[1]), transform(b[0], b[1]), .015, 32, (color[0] / 255.0, color[1] / 255.0, color[2] / 255.0, 0.5)) - for b in Bone: - bone, side = b.identify(object.data, object.pose) - color = b.color() - if bone is None: continue - tail = bpy_extras.object_utils.world_to_camera_view(context.scene, context.scene.camera, object.matrix_world @ (bone.tail if side == Side.TAIL else bone.head)) - draw_circle_2d(transform(tail[0], tail[1]), .015, 16, (color[0] / 255.0, color[1] / 255.0, color[2] / 255.0, 0.5)) - - depth = np.array(fb.read_color(0, 0, width, height, 4, 0, 'FLOAT').to_list()) - offscreen.free() - return depth + if object.pose is None: + continue + for connection, color in lines.items(): + a, a_side = connection[0].identify(object.data, object.pose) + b, b_side = connection[1].identify(object.data, object.pose) + if a is None or b is None: + continue + a = bpy_extras.object_utils.world_to_camera_view(context.scene, context.scene.camera, object.matrix_world @ (a.tail if a_side == Side.TAIL else a.head)) + b = bpy_extras.object_utils.world_to_camera_view(context.scene, context.scene.camera, object.matrix_world @ (b.tail if b_side == Side.TAIL else b.head)) + draw_ellipse_2d(transform(a[0], a[1]), transform(b[0], b[1]), .015, 32, (color[0] / 255.0, color[1] / 255.0, color[2] / 255.0, 0.5)) + for b in Bone: + bone, side = b.identify(object.data, object.pose) + color = b.color() + if bone is None: continue + tail = bpy_extras.object_utils.world_to_camera_view(context.scene, context.scene.camera, object.matrix_world @ (bone.tail if side == Side.TAIL else bone.head)) + draw_circle_2d(transform(tail[0], tail[1]), .015, 16, (color[0] / 255.0, color[1] / 255.0, color[2] / 255.0, 0.5)) + + depth = np.array(fb.read_color(0, 0, width, height, 4, 0, 'FLOAT').to_list()) + offscreen.free() + result = depth + e.set() + bpy.app.timers.register(_execute, first_interval=0) + e.wait() + return result def draw_circle_2d(center, radius, segments, color): m = (1.0 / (segments - 1)) * (math.pi * 2) diff --git a/engine/engine.py b/engine/engine.py index 28cbfa82..5e3d6dc9 100644 --- a/engine/engine.py +++ b/engine/engine.py @@ -7,6 +7,8 @@ from .node_tree import DreamTexturesNodeTree from ..engine import node_executor from .annotations import openpose +import time +from threading import Event class DreamTexturesRenderEngine(bpy.types.RenderEngine): """A custom Dream Textures render engine, that uses Stable Diffusion and scene data to render images, instead of as a pass on top of Cycles.""" @@ -14,7 +16,7 @@ class DreamTexturesRenderEngine(bpy.types.RenderEngine): bl_idname = "DREAM_TEXTURES" bl_label = "Dream Textures" bl_use_preview = False - bl_use_gpu_context = True + # bl_use_gpu_context = True def __init__(self): pass @@ -39,19 +41,22 @@ def prepare_result(result): result = self.begin_result(0, 0, scene.render.resolution_x, scene.render.resolution_y) layer = result.layers[0].passes["Combined"] + self.update_result(result) try: progress = 0 - def update_result(node, result): - nonlocal progress - progress += 1 - if isinstance(result, np.ndarray): - node_result = prepare_result(result) + def node_begin(node): + self.update_stats("Node", node.name) + def node_update(response): + if isinstance(response, np.ndarray): + node_result = prepare_result(response) layer.rect = node_result.reshape(-1, node_result.shape[-1]) self.update_result(result) - self.update_stats("Node", node.name) + def node_end(_): + nonlocal progress + progress += 1 self.update_progress(progress / len(scene.dream_textures_render_engine.node_tree.nodes)) - node_result = node_executor.execute(scene.dream_textures_render_engine.node_tree, depsgraph, on_execute=update_result) + node_result = node_executor.execute(scene.dream_textures_render_engine.node_tree, depsgraph, node_begin=node_begin, node_update=node_update, node_end=node_end) node_result = prepare_result(node_result) except Exception as error: self.report({'ERROR'}, str(error)) diff --git a/engine/node_executor.py b/engine/node_executor.py index 815ae244..7bf8a52b 100644 --- a/engine/node_executor.py +++ b/engine/node_executor.py @@ -1,28 +1,49 @@ import bpy import numpy as np +from threading import Event +import graphlib # from dream_textures.engine import node_executor # node_executor.execute(bpy.data.node_groups["NodeTree"], bpy.context.evaluated_depsgraph_get()) -def execute_node(node, context, cache, on_execute=lambda _, __: None): - if node in cache: - return cache[node] - kwargs = { - input.name.lower().replace(' ', '_'): ([ - execute_node(link.from_socket.node, context, cache)[link.from_socket.name] - for link in input.links - ] if len(input.links) > 1 else execute_node(input.links[0].from_socket.node, context, cache)[input.links[0].from_socket.name]) - if input.is_linked else getattr(input, 'default_value', None) - for input in node.inputs - } - if node.type == 'GROUP_OUTPUT': - return list(kwargs.values())[0] - result = node.execute(context, **kwargs) - on_execute(node, result) - cache[node] = result - return result +class NodeExecutionContext: + def __init__(self, depsgraph, update): + self.depsgraph = depsgraph + self.update = update + +def execute_node(node, context, cache): + result = None + match node.type: + case 'GROUP_OUTPUT': + return cache[node.inputs[0].links[0].from_socket.node] + case _: + if node in cache: + return cache[node] + kwargs = { + input.name.lower().replace(' ', '_'): ([ + cache[link.from_socket.node][link.from_socket.name] + for link in input.links + ] if len(input.links) > 1 else cache[input.links[0].from_socket.node][input.links[0].from_socket.name]) + if input.is_linked else getattr(input, 'default_value', None) + for input in node.inputs + } + if node.type == 'GROUP_OUTPUT': + return list(kwargs.values())[0] + result = node.execute(context, **kwargs) + return result -def execute(node_tree, context, on_execute=lambda _, __: None): +def execute(node_tree, depsgraph, node_begin=lambda node: None, node_update=lambda result: None, node_end=lambda node: None): output = next(n for n in node_tree.nodes if n.type == 'GROUP_OUTPUT') cache = {} - result = execute_node(output, context, cache, on_execute) - return result \ No newline at end of file + graph = { + node: [link.from_socket.node for input in node.inputs for link in input.links] + for node in node_tree.nodes + } + sort = graphlib.TopologicalSorter(graph) + for node in sort.static_order(): + if len(node.outputs) > 0 and next((l for i in node.outputs for l in i.links), None) is None: + continue # node outputs are unused + node_begin(node) + result = execute_node(node, NodeExecutionContext(depsgraph, node_update), cache) + cache[node] = result + node_end(node) + return next(iter(cache[output].values())) \ No newline at end of file diff --git a/engine/nodes/annotation_nodes.py b/engine/nodes/annotation_nodes.py index fcd03e0f..5437dbf2 100644 --- a/engine/nodes/annotation_nodes.py +++ b/engine/nodes/annotation_nodes.py @@ -17,8 +17,10 @@ def draw_buttons(self, context, layout): pass def execute(self, context, collection, invert): + depth_map = depth.render_depth_map(context.depsgraph, collection=collection, invert=invert) + context.update(depth_map) return { - 'Depth Map': depth.render_depth_map(context, collection=collection, invert=invert), + 'Depth Map': depth_map, } class NodeAnnotationOpenPose(DreamTexturesNode): @@ -34,6 +36,8 @@ def draw_buttons(self, context, layout): pass def execute(self, context, collection): + openpose_map = openpose.render_openpose_map(context.depsgraph, collection=collection) + context.update(openpose_map) return { - 'OpenPose Map': openpose.render_openpose_map(context, collection=collection) + 'OpenPose Map': openpose_map } \ No newline at end of file diff --git a/engine/nodes/pipeline_nodes.py b/engine/nodes/pipeline_nodes.py index c2e127f4..013c4f3a 100644 --- a/engine/nodes/pipeline_nodes.py +++ b/engine/nodes/pipeline_nodes.py @@ -9,6 +9,7 @@ from ...property_groups.dream_prompt import DreamPrompt, control_net_options from ..annotations import openpose from ..annotations import depth +import threading class NodeSocketControlNet(bpy.types.NodeSocket): bl_idname = "NodeSocketControlNet" @@ -102,21 +103,22 @@ def execute(self, context, prompt, negative_prompt, width, height, steps, seed, self.prompt.cfg_scale = cfg_scale args = self.prompt.generate_args() - shared_args = context.scene.dream_textures_engine_prompt.generate_args() + shared_args = context.depsgraph.scene.dream_textures_engine_prompt.generate_args() if controlnets is not None: if not isinstance(controlnets, list): controlnets = [controlnets] - result = Generator.shared().control_net( + future = Generator.shared().control_net( pipeline=args['pipeline'], model=args['model'], scheduler=args['scheduler'], optimizations=shared_args['optimizations'], seamless_axes=args['seamless_axes'], iterations=args['iterations'], + step_preview_mode=args['step_preview_mode'], control_net=[c.model for c in controlnets], - control=[c.control(context) for c in controlnets], + control=[c.control(context.depsgraph) for c in controlnets], controlnet_conditioning_scale=[c.conditioning_scale for c in controlnets], image=np.uint8(source_image * 255) if self.task == 'image_to_image' else None, @@ -129,19 +131,19 @@ def execute(self, context, prompt, negative_prompt, width, height, steps, seed, height=height, cfg_scale=cfg_scale, use_negative_prompt=True, - negative_prompt=negative_prompt, - step_preview_mode=StepPreviewMode.NONE - ).result() + negative_prompt=negative_prompt + ) else: match self.task: case 'prompt_to_image': - result = Generator.shared().prompt_to_image( + future = Generator.shared().prompt_to_image( pipeline=args['pipeline'], model=args['model'], scheduler=args['scheduler'], optimizations=shared_args['optimizations'], seamless_axes=args['seamless_axes'], iterations=args['iterations'], + step_preview_mode=args['step_preview_mode'], prompt=prompt, steps=steps, seed=seed, @@ -149,17 +151,17 @@ def execute(self, context, prompt, negative_prompt, width, height, steps, seed, height=height, cfg_scale=cfg_scale, use_negative_prompt=True, - negative_prompt=negative_prompt, - step_preview_mode=StepPreviewMode.NONE - ).result() + negative_prompt=negative_prompt + ) case 'image_to_image': - result = Generator.shared().image_to_image( + future = Generator.shared().image_to_image( pipeline=args['pipeline'], model=args['model'], scheduler=args['scheduler'], optimizations=shared_args['optimizations'], seamless_axes=args['seamless_axes'], iterations=args['iterations'], + step_preview_mode=args['step_preview_mode'], image=np.uint8(source_image * 255), strength=noise_strength, @@ -172,17 +174,17 @@ def execute(self, context, prompt, negative_prompt, width, height, steps, seed, height=height, cfg_scale=cfg_scale, use_negative_prompt=True, - negative_prompt=negative_prompt, - step_preview_mode=StepPreviewMode.NONE - ).result() + negative_prompt=negative_prompt + ) case 'depth_to_image': - result = Generator.shared().depth_to_image( + future = Generator.shared().depth_to_image( pipeline=args['pipeline'], model=args['model'], scheduler=args['scheduler'], optimizations=shared_args['optimizations'], seamless_axes=args['seamless_axes'], iterations=args['iterations'], + step_preview_mode=args['step_preview_mode'], depth=depth_map, image=np.uint8(source_image * 255) if source_image is not None else None, @@ -195,9 +197,21 @@ def execute(self, context, prompt, negative_prompt, width, height, steps, seed, height=height, cfg_scale=cfg_scale, use_negative_prompt=True, - negative_prompt=negative_prompt, - step_preview_mode=StepPreviewMode.NONE - ).result() + negative_prompt=negative_prompt + ) + event = threading.Event() + result = None + def on_response(_, response): + context.update(response.images[0]) + + def on_done(future): + nonlocal result + result = future.result() + event.set() + + future.add_response_callback(on_response) + future.add_done_callback(on_done) + event.wait() return { 'Image': result[-1].images[-1] } diff --git a/generator_process/actions/control_net.py b/generator_process/actions/control_net.py index 7cf15971..4ad3f5bf 100644 --- a/generator_process/actions/control_net.py +++ b/generator_process/actions/control_net.py @@ -286,7 +286,6 @@ def __call__( int(8 * (width // 8)), int(8 * (height // 8)), ) - print(control) control_image = [PIL.Image.fromarray(np.uint8(c * 255)).convert('RGB').resize(rounded_size) for c in control] if control is not None else None init_image = None if image is None else (PIL.Image.open(image) if isinstance(image, str) else PIL.Image.fromarray(image.astype(np.uint8))).convert('RGB').resize(rounded_size) From 76f883dac05a4d939a11514c656b93d4a66e8d77 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Fri, 17 Mar 2023 17:59:17 -0400 Subject: [PATCH 17/28] Support the Group Input node --- engine/engine.py | 70 +++++++++-------------------------------- engine/node_executor.py | 5 +++ 2 files changed, 19 insertions(+), 56 deletions(-) diff --git a/engine/engine.py b/engine/engine.py index 5e3d6dc9..9f26a8d5 100644 --- a/engine/engine.py +++ b/engine/engine.py @@ -64,62 +64,6 @@ def node_end(_): layer.rect = node_result.reshape(-1, node_result.shape[-1]) self.end_result(result) - - def view_update(self, context, depsgraph): - region = context.region - view3d = context.space_data - scene = depsgraph.scene - - # Get viewport dimensions - dimensions = region.width, region.height - - if not self.scene_data: - # First time initialization - self.scene_data = [] - first_time = True - - # Loop over all datablocks used in the scene. - for datablock in depsgraph.ids: - pass - else: - first_time = False - - # Test which datablocks changed - for update in depsgraph.updates: - print("Datablock updated: ", update.id.name) - - # Test if any material was added, removed or changed. - if depsgraph.id_type_updated('MATERIAL'): - print("Materials updated") - - # Loop over all object instances in the scene. - if first_time or depsgraph.id_type_updated('OBJECT'): - for instance in depsgraph.object_instances: - pass - - # For viewport renders, this method is called whenever Blender redraws - # the 3D viewport. The renderer is expected to quickly draw the render - # with OpenGL, and not perform other expensive work. - # Blender will draw overlays for selection and editing on top of the - # rendered image automatically. - def view_draw(self, context, depsgraph): - region = context.region - scene = depsgraph.scene - - # Get viewport dimensions - dimensions = region.width, region.height - - # Bind shader that converts from scene linear to display space, - gpu.state.blend_set('ALPHA_PREMULT') - self.bind_display_space_shader(scene) - - if not self.draw_data or self.draw_data.dimensions != dimensions: - self.draw_data = CustomDrawData(dimensions) - - self.draw_data.draw() - - self.unbind_display_space_shader() - gpu.state.blend_set('NONE') class NewEngineNodeTree(bpy.types.Operator): bl_idname = "dream_textures.new_engine_node_tree" @@ -178,6 +122,20 @@ def draw(self, context): col.prop(context.scene.render, "resolution_y", text="Y") yield FormatPanel + class NodeTreeInputsPanel(RenderPanel): + """Create a subpanel for format options""" + bl_idname = f"DREAM_PT_dream_panel_node_tree_inputs_engine" + bl_label = "Inputs" + + def draw(self, context): + super().draw(context) + layout = self.layout + layout.use_property_split = True + + for input in context.scene.dream_textures_render_engine.node_tree.inputs: + layout.prop(input, "default_value", text=input.name) + yield NodeTreeInputsPanel + # Bone properties class OpenPoseArmaturePanel(bpy.types.Panel): bl_idname = "DREAM_PT_dream_textures_armature_openpose" diff --git a/engine/node_executor.py b/engine/node_executor.py index 7bf8a52b..02158468 100644 --- a/engine/node_executor.py +++ b/engine/node_executor.py @@ -13,6 +13,11 @@ def __init__(self, depsgraph, update): def execute_node(node, context, cache): result = None match node.type: + case 'GROUP_INPUT': + return { + input.name: input.default_value + for input in context.depsgraph.scene.dream_textures_render_engine.node_tree.inputs + } case 'GROUP_OUTPUT': return cache[node.inputs[0].links[0].from_socket.node] case _: From 0c5b941cb9931ce4aea1dd88aa9c450f837b3f7f Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Fri, 17 Mar 2023 18:11:17 -0400 Subject: [PATCH 18/28] Use test_break to support cancellation --- engine/__init__.py | 3 +++ engine/engine.py | 5 +---- engine/node_executor.py | 12 ++++++------ engine/nodes/input_nodes.py | 17 +++++++++++++++++ engine/nodes/pipeline_nodes.py | 6 +++++- 5 files changed, 32 insertions(+), 11 deletions(-) diff --git a/engine/__init__.py b/engine/__init__.py index dbbdc72c..5bc83201 100644 --- a/engine/__init__.py +++ b/engine/__init__.py @@ -25,6 +25,7 @@ def poll(cls, context): nodeitems_utils.NodeItem(NodeString.bl_idname), nodeitems_utils.NodeItem(NodeImage.bl_idname), nodeitems_utils.NodeItem(NodeCollection.bl_idname), + nodeitems_utils.NodeItem(NodeRenderProperties.bl_idname), ]), DreamTexturesNodeCategory("DREAM_TEXTURES_UTILITY", "Utilities", items = [ nodeitems_utils.NodeItem(NodeMath.bl_idname), @@ -65,6 +66,7 @@ def register(): bpy.utils.register_class(NodeString) bpy.utils.register_class(NodeCollection) bpy.utils.register_class(NodeImage) + bpy.utils.register_class(NodeRenderProperties) bpy.utils.register_class(NodeAnnotationDepth) bpy.utils.register_class(NodeAnnotationOpenPose) @@ -93,6 +95,7 @@ def unregister(): bpy.utils.unregister_class(NodeString) bpy.utils.unregister_class(NodeCollection) bpy.utils.unregister_class(NodeImage) + bpy.utils.unregister_class(NodeRenderProperties) bpy.utils.unregister_class(NodeAnnotationDepth) bpy.utils.unregister_class(NodeAnnotationOpenPose) diff --git a/engine/engine.py b/engine/engine.py index 9f26a8d5..a3dbf25e 100644 --- a/engine/engine.py +++ b/engine/engine.py @@ -6,9 +6,6 @@ from ..ui.panels.dream_texture import optimization_panels from .node_tree import DreamTexturesNodeTree from ..engine import node_executor -from .annotations import openpose -import time -from threading import Event class DreamTexturesRenderEngine(bpy.types.RenderEngine): """A custom Dream Textures render engine, that uses Stable Diffusion and scene data to render images, instead of as a pass on top of Cycles.""" @@ -56,7 +53,7 @@ def node_end(_): nonlocal progress progress += 1 self.update_progress(progress / len(scene.dream_textures_render_engine.node_tree.nodes)) - node_result = node_executor.execute(scene.dream_textures_render_engine.node_tree, depsgraph, node_begin=node_begin, node_update=node_update, node_end=node_end) + node_result = node_executor.execute(scene.dream_textures_render_engine.node_tree, depsgraph, node_begin=node_begin, node_update=node_update, node_end=node_end, test_break=self.test_break) node_result = prepare_result(node_result) except Exception as error: self.report({'ERROR'}, str(error)) diff --git a/engine/node_executor.py b/engine/node_executor.py index 02158468..a7aaa4d2 100644 --- a/engine/node_executor.py +++ b/engine/node_executor.py @@ -1,14 +1,12 @@ -import bpy -import numpy as np -from threading import Event import graphlib # from dream_textures.engine import node_executor # node_executor.execute(bpy.data.node_groups["NodeTree"], bpy.context.evaluated_depsgraph_get()) class NodeExecutionContext: - def __init__(self, depsgraph, update): + def __init__(self, depsgraph, update, test_break): self.depsgraph = depsgraph self.update = update + self.test_break = test_break def execute_node(node, context, cache): result = None @@ -36,7 +34,7 @@ def execute_node(node, context, cache): result = node.execute(context, **kwargs) return result -def execute(node_tree, depsgraph, node_begin=lambda node: None, node_update=lambda result: None, node_end=lambda node: None): +def execute(node_tree, depsgraph, node_begin=lambda node: None, node_update=lambda result: None, node_end=lambda node: None, test_break=lambda: False): output = next(n for n in node_tree.nodes if n.type == 'GROUP_OUTPUT') cache = {} graph = { @@ -45,10 +43,12 @@ def execute(node_tree, depsgraph, node_begin=lambda node: None, node_update=lamb } sort = graphlib.TopologicalSorter(graph) for node in sort.static_order(): + if test_break(): + return None if len(node.outputs) > 0 and next((l for i in node.outputs for l in i.links), None) is None: continue # node outputs are unused node_begin(node) - result = execute_node(node, NodeExecutionContext(depsgraph, node_update), cache) + result = execute_node(node, NodeExecutionContext(depsgraph, node_update, test_break), cache) cache[node] = result node_end(node) return next(iter(cache[output].values())) \ No newline at end of file diff --git a/engine/nodes/input_nodes.py b/engine/nodes/input_nodes.py index 19282011..2bed9e74 100644 --- a/engine/nodes/input_nodes.py +++ b/engine/nodes/input_nodes.py @@ -72,4 +72,21 @@ def draw_buttons(self, context, layout): def execute(self, context): return { 'Image': np.array(self.value.pixels).reshape((*self.value.size, self.value.channels)) + } + +class NodeRenderProperties(DreamTexturesNode): + bl_idname = "dream_textures.node_render_properties" + bl_label = "Render Properties" + + def init(self, context): + self.outputs.new("NodeSocketInt", "Resolution X") + self.outputs.new("NodeSocketInt", "Resolution Y") + + def draw_buttons(self, context, layout): + pass + + def execute(self, context): + return { + 'Resolution X': context.depsgraph.scene.render.resolution_x, + 'Resolution Y': context.depsgraph.scene.render.resolution_y, } \ No newline at end of file diff --git a/engine/nodes/pipeline_nodes.py b/engine/nodes/pipeline_nodes.py index 013c4f3a..07f0d3f0 100644 --- a/engine/nodes/pipeline_nodes.py +++ b/engine/nodes/pipeline_nodes.py @@ -5,7 +5,6 @@ import enum from ..node import DreamTexturesNode from ...generator_process import Generator -from ...generator_process.actions.prompt_to_image import StepPreviewMode from ...property_groups.dream_prompt import DreamPrompt, control_net_options from ..annotations import openpose from ..annotations import depth @@ -203,6 +202,11 @@ def execute(self, context, prompt, negative_prompt, width, height, steps, seed, result = None def on_response(_, response): context.update(response.images[0]) + if context.test_break(): + nonlocal result + future.cancel() + result = [response] + event.set() def on_done(future): nonlocal result From cb31d6166f50562f2557bbbfdf8f30c996bc00ee Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Sat, 18 Mar 2023 21:47:07 -0400 Subject: [PATCH 19/28] Update for new huggingface_hub version and simplify implementation --- generator_process/actions/huggingface_hub.py | 521 ++----------------- generator_process/actor.py | 129 +---- generator_process/future.py | 122 +++++ 3 files changed, 162 insertions(+), 610 deletions(-) create mode 100644 generator_process/future.py diff --git a/generator_process/actions/huggingface_hub.py b/generator_process/actions/huggingface_hub.py index 804a334b..18854cfc 100644 --- a/generator_process/actions/huggingface_hub.py +++ b/generator_process/actions/huggingface_hub.py @@ -14,6 +14,7 @@ import requests import json import enum +from ..future import Future class ModelType(enum.IntEnum): """ @@ -158,23 +159,27 @@ def hf_snapshot_download( model: str, token: str, revision: str | None = None -) -> Generator[DownloadStatus, None, None]: - from filelock import FileLock - from huggingface_hub.constants import ( - DEFAULT_REVISION, - HUGGINGFACE_HEADER_X_REPO_COMMIT, - HUGGINGFACE_HUB_CACHE, - REPO_TYPES, - ) - from huggingface_hub.file_download import REGEX_COMMIT_HASH, repo_folder_name, hf_hub_url, _request_wrapper, hf_raise_for_status, logger, cached_download, build_hf_headers, get_hf_file_metadata, _cache_commit_hash_for_specific_revision, OfflineModeIsEnabled, _create_relative_symlink - from huggingface_hub.hf_api import HfApi - from huggingface_hub.utils import filter_repo_objects, validate_hf_hub_args, tqdm, logging, EntryNotFoundError, LocalEntryNotFoundError, RevisionNotFoundError - +): + from huggingface_hub._snapshot_download import snapshot_download + from huggingface_hub.utils.tqdm import tqdm + from huggingface_hub.utils._errors import RevisionNotFoundError + from diffusers import StableDiffusionPipeline from diffusers.utils import DIFFUSERS_CACHE, WEIGHTS_NAME, CONFIG_NAME, ONNX_WEIGHTS_NAME from diffusers.schedulers.scheduling_utils import SCHEDULER_CONFIG_NAME - from diffusers.utils.hub_utils import http_user_agent + future = Future() + yield future + + class future_tqdm(tqdm): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + future.add_response(DownloadStatus(self.desc, 0, self.total)) + + def update(self, n=1): + future.add_response(DownloadStatus(self.desc, self.last_print_n + n, self.total)) + return super().update(n=n) + try: config_dict = StableDiffusionPipeline.load_config( model, @@ -183,499 +188,35 @@ def hf_snapshot_download( force_download=False, use_auth_token=token ) - # make sure we only download sub-folders and `diffusers` filenames folder_names = [k for k in config_dict.keys() if not k.startswith("_")] allow_patterns = [os.path.join(k, "*") for k in folder_names] allow_patterns += [WEIGHTS_NAME, SCHEDULER_CONFIG_NAME, CONFIG_NAME, ONNX_WEIGHTS_NAME, StableDiffusionPipeline.config_name] - - # make sure we don't download flax, safetensors, or ckpt weights. - ignore_patterns = ["*.msgpack", "*.safetensors", "*.ckpt"] - - requested_pipeline_class = config_dict.get("_class_name", StableDiffusionPipeline.__name__) - user_agent = {"pipeline_class": requested_pipeline_class} - user_agent = http_user_agent(user_agent) except: allow_patterns = None - ignore_patterns = None - user_agent = None - - # download all allow_patterns - - # NOTE: Modified to yield the progress as an int from 0-100. - def http_get( - url: str, - temp_file: BinaryIO, - *, - proxies=None, - resume_size=0, - headers: Optional[Dict[str, str]] = None, - timeout=10.0, - max_retries=0, - ): - headers = copy.deepcopy(headers) - if resume_size > 0: - headers["Range"] = "bytes=%d-" % (resume_size,) - r = _request_wrapper( - method="GET", - url=url, - stream=True, - proxies=proxies, - headers=headers, - timeout=timeout, - max_retries=max_retries, - ) - hf_raise_for_status(r) - content_length = r.headers.get("Content-Length") - total = resume_size + int(content_length) if content_length is not None else None - progress = 0 - previous_value = 0 - for chunk in r.iter_content(chunk_size=1024): - if chunk: # filter out keep-alive new chunks - progress += len(chunk) - value = progress / total - if value - previous_value > 0.01: - previous_value = value - yield value - temp_file.write(chunk) - - def hf_hub_download( - repo_id: str, - filename: str, - *, - subfolder: Optional[str] = None, - repo_type: Optional[str] = None, - revision: Optional[str] = None, - library_name: Optional[str] = None, - library_version: Optional[str] = None, - cache_dir: Union[str, Path, None] = None, - user_agent: Union[Dict, str, None] = None, - force_download: Optional[bool] = False, - force_filename: Optional[str] = None, - proxies: Optional[Dict] = None, - etag_timeout: Optional[float] = 10, - resume_download: Optional[bool] = False, - use_auth_token: Union[bool, str, None] = None, - local_files_only: Optional[bool] = False, - legacy_cache_layout: Optional[bool] = False, - ): - if force_filename is not None: - warnings.warn( - "The `force_filename` parameter is deprecated as a new caching system, " - "which keeps the filenames as they are on the Hub, is now in place.", - FutureWarning, - ) - legacy_cache_layout = True - - if legacy_cache_layout: - url = hf_hub_url( - repo_id, - filename, - subfolder=subfolder, - repo_type=repo_type, - revision=revision, - ) - - return cached_download( - url, - library_name=library_name, - library_version=library_version, - cache_dir=cache_dir, - user_agent=user_agent, - force_download=force_download, - force_filename=force_filename, - proxies=proxies, - etag_timeout=etag_timeout, - resume_download=resume_download, - use_auth_token=use_auth_token, - local_files_only=local_files_only, - legacy_cache_layout=legacy_cache_layout, - ) - - if cache_dir is None: - cache_dir = HUGGINGFACE_HUB_CACHE - if revision is None: - revision = DEFAULT_REVISION - if isinstance(cache_dir, Path): - cache_dir = str(cache_dir) - - if subfolder == "": - subfolder = None - if subfolder is not None: - # This is used to create a URL, and not a local path, hence the forward slash. - filename = f"{subfolder}/{filename}" - - if repo_type is None: - repo_type = "model" - if repo_type not in REPO_TYPES: - raise ValueError( - f"Invalid repo type: {repo_type}. Accepted repo types are:" - f" {str(REPO_TYPES)}" - ) - - storage_folder = os.path.join( - cache_dir, repo_folder_name(repo_id=repo_id, repo_type=repo_type) - ) - os.makedirs(storage_folder, exist_ok=True) - - # cross platform transcription of filename, to be used as a local file path. - relative_filename = os.path.join(*filename.split("/")) - - # if user provides a commit_hash and they already have the file on disk, - # shortcut everything. - if REGEX_COMMIT_HASH.match(revision): - pointer_path = os.path.join( - storage_folder, "snapshots", revision, relative_filename - ) - if os.path.exists(pointer_path): - return pointer_path - - url = hf_hub_url(repo_id, filename, repo_type=repo_type, revision=revision) - - headers = build_hf_headers( - use_auth_token=use_auth_token, - library_name=library_name, - library_version=library_version, - user_agent=user_agent, - ) - - url_to_download = url - etag = None - commit_hash = None - if not local_files_only: - try: - try: - metadata = get_hf_file_metadata( - url=url, - use_auth_token=use_auth_token, - proxies=proxies, - timeout=etag_timeout, - ) - except EntryNotFoundError as http_error: - # Cache the non-existence of the file and raise - commit_hash = http_error.response.headers.get( - HUGGINGFACE_HEADER_X_REPO_COMMIT - ) - if commit_hash is not None and not legacy_cache_layout: - no_exist_file_path = ( - Path(storage_folder) - / ".no_exist" - / commit_hash - / relative_filename - ) - no_exist_file_path.parent.mkdir(parents=True, exist_ok=True) - no_exist_file_path.touch() - _cache_commit_hash_for_specific_revision( - storage_folder, revision, commit_hash - ) - raise - - # Commit hash must exist - commit_hash = metadata.commit_hash - if commit_hash is None: - raise OSError( - "Distant resource does not seem to be on huggingface.co (missing" - " commit header)." - ) - - # Etag must exist - etag = metadata.etag - # We favor a custom header indicating the etag of the linked resource, and - # we fallback to the regular etag header. - # If we don't have any of those, raise an error. - if etag is None: - raise OSError( - "Distant resource does not have an ETag, we won't be able to" - " reliably ensure reproducibility." - ) - - # In case of a redirect, save an extra redirect on the request.get call, - # and ensure we download the exact atomic version even if it changed - # between the HEAD and the GET (unlikely, but hey). - # Useful for lfs blobs that are stored on a CDN. - if metadata.location != url: - url_to_download = metadata.location - if ( - "lfs.huggingface.co" in url_to_download - or "lfs-staging.huggingface.co" in url_to_download - ): - # Remove authorization header when downloading a LFS blob - headers.pop("authorization", None) - except (requests.exceptions.SSLError, requests.exceptions.ProxyError): - # Actually raise for those subclasses of ConnectionError - raise - except ( - requests.exceptions.ConnectionError, - requests.exceptions.Timeout, - OfflineModeIsEnabled, - ): - # Otherwise, our Internet connection is down. - # etag is None - pass - - # etag is None == we don't have a connection or we passed local_files_only. - # try to get the last downloaded one from the specified revision. - # If the specified revision is a commit hash, look inside "snapshots". - # If the specified revision is a branch or tag, look inside "refs". - if etag is None: - # In those cases, we cannot force download. - if force_download: - raise ValueError( - "We have no connection or you passed local_files_only, so" - " force_download is not an accepted option." - ) - if REGEX_COMMIT_HASH.match(revision): - commit_hash = revision - else: - ref_path = os.path.join(storage_folder, "refs", revision) - with open(ref_path) as f: - commit_hash = f.read() - - pointer_path = os.path.join( - storage_folder, "snapshots", commit_hash, relative_filename - ) - if os.path.exists(pointer_path): - return pointer_path - - # If we couldn't find an appropriate file on disk, - # raise an error. - # If files cannot be found and local_files_only=True, - # the models might've been found if local_files_only=False - # Notify the user about that - if local_files_only: - raise LocalEntryNotFoundError( - "Cannot find the requested files in the disk cache and" - " outgoing traffic has been disabled. To enable hf.co look-ups" - " and downloads online, set 'local_files_only' to False." - ) - else: - raise LocalEntryNotFoundError( - "Connection error, and we cannot find the requested files in" - " the disk cache. Please try again or make sure your Internet" - " connection is on." - ) - - # From now on, etag and commit_hash are not None. - blob_path = os.path.join(storage_folder, "blobs", etag) - pointer_path = os.path.join( - storage_folder, "snapshots", commit_hash, relative_filename - ) - - os.makedirs(os.path.dirname(blob_path), exist_ok=True) - os.makedirs(os.path.dirname(pointer_path), exist_ok=True) - # if passed revision is not identical to commit_hash - # then revision has to be a branch name or tag name. - # In that case store a ref. - _cache_commit_hash_for_specific_revision(storage_folder, revision, commit_hash) - - if os.path.exists(pointer_path) and not force_download: - return pointer_path - - if os.path.exists(blob_path) and not force_download: - # we have the blob already, but not the pointer - logger.info("creating pointer to %s from %s", blob_path, pointer_path) - _create_relative_symlink(blob_path, pointer_path, new_blob=False) - return pointer_path - - # Prevent parallel downloads of the same file with a lock. - lock_path = blob_path + ".lock" - - # Some Windows versions do not allow for paths longer than 255 characters. - # In this case, we must specify it is an extended path by using the "\\?\" prefix. - if os.name == "nt" and len(os.path.abspath(lock_path)) > 255: - lock_path = "\\\\?\\" + os.path.abspath(lock_path) - - if os.name == "nt" and len(os.path.abspath(blob_path)) > 255: - blob_path = "\\\\?\\" + os.path.abspath(blob_path) - - with FileLock(lock_path): - # If the download just completed while the lock was activated. - if os.path.exists(pointer_path) and not force_download: - # Even if returning early like here, the lock will be released. - return pointer_path - - if resume_download: - incomplete_path = blob_path + ".incomplete" - - @contextmanager - def _resumable_file_manager() -> "io.BufferedWriter": - with open(incomplete_path, "ab") as f: - yield f - - temp_file_manager = _resumable_file_manager - if os.path.exists(incomplete_path): - resume_size = os.stat(incomplete_path).st_size - else: - resume_size = 0 - else: - temp_file_manager = partial( - tempfile.NamedTemporaryFile, mode="wb", dir=cache_dir, delete=False - ) - resume_size = 0 - - # Download to temporary file, then copy to cache dir once finished. - # Otherwise you get corrupt cache entries if the download gets interrupted. - with temp_file_manager() as temp_file: - logger.info("downloading %s to %s", url, temp_file.name) - - yield from http_get( - url_to_download, - temp_file, - proxies=proxies, - resume_size=resume_size, - headers=headers, - ) - - logger.info("storing %s in cache at %s", url, blob_path) - os.replace(temp_file.name, blob_path) - - logger.info("creating pointer to %s from %s", blob_path, pointer_path) - _create_relative_symlink(blob_path, pointer_path, new_blob=True) - - try: - os.remove(lock_path) - except OSError: - pass - - @validate_hf_hub_args - def snapshot_download( - repo_id: str, - *, - revision: Optional[str] = None, - repo_type: Optional[str] = None, - cache_dir: Union[str, Path, None] = None, - library_name: Optional[str] = None, - library_version: Optional[str] = None, - user_agent: Optional[Union[Dict, str]] = None, - proxies: Optional[Dict] = None, - etag_timeout: Optional[float] = 10, - resume_download: Optional[bool] = False, - use_auth_token: Optional[Union[bool, str]] = None, - local_files_only: Optional[bool] = False, - allow_regex: Optional[Union[List[str], str]] = None, - ignore_regex: Optional[Union[List[str], str]] = None, - allow_patterns: Optional[Union[List[str], str]] = None, - ignore_patterns: Optional[Union[List[str], str]] = None, - ): - if cache_dir is None: - cache_dir = HUGGINGFACE_HUB_CACHE - if revision is None: - revision = DEFAULT_REVISION - if isinstance(cache_dir, Path): - cache_dir = str(cache_dir) - - if repo_type is None: - repo_type = "model" - if repo_type not in REPO_TYPES: - raise ValueError( - f"Invalid repo type: {repo_type}. Accepted repo types are:" - f" {str(REPO_TYPES)}" - ) - - storage_folder = os.path.join( - cache_dir, repo_folder_name(repo_id=repo_id, repo_type=repo_type) - ) - - # TODO: remove these 4 lines in version 0.12 - # Deprecated code to ensure backward compatibility. - if allow_regex is not None: - allow_patterns = allow_regex - if ignore_regex is not None: - ignore_patterns = ignore_regex - - # if we have no internet connection we will look for an - # appropriate folder in the cache - # If the specified revision is a commit hash, look inside "snapshots". - # If the specified revision is a branch or tag, look inside "refs". - if local_files_only: - if REGEX_COMMIT_HASH.match(revision): - commit_hash = revision - else: - # retrieve commit_hash from file - ref_path = os.path.join(storage_folder, "refs", revision) - with open(ref_path) as f: - commit_hash = f.read() - - snapshot_folder = os.path.join(storage_folder, "snapshots", commit_hash) - - if os.path.exists(snapshot_folder): - return snapshot_folder - - raise ValueError( - "Cannot find an appropriate cached snapshot folder for the specified" - " revision on the local disk and outgoing traffic has been disabled. To" - " enable repo look-ups and downloads online, set 'local_files_only' to" - " False." - ) - - # if we have internet connection we retrieve the correct folder name from the huggingface api - _api = HfApi() - repo_info = _api.repo_info( - repo_id=repo_id, - repo_type=repo_type, - revision=revision, - use_auth_token=use_auth_token, - ) - filtered_repo_files = list( - filter_repo_objects( - items=[f.rfilename for f in repo_info.siblings], - allow_patterns=allow_patterns, - ignore_patterns=ignore_patterns, - ) - ) - commit_hash = repo_info.sha - snapshot_folder = os.path.join(storage_folder, "snapshots", commit_hash) - # if passed revision is not identical to commit_hash - # then revision has to be a branch name or tag name. - # In that case store a ref. - if revision != commit_hash: - ref_path = os.path.join(storage_folder, "refs", revision) - os.makedirs(os.path.dirname(ref_path), exist_ok=True) - with open(ref_path, "w") as f: - f.write(commit_hash) - - # we pass the commit_hash to hf_hub_download - # so no network call happens if we already - # have the file locally. - - for i, repo_file in tqdm( - enumerate(filtered_repo_files), f"Fetching {len(filtered_repo_files)} files" - ): - yield DownloadStatus(repo_file, i, len(filtered_repo_files)) - for status in hf_hub_download( - repo_id, - filename=repo_file, - repo_type=repo_type, - revision=commit_hash, - cache_dir=cache_dir, - library_name=library_name, - library_version=library_version, - user_agent=user_agent, - proxies=proxies, - etag_timeout=etag_timeout, - resume_download=resume_download, - use_auth_token=use_auth_token, - ): - yield DownloadStatus(repo_file, status, 1) - yield DownloadStatus(repo_file, i + 1, len(filtered_repo_files)) + + # make sure we don't download flax, safetensors, or ckpt weights. + ignore_patterns = ["*.msgpack", "*.safetensors", "*.ckpt"] try: - yield from snapshot_download( + snapshot_download( model, - revision=revision, cache_dir=DIFFUSERS_CACHE, + token=token, + revision=revision, resume_download=True, - use_auth_token=token, allow_patterns=allow_patterns, ignore_patterns=ignore_patterns, - user_agent=user_agent, + tqdm_class=future_tqdm ) except RevisionNotFoundError: - yield from snapshot_download( + snapshot_download( model, cache_dir=DIFFUSERS_CACHE, + token=token, resume_download=True, - use_auth_token=token, allow_patterns=allow_patterns, ignore_patterns=ignore_patterns, - user_agent=user_agent, - ) \ No newline at end of file + tqdm_class=future_tqdm + ) + + future.set_done() \ No newline at end of file diff --git a/generator_process/actor.py b/generator_process/actor.py index 2c9d0923..3fabfc8f 100644 --- a/generator_process/actor.py +++ b/generator_process/actor.py @@ -1,12 +1,13 @@ -from multiprocessing import Queue, Process, Lock, current_process, get_context +from multiprocessing import Queue, Lock, current_process, get_context import multiprocessing.synchronize import enum import traceback import threading -from typing import Type, TypeVar, Callable, Any, MutableSet, Generator +from typing import Type, TypeVar, Generator import site import sys from ..absolute_path import absolute_path +from .future import Future def _load_dependencies(): site.addsitedir(absolute_path(".python_dependencies")) @@ -15,123 +16,6 @@ def _load_dependencies(): if current_process().name == "__actor__": _load_dependencies() -class Future: - """ - Object that represents a value that has not completed processing, but will in the future. - - Add callbacks to be notified when values become available, or use `.result()` and `.exception()` to wait for the value. - """ - _response_callbacks: MutableSet[Callable[['Future', Any], None]] = set() - _exception_callbacks: MutableSet[Callable[['Future', BaseException], None]] = set() - _done_callbacks: MutableSet[Callable[['Future'], None]] = set() - _responses: list = [] - _exception: BaseException | None = None - _done_event: threading.Event - done: bool = False - cancelled: bool = False - call_done_on_exception: bool = True - - def __init__(self): - self._response_callbacks = set() - self._exception_callbacks = set() - self._done_callbacks = set() - self._responses = [] - self._exception = None - self._done_event = threading.Event() - self.done = False - self.cancelled = False - self.call_done_on_exception = True - - def result(self, last_only=False): - """ - Get the result value (blocking). - """ - def _response(): - match len(self._responses): - case 0: - return None - case 1: - return self._responses[0] - case _: - return self._responses[-1] if last_only else self._responses - if self._exception is not None: - raise self._exception - if self.done: - return _response() - else: - self._done_event.wait() - if self._exception is not None: - raise self._exception - return _response() - - def exception(self): - if self.done: - return self._exception - else: - self._done_event.wait() - return self._exception - - def cancel(self): - self.cancelled = True - - def _run_on_main_thread(self, func): - import bpy - bpy.app.timers.register(func) - - def add_response(self, response): - """ - Add a response value and notify all consumers. - """ - self._responses.append(response) - def run_callbacks(): - for response_callback in self._response_callbacks: - response_callback(self, response) - self._run_on_main_thread(run_callbacks) - - def set_exception(self, exception: BaseException): - """ - Set the exception. - """ - self._exception = exception - def run_callbacks(): - for exception_callback in self._exception_callbacks: - exception_callback(self, exception) - self._run_on_main_thread(run_callbacks) - - def set_done(self): - """ - Mark the future as done. - """ - assert not self.done - self.done = True - self._done_event.set() - if self._exception is None or self.call_done_on_exception: - def run_callbacks(): - for done_callback in self._done_callbacks: - done_callback(self) - self._run_on_main_thread(run_callbacks) - - def add_response_callback(self, callback: Callable[['Future', Any], None]): - """ - Add a callback to run whenever a response is received. - Will be called multiple times by generator functions. - """ - self._response_callbacks.add(callback) - - def add_exception_callback(self, callback: Callable[['Future', BaseException], None]): - """ - Add a callback to run when the future errors. - Will only be called once at the first exception. - """ - self._exception_callbacks.add(callback) - - def add_done_callback(self, callback: Callable[['Future'], None]): - """ - Add a callback to run when the future is marked as done. - Will only be called once. - """ - self._done_callbacks.add(callback) - class ActorContext(enum.IntEnum): """ The context of an `Actor` object. @@ -278,7 +162,12 @@ def _receive(self, message: Message): pass if extra_message == Message.CANCEL: break - self._response_queue.put(res) + if isinstance(res, Future): + res.add_response_callback(lambda _, res: self._response_queue.put(res)) + res.add_exception_callback(lambda _, e: self._response_queue.put(RuntimeError(repr(e)))) + res.add_done_callback(lambda _: None) + else: + self._response_queue.put(res) else: self._response_queue.put(response) except Exception as e: diff --git a/generator_process/future.py b/generator_process/future.py new file mode 100644 index 00000000..c9f55f7c --- /dev/null +++ b/generator_process/future.py @@ -0,0 +1,122 @@ +import threading +from typing import Callable, Any, MutableSet + +class Future: + """ + Object that represents a value that has not completed processing, but will in the future. + + Add callbacks to be notified when values become available, or use `.result()` and `.exception()` to wait for the value. + """ + _response_callbacks: MutableSet[Callable[['Future', Any], None]] = set() + _exception_callbacks: MutableSet[Callable[['Future', BaseException], None]] = set() + _done_callbacks: MutableSet[Callable[['Future'], None]] = set() + _responses: list = [] + _exception: BaseException | None = None + _done_event: threading.Event + done: bool = False + cancelled: bool = False + call_done_on_exception: bool = True + + def __init__(self): + self._response_callbacks = set() + self._exception_callbacks = set() + self._done_callbacks = set() + self._responses = [] + self._exception = None + self._done_event = threading.Event() + self.done = False + self.cancelled = False + self.call_done_on_exception = True + + def result(self, last_only=False): + """ + Get the result value (blocking). + """ + def _response(): + match len(self._responses): + case 0: + return None + case 1: + return self._responses[0] + case _: + return self._responses[-1] if last_only else self._responses + if self._exception is not None: + raise self._exception + if self.done: + return _response() + else: + self._done_event.wait() + if self._exception is not None: + raise self._exception + return _response() + + def exception(self): + if self.done: + return self._exception + else: + self._done_event.wait() + return self._exception + + def cancel(self): + self.cancelled = True + + def _run_on_main_thread(self, func): + try: + import bpy + bpy.app.timers.register(func) + except ModuleNotFoundError: + func() + + def add_response(self, response): + """ + Add a response value and notify all consumers. + """ + self._responses.append(response) + def run_callbacks(): + for response_callback in self._response_callbacks: + response_callback(self, response) + self._run_on_main_thread(run_callbacks) + + def set_exception(self, exception: BaseException): + """ + Set the exception. + """ + self._exception = exception + def run_callbacks(): + for exception_callback in self._exception_callbacks: + exception_callback(self, exception) + self._run_on_main_thread(run_callbacks) + + def set_done(self): + """ + Mark the future as done. + """ + assert not self.done + self.done = True + self._done_event.set() + if self._exception is None or self.call_done_on_exception: + def run_callbacks(): + for done_callback in self._done_callbacks: + done_callback(self) + self._run_on_main_thread(run_callbacks) + + def add_response_callback(self, callback: Callable[['Future', Any], None]): + """ + Add a callback to run whenever a response is received. + Will be called multiple times by generator functions. + """ + self._response_callbacks.add(callback) + + def add_exception_callback(self, callback: Callable[['Future', BaseException], None]): + """ + Add a callback to run when the future errors. + Will only be called once at the first exception. + """ + self._exception_callbacks.add(callback) + + def add_done_callback(self, callback: Callable[['Future'], None]): + """ + Add a callback to run when the future is marked as done. + Will only be called once. + """ + self._done_callbacks.add(callback) \ No newline at end of file From 507c2956d3f88439f5aec01d3cbc4d8f0e4dd4a6 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Sat, 18 Mar 2023 21:48:46 -0400 Subject: [PATCH 20/28] Add ControlNet node to category --- engine/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/engine/__init__.py b/engine/__init__.py index 5bc83201..6e42b94b 100644 --- a/engine/__init__.py +++ b/engine/__init__.py @@ -19,6 +19,7 @@ def poll(cls, context): categories = [ DreamTexturesNodeCategory("DREAM_TEXTURES_PIPELINE", "Pipeline", items = [ nodeitems_utils.NodeItem(NodeStableDiffusion.bl_idname), + nodeitems_utils.NodeItem(NodeControlNet.bl_idname), ]), DreamTexturesNodeCategory("DREAM_TEXTURES_INPUT", "Input", items = [ nodeitems_utils.NodeItem(NodeInteger.bl_idname), From b6c82c788306f85e24d97e4431a10f7f31eee3e6 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Sat, 18 Mar 2023 22:05:44 -0400 Subject: [PATCH 21/28] Fix UI error --- engine/engine.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/engine/engine.py b/engine/engine.py index a3dbf25e..2dd209ed 100644 --- a/engine/engine.py +++ b/engine/engine.py @@ -129,8 +129,9 @@ def draw(self, context): layout = self.layout layout.use_property_split = True - for input in context.scene.dream_textures_render_engine.node_tree.inputs: - layout.prop(input, "default_value", text=input.name) + if context.scene.dream_textures_render_engine.node_tree is not None: + for input in context.scene.dream_textures_render_engine.node_tree.inputs: + layout.prop(input, "default_value", text=input.name) yield NodeTreeInputsPanel # Bone properties From 7d6375985c12e1b15f3432c527d09da38a0dc203 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Sat, 25 Mar 2023 16:33:56 -0400 Subject: [PATCH 22/28] Fix ControlNet + img2img --- generator_process/actions/control_net.py | 114 ++++++++++++++++++++--- operators/dream_texture.py | 8 ++ property_groups/dream_prompt.py | 2 + 3 files changed, 111 insertions(+), 13 deletions(-) diff --git a/generator_process/actions/control_net.py b/generator_process/actions/control_net.py index 4ad3f5bf..24fb92c3 100644 --- a/generator_process/actions/control_net.py +++ b/generator_process/actions/control_net.py @@ -44,16 +44,85 @@ def control_net( case Pipeline.STABLE_DIFFUSION: import diffusers from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import MultiControlNetModel, ControlNetModel + from diffusers.utils import deprecate, randn_tensor import torch import PIL.Image import PIL.ImageOps class GeneratorPipeline(diffusers.StableDiffusionControlNetPipeline): + # copied from diffusers.StableDiffusionImg2ImgPipeline + def get_timesteps(self, num_inference_steps, strength, device): + # get the original timestep using init_timestep + init_timestep = min(int(num_inference_steps * strength), num_inference_steps) + + t_start = max(num_inference_steps - init_timestep, 0) + timesteps = self.scheduler.timesteps[t_start:] + + return timesteps, num_inference_steps - t_start + + # copied from diffusers.StableDiffusionImg2ImgPipeline + def prepare_img2img_latents(self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None): + if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)): + raise ValueError( + f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}" + ) + + image = image.to(device=device, dtype=dtype) + + batch_size = batch_size * num_images_per_prompt + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + if isinstance(generator, list): + init_latents = [ + self.vae.encode(image[i : i + 1]).latent_dist.sample(generator[i]) for i in range(batch_size) + ] + init_latents = torch.cat(init_latents, dim=0) + else: + init_latents = self.vae.encode(image).latent_dist.sample(generator) + + init_latents = self.vae.config.scaling_factor * init_latents + + if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0: + # expand init_latents for batch_size + deprecation_message = ( + f"You have passed {batch_size} text prompts (`prompt`), but only {init_latents.shape[0]} initial" + " images (`image`). Initial images are now duplicating to match the number of text prompts. Note" + " that this behavior is deprecated and will be removed in a version 1.0.0. Please make sure to update" + " your script to pass as many initial images as text prompts to suppress this warning." + ) + deprecate("len(prompt) != len(image)", "1.0.0", deprecation_message, standard_warn=False) + additional_image_per_prompt = batch_size // init_latents.shape[0] + init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0) + elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0: + raise ValueError( + f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts." + ) + else: + init_latents = torch.cat([init_latents], dim=0) + + shape = init_latents.shape + noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + + # get latents + init_latents = self.scheduler.add_noise(init_latents, noise, timestep) + latents = init_latents + + return latents + @torch.no_grad() def __call__( self, prompt: Union[str, List[str]] = None, image: Union[torch.FloatTensor, PIL.Image.Image, List[torch.FloatTensor], List[PIL.Image.Image]] = None, + + # NOTE: Modified to support initial image. + init_image: Union[torch.FloatTensor, PIL.Image.Image, List[torch.FloatTensor], List[PIL.Image.Image]] = None, + strength: float = 1.0, + height: Optional[int] = None, width: Optional[int] = None, num_inference_steps: int = 50, @@ -152,21 +221,40 @@ def __call__( assert False # 5. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps + # NOTE: Modified to support initial image + if init_image is not None: + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) + latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) + else: + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = self.scheduler.timesteps # 6. Prepare latent variables num_channels_latents = self.unet.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - device, - generator, - latents, - ) + # NOTE: Modified to support initial image + if init_image is not None: + init_image = diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.preprocess(init_image) + latents = self.prepare_img2img_latents( + init_image, + latent_timestep, + batch_size, + num_images_per_prompt, + prompt_embeds.dtype, + device, + generator + ) + else: + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + prompt_embeds.dtype, + device, + generator, + latents, + ) # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) @@ -309,7 +397,7 @@ def __call__( prompt=prompt, image=control_image, controlnet_conditioning_scale=controlnet_conditioning_scale, - # image=init_image, + init_image=init_image, strength=strength, width=rounded_size[0], height=rounded_size[1], diff --git a/operators/dream_texture.py b/operators/dream_texture.py index c8b9e4c2..44235086 100644 --- a/operators/dream_texture.py +++ b/operators/dream_texture.py @@ -216,6 +216,14 @@ def require_depth(): control=[init_image], **generated_args ) + case 'control_net_color': + f = gen.control_net( + image=init_image, + control=[np.flipud(np.array(scene.init_depth.pixels) + .astype(np.float32) + .reshape((scene.init_depth.size[1], scene.init_depth.size[0], scene.init_depth.channels)))], + **generated_args + ) case 'inpaint': f = gen.inpaint( image=init_image, diff --git a/property_groups/dream_prompt.py b/property_groups/dream_prompt.py index c64001d7..4e3bd8a7 100644 --- a/property_groups/dream_prompt.py +++ b/property_groups/dream_prompt.py @@ -59,9 +59,11 @@ def inpaint_mask_sources_filtered(self, context): def modify_action_source_type(self, context): return [ ('color', 'Color', 'Use the color information from the image', 1), + None, ('depth_generated', 'Color and Generated Depth', 'Use MiDaS to infer the depth of the initial image and include it in the conditioning. Can give results that more closely match the composition of the source image', 2), ('depth_map', 'Color and Depth Map', 'Specify a secondary image to use as the depth map. Can give results that closely match the composition of the depth map', 3), ('depth', 'Depth', 'Treat the initial image as a depth map, and ignore any color. Matches the composition of the source image without any color influence', 4), + None, ('control_net', 'ControlNet', 'Treat the initial image as the input to a ControlNet model', 5), ('control_net_color', 'Color and ControlNet', 'Specify a secondary image to use with a ControlNet model', 6), ] From 7498c6aaa713d3677c0f534b39877bfd1b7a49cd Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Sat, 25 Mar 2023 17:46:42 -0400 Subject: [PATCH 23/28] Support multi-ControlNet without nodes --- classes.py | 6 ++++ engine/nodes/pipeline_nodes.py | 3 +- generator_process/actions/control_net.py | 4 ++- operators/dream_texture.py | 30 +++++++++---------- operators/view_history.py | 10 ++++++- property_groups/control_net.py | 38 ++++++++++++++++++++++++ property_groups/dream_prompt.py | 31 ++++++++++--------- ui/panels/dream_texture.py | 24 ++++++++++++--- 8 files changed, 109 insertions(+), 37 deletions(-) create mode 100644 property_groups/control_net.py diff --git a/classes.py b/classes.py index 83ea7af3..5f6adefd 100644 --- a/classes.py +++ b/classes.py @@ -6,6 +6,7 @@ from .operators.upscale import Upscale from .operators.project import ProjectDreamTexture, dream_texture_projection_panels from .operators.notify_result import NotifyResult +from .property_groups.control_net import ControlNet, SCENE_UL_ControlNetList, ControlNetsAdd, ControlNetsRemove from .property_groups.dream_prompt import DreamPrompt from .property_groups.seamless_result import SeamlessResult from .ui.panels import dream_texture, history, upscaling, render_properties @@ -31,6 +32,10 @@ InpaintAreaBrushActivated, Upscale, ProjectDreamTexture, + + SCENE_UL_ControlNetList, + ControlNetsAdd, + ControlNetsRemove, DREAM_PT_AdvancedPresets, DREAM_MT_AdvancedPresets, @@ -55,6 +60,7 @@ ModelSearch, InstallModel, Model, + ControlNet, DreamPrompt, SeamlessResult, UninstallDependencies, diff --git a/engine/nodes/pipeline_nodes.py b/engine/nodes/pipeline_nodes.py index 07f0d3f0..a9a5be7c 100644 --- a/engine/nodes/pipeline_nodes.py +++ b/engine/nodes/pipeline_nodes.py @@ -5,7 +5,8 @@ import enum from ..node import DreamTexturesNode from ...generator_process import Generator -from ...property_groups.dream_prompt import DreamPrompt, control_net_options +from ...property_groups.control_net import control_net_options +from ...property_groups.dream_prompt import DreamPrompt from ..annotations import openpose from ..annotations import depth import threading diff --git a/generator_process/actions/control_net.py b/generator_process/actions/control_net.py index 24fb92c3..8feb196d 100644 --- a/generator_process/actions/control_net.py +++ b/generator_process/actions/control_net.py @@ -21,7 +21,9 @@ def control_net( control_net: list[str], control: list[NDArray] | None, controlnet_conditioning_scale: list[float], - image: NDArray | str | None, + + image: NDArray | str | None, # image to image + strength: float, prompt: str | list[str], steps: int, diff --git a/operators/dream_texture.py b/operators/dream_texture.py index 44235086..2da701d6 100644 --- a/operators/dream_texture.py +++ b/operators/dream_texture.py @@ -135,7 +135,14 @@ def done_callback(future): scene.dream_textures_prompt.hash = image_hash history_entry = context.scene.dream_textures_history.add() for key, value in history_template.items(): - setattr(history_entry, key, value) + match key: + case 'control_nets': + for net in value: + n = history_entry.control_nets.add() + for prop in n.__annotations__.keys(): + setattr(n, prop, getattr(net, prop)) + case _: + setattr(history_entry, key, value) history_entry.seed = str(seed) history_entry.hash = image_hash if is_file_batch: @@ -170,7 +177,12 @@ def generate_next(): else: generated_args["prompt"] = [original_prompt] * batch_size generated_args["negative_prompt"] = [original_negative_prompt] * batch_size - if init_image is not None: + if len(generated_args['control_net']) > 0: + f = gen.control_net( + image=init_image, + **generated_args + ) + elif init_image is not None: match generated_args['init_img_action']: case 'modify': models = list(filter( @@ -210,20 +222,6 @@ def require_depth(): depth=np.flipud(init_image.astype(np.float32) / 255.), **generated_args, ) - case 'control_net': - f = gen.control_net( - image=None, - control=[init_image], - **generated_args - ) - case 'control_net_color': - f = gen.control_net( - image=init_image, - control=[np.flipud(np.array(scene.init_depth.pixels) - .astype(np.float32) - .reshape((scene.init_depth.size[1], scene.init_depth.size[0], scene.init_depth.channels)))], - **generated_args - ) case 'inpaint': f = gen.inpaint( image=init_image, diff --git a/operators/view_history.py b/operators/view_history.py index f86a3bff..45440c03 100644 --- a/operators/view_history.py +++ b/operators/view_history.py @@ -31,7 +31,15 @@ def execute(self, context): selection = context.scene.dream_textures_history[context.scene.dream_textures_history_selection] for prop in selection.__annotations__.keys(): if hasattr(context.scene.dream_textures_prompt, prop): - setattr(context.scene.dream_textures_prompt, prop, getattr(selection, prop)) + match prop: + case 'control_nets': + context.scene.dream_textures_prompt.control_nets.clear() + for net in selection.control_nets: + n = context.scene.dream_textures_prompt.control_nets.add() + for k in n.__annotations__.keys(): + setattr(n, k, getattr(net, k)) + case _: + setattr(context.scene.dream_textures_prompt, prop, getattr(selection, prop)) # when the seed of the promt is found in the available image datablocks, use that one in the open image editor # note: when there is more than one image with the seed in it's name, do nothing. Same when no image with that seed is available. if prop == 'hash': diff --git a/property_groups/control_net.py b/property_groups/control_net.py new file mode 100644 index 00000000..db47b2dd --- /dev/null +++ b/property_groups/control_net.py @@ -0,0 +1,38 @@ +import bpy +from bpy.props import FloatProperty, EnumProperty, PointerProperty + +from ..generator_process.actions.huggingface_hub import ModelType +from ..preferences import StableDiffusionPreferences + +def control_net_options(self, context): + return [ + (model.model_base, model.model_base.replace('models--', '').replace('--', '/'), '') for model in context.preferences.addons[StableDiffusionPreferences.bl_idname].preferences.installed_models + if model.model_type == ModelType.CONTROL_NET.name + ] + +class ControlNet(bpy.types.PropertyGroup): + control_net: EnumProperty(name="ControlNet", items=control_net_options, description="Specify which ControlNet to use") + conditioning_scale: FloatProperty(name="ControlNet Conditioning Scale", default=1.0, description="Increases the strength of the ControlNet's effect") + control_image: PointerProperty(type=bpy.types.Image) + +class SCENE_UL_ControlNetList(bpy.types.UIList): + def draw_item(self, context, layout, data, item, icon, active_data, active_propname): + layout.separator() + layout.prop(item, "control_net", text="") + layout.prop(item, "conditioning_scale", text="") + layout.template_ID(item, "control_image", open="image.open") + +class ControlNetsAdd(bpy.types.Operator): + bl_idname = "dream_textures.control_nets_add" + bl_label = "Add ControlNet" + + def execute(self, context): + context.scene.dream_textures_prompt.control_nets.add() + return {'FINISHED'} +class ControlNetsRemove(bpy.types.Operator): + bl_idname = "dream_textures.control_nets_remove" + bl_label = "Add ControlNet" + + def execute(self, context): + context.scene.dream_textures_prompt.control_nets.remove(context.scene.dream_textures_prompt.active_control_net) + return {'FINISHED'} \ No newline at end of file diff --git a/property_groups/dream_prompt.py b/property_groups/dream_prompt.py index 4e3bd8a7..fe8893fd 100644 --- a/property_groups/dream_prompt.py +++ b/property_groups/dream_prompt.py @@ -1,5 +1,5 @@ import bpy -from bpy.props import FloatProperty, IntProperty, EnumProperty, BoolProperty, StringProperty, IntVectorProperty +from bpy.props import FloatProperty, IntProperty, EnumProperty, BoolProperty, StringProperty, IntVectorProperty, CollectionProperty import os import sys from typing import _AnnotatedAlias @@ -10,6 +10,9 @@ from ..prompt_engineering import * from ..preferences import StableDiffusionPreferences from .dream_prompt_validation import validate +from .control_net import ControlNet + +import numpy as np from functools import reduce @@ -63,9 +66,6 @@ def modify_action_source_type(self, context): ('depth_generated', 'Color and Generated Depth', 'Use MiDaS to infer the depth of the initial image and include it in the conditioning. Can give results that more closely match the composition of the source image', 2), ('depth_map', 'Color and Depth Map', 'Specify a secondary image to use as the depth map. Can give results that closely match the composition of the depth map', 3), ('depth', 'Depth', 'Treat the initial image as a depth map, and ignore any color. Matches the composition of the source image without any color influence', 4), - None, - ('control_net', 'ControlNet', 'Treat the initial image as the input to a ControlNet model', 5), - ('control_net_color', 'Color and ControlNet', 'Specify a secondary image to use with a ControlNet model', 6), ] def model_options(self, context): @@ -113,12 +113,6 @@ def pipeline_options(self, context): (Pipeline.STABILITY_SDK.name, 'DreamStudio', 'Cloud compute via DreamStudio', 2), ] -def control_net_options(self, context): - return [ - (model.model_base, model.model_base.replace('models--', '').replace('--', '/'), '') for model in context.preferences.addons[StableDiffusionPreferences.bl_idname].preferences.installed_models - if model.model_type == ModelType.CONTROL_NET.name - ] - def seed_clamp(self, ctx): # clamp seed right after input to make it clear what the limits are try: @@ -132,8 +126,8 @@ def seed_clamp(self, ctx): "pipeline": EnumProperty(name="Pipeline", items=pipeline_options, default=1 if Pipeline.local_available() else 2, description="Specify which model and target should be used."), "model": EnumProperty(name="Model", items=model_options, description="Specify which model to use for inference"), - "control_net": EnumProperty(name="ControlNet", items=control_net_options, description="Specify which ControlNet to use"), - "controlnet_conditioning_scale": FloatProperty(name="ControlNet Conditioning Scale", default=1.0, description="Increases the strength of the ControlNet's effect"), + "control_nets": CollectionProperty(type=ControlNet), + "active_control_net": IntProperty(name="Active ControlNet"), # Prompt "prompt_structure": EnumProperty(name="Preset", items=prompt_structures_items, description="Fill in a few simple options to create interesting images quickly"), @@ -281,8 +275,17 @@ def generate_args(self): args['width'] = args['width'] if args['use_size'] else None args['height'] = args['height'] if args['use_size'] else None - args['control_net'] = [args['control_net']] - args['controlnet_conditioning_scale'] = [args['controlnet_conditioning_scale']] + args['control_net'] = [net.control_net for net in args['control_nets']] + args['controlnet_conditioning_scale'] = [net.conditioning_scale for net in args['control_nets']] + args['control'] = [ + np.flipud( + (np.array(net.control_image.pixels) * 255) + .astype(np.uint8) + .reshape((net.control_image.size[1], net.control_image.size[0], net.control_image.channels)) + ) + for net in args['control_nets'] + ] + del args['control_nets'] return args DreamPrompt.generate_prompt = generate_prompt diff --git a/ui/panels/dream_texture.py b/ui/panels/dream_texture.py index 6e543ae6..1a7beb75 100644 --- a/ui/panels/dream_texture.py +++ b/ui/panels/dream_texture.py @@ -78,6 +78,7 @@ def get_seamless_result(context, prompt): get_seamless_result=get_seamless_result) yield create_panel(space_type, 'UI', DreamTexturePanel.bl_idname, size_panel, get_prompt) yield from create_panel(space_type, 'UI', DreamTexturePanel.bl_idname, init_image_panels, get_prompt) + yield create_panel(space_type, 'UI', DreamTexturePanel.bl_idname, control_net_panel, get_prompt) yield from create_panel(space_type, 'UI', DreamTexturePanel.bl_idname, advanced_panel, get_prompt) yield create_panel(space_type, 'UI', DreamTexturePanel.bl_idname, actions_panel, get_prompt) @@ -232,13 +233,28 @@ def _outpaint_warning_box(warning): layout.prop(prompt, "use_init_img_color") if prompt.init_img_action == 'modify': layout.prop(prompt, "modify_action_source_type") - if prompt.modify_action_source_type in {'control_net', 'control_net_color'}: - layout.prop(context.scene.dream_textures_prompt, 'control_net') - layout.prop(context.scene.dream_textures_prompt, 'controlnet_conditioning_scale') - if prompt.modify_action_source_type == 'depth_map' or prompt.modify_action_source_type == 'control_net_color': + if prompt.modify_action_source_type == 'depth_map': layout.template_ID(context.scene, "init_depth", open="image.open") yield InitImagePanel +def control_net_panel(sub_panel, space_type, get_prompt): + class ControlNetPanel(sub_panel): + """Create a subpanel for ControlNet options""" + bl_idname = f"DREAM_PT_dream_panel_control_net_{space_type}" + bl_label = "ControlNet" + bl_options = {'DEFAULT_CLOSED'} + + def draw(self, context): + layout = self.layout + prompt = get_prompt(context) + + row = layout.row() + row.template_list("SCENE_UL_ControlNetList", "", prompt, "control_nets", prompt, "active_control_net") + col = row.column(align=True) + col.operator("dream_textures.control_nets_add", icon='ADD', text="") + col.operator("dream_textures.control_nets_remove", icon='REMOVE', text="") + return ControlNetPanel + def advanced_panel(sub_panel, space_type, get_prompt): class AdvancedPanel(sub_panel): """Create a subpanel for advanced options""" From 316ec86cfd78b8786cdb99d623dbce769f788544 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Sat, 25 Mar 2023 18:08:45 -0400 Subject: [PATCH 24/28] Fix control image processing --- property_groups/dream_prompt.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/property_groups/dream_prompt.py b/property_groups/dream_prompt.py index fe8893fd..dd828811 100644 --- a/property_groups/dream_prompt.py +++ b/property_groups/dream_prompt.py @@ -279,8 +279,7 @@ def generate_args(self): args['controlnet_conditioning_scale'] = [net.conditioning_scale for net in args['control_nets']] args['control'] = [ np.flipud( - (np.array(net.control_image.pixels) * 255) - .astype(np.uint8) + np.array(net.control_image.pixels) .reshape((net.control_image.size[1], net.control_image.size[0], net.control_image.channels)) ) for net in args['control_nets'] From 9bbd3de6fbd72fc6e0ce5ad2a3c74f2dc3cafcaf Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Sat, 25 Mar 2023 21:16:06 -0400 Subject: [PATCH 25/28] Fix ControlNet projection --- __init__.py | 8 +++++++- operators/project.py | 10 ++++++---- property_groups/dream_prompt.py | 1 + 3 files changed, 14 insertions(+), 5 deletions(-) diff --git a/__init__.py b/__init__.py index caa39279..3b809fac 100644 --- a/__init__.py +++ b/__init__.py @@ -102,7 +102,13 @@ def get_selection_preview(self): bpy.types.Scene.dream_textures_project_prompt = PointerProperty(type=DreamPrompt) bpy.types.Scene.dream_textures_project_framebuffer_arguments = EnumProperty(name="Inputs", items=framebuffer_arguments) bpy.types.Scene.dream_textures_project_bake = BoolProperty(name="Bake", default=False, description="Re-maps the generated texture onto the specified UV map") - bpy.types.Scene.dream_textures_project_use_control_net = BoolProperty(name="Use ControlNet", default=False, description="Use a depth ControlNet instead of a depth model") + def project_use_controlnet(self, context): + if self.dream_textures_project_use_control_net: + if len(self.dream_textures_project_prompt.control_nets) < 1: + self.dream_textures_project_prompt.control_nets.add() + else: + self.dream_textures_project_prompt.control_nets.clear() + bpy.types.Scene.dream_textures_project_use_control_net = BoolProperty(name="Use ControlNet", default=False, description="Use a depth ControlNet instead of a depth model", update=project_use_controlnet) engine.register() diff --git a/operators/project.py b/operators/project.py index e6cb483d..948ce13c 100644 --- a/operators/project.py +++ b/operators/project.py @@ -126,9 +126,9 @@ def draw(self, context): col = layout.column() col.prop(context.scene, "dream_textures_project_use_control_net") - if context.scene.dream_textures_project_use_control_net: - col.prop(prompt, "control_net", text="Depth ControlNet") - col.prop(prompt, "controlnet_conditioning_scale") + if context.scene.dream_textures_project_use_control_net and len(prompt.control_nets) > 0: + col.prop(prompt.control_nets[0], "control_net", text="Depth ControlNet") + col.prop(prompt.control_nets[0], "conditioning_scale", text="ControlNet Conditioning Scale") col.prop(context.scene, "dream_textures_project_bake") if context.scene.dream_textures_project_bake: @@ -439,10 +439,12 @@ def on_exception(_, exception): context.scene.dream_textures_info = "Starting..." if context.scene.dream_textures_project_use_control_net: + generated_args = context.scene.dream_textures_project_prompt.generate_args() + del generated_args['control'] future = gen.control_net( control=[np.flipud(depth)], # the depth control needs to be flipped. image=init_img_path, - **context.scene.dream_textures_project_prompt.generate_args() + **generated_args ) else: future = gen.depth_to_image( diff --git a/property_groups/dream_prompt.py b/property_groups/dream_prompt.py index dd828811..8c6c1268 100644 --- a/property_groups/dream_prompt.py +++ b/property_groups/dream_prompt.py @@ -283,6 +283,7 @@ def generate_args(self): .reshape((net.control_image.size[1], net.control_image.size[0], net.control_image.channels)) ) for net in args['control_nets'] + if net.control_image is not None ] del args['control_nets'] return args From 1544714175d601d0005b41a35e73cd6d37053b7a Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Sun, 26 Mar 2023 22:58:35 -0400 Subject: [PATCH 26/28] Rough inpainting-ControlNet approach --- generator_process/actions/control_net.py | 116 +++++++++- operators/dream_texture.py | 1 + operators/project.py | 265 ++++++++++++++++++++++- 3 files changed, 367 insertions(+), 15 deletions(-) diff --git a/generator_process/actions/control_net.py b/generator_process/actions/control_net.py index 8feb196d..aa97da5a 100644 --- a/generator_process/actions/control_net.py +++ b/generator_process/actions/control_net.py @@ -23,6 +23,11 @@ def control_net( controlnet_conditioning_scale: list[float], image: NDArray | str | None, # image to image + # inpaint + inpaint: bool, + inpaint_mask_src: str, + text_mask: str, + text_mask_confidence: float, strength: float, prompt: str | list[str], @@ -115,15 +120,68 @@ def prepare_img2img_latents(self, image, timestep, batch_size, num_images_per_pr return latents + # copied from diffusers.StableDiffusionInpaintPipeline + def prepare_mask_latents( + self, mask, masked_image, batch_size, height, width, dtype, device, generator, do_classifier_free_guidance + ): + # resize the mask to latents shape as we concatenate the mask to the latents + # we do that before converting to dtype to avoid breaking in case we're using cpu_offload + # and half precision + mask = torch.nn.functional.interpolate( + mask, size=(height // self.vae_scale_factor, width // self.vae_scale_factor) + ) + mask = mask.to(device=device, dtype=dtype) + + masked_image = masked_image.to(device=device, dtype=dtype) + + # encode the mask image into latents space so we can concatenate it to the latents + if isinstance(generator, list): + masked_image_latents = [ + self.vae.encode(masked_image[i : i + 1]).latent_dist.sample(generator=generator[i]) + for i in range(batch_size) + ] + masked_image_latents = torch.cat(masked_image_latents, dim=0) + else: + masked_image_latents = self.vae.encode(masked_image).latent_dist.sample(generator=generator) + masked_image_latents = self.vae.config.scaling_factor * masked_image_latents + + # duplicate mask and masked_image_latents for each generation per prompt, using mps friendly method + if mask.shape[0] < batch_size: + if not batch_size % mask.shape[0] == 0: + raise ValueError( + "The passed mask and the required batch size don't match. Masks are supposed to be duplicated to" + f" a total batch size of {batch_size}, but {mask.shape[0]} masks were passed. Make sure the number" + " of masks that you pass is divisible by the total requested batch size." + ) + mask = mask.repeat(batch_size // mask.shape[0], 1, 1, 1) + if masked_image_latents.shape[0] < batch_size: + if not batch_size % masked_image_latents.shape[0] == 0: + raise ValueError( + "The passed images and the required batch size don't match. Images are supposed to be duplicated" + f" to a total batch size of {batch_size}, but {masked_image_latents.shape[0]} images were passed." + " Make sure the number of images that you pass is divisible by the total requested batch size." + ) + masked_image_latents = masked_image_latents.repeat(batch_size // masked_image_latents.shape[0], 1, 1, 1) + + mask = torch.cat([mask] * 2) if do_classifier_free_guidance else mask + masked_image_latents = ( + torch.cat([masked_image_latents] * 2) if do_classifier_free_guidance else masked_image_latents + ) + + # aligning device to prevent device errors when concating it with the latent model input + masked_image_latents = masked_image_latents.to(device=device, dtype=dtype) + return mask, masked_image_latents + @torch.no_grad() def __call__( self, prompt: Union[str, List[str]] = None, image: Union[torch.FloatTensor, PIL.Image.Image, List[torch.FloatTensor], List[PIL.Image.Image]] = None, - # NOTE: Modified to support initial image. + # NOTE: Modified to support initial image and inpaint. init_image: Union[torch.FloatTensor, PIL.Image.Image, List[torch.FloatTensor], List[PIL.Image.Image]] = None, strength: float = 1.0, + mask: Union[torch.FloatTensor, PIL.Image.Image, List[torch.FloatTensor], List[PIL.Image.Image]] = None, height: Optional[int] = None, width: Optional[int] = None, @@ -224,7 +282,7 @@ def __call__( # 5. Prepare timesteps # NOTE: Modified to support initial image - if init_image is not None: + if init_image is not None and not inpaint: self.scheduler.set_timesteps(num_inference_steps, device=device) timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) @@ -235,7 +293,37 @@ def __call__( # 6. Prepare latent variables num_channels_latents = self.unet.in_channels # NOTE: Modified to support initial image - if init_image is not None: + if mask is not None: + num_channels_latents = self.vae.config.latent_channels + mask, masked_image = diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_inpaint.prepare_mask_and_masked_image(init_image, mask) + mask, masked_image_latents = self.prepare_mask_latents( + mask, + masked_image, + batch_size * num_images_per_prompt, + height, + width, + prompt_embeds.dtype, + device, + generator, + do_classifier_free_guidance, + ) + num_channels_mask = mask.shape[1] + num_channels_masked_image = masked_image_latents.shape[1] + if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels: + raise ValueError( + f"Select an inpainting model, such as 'stabilityai/stable-diffusion-2-inpainting'" + ) + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + prompt_embeds.dtype, + device, + generator, + latents, + ) + elif init_image is not None: init_image = diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.preprocess(init_image) latents = self.prepare_img2img_latents( init_image, @@ -279,6 +367,9 @@ def __call__( return_dict=False, ) + if mask is not None: + latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1) + # predict the noise residual noise_pred = self.unet( latent_model_input, @@ -377,7 +468,21 @@ def __call__( int(8 * (height // 8)), ) control_image = [PIL.Image.fromarray(np.uint8(c * 255)).convert('RGB').resize(rounded_size) for c in control] if control is not None else None - init_image = None if image is None else (PIL.Image.open(image) if isinstance(image, str) else PIL.Image.fromarray(image.astype(np.uint8))).convert('RGB').resize(rounded_size) + init_image = None if image is None else (PIL.Image.open(image) if isinstance(image, str) else PIL.Image.fromarray(image.astype(np.uint8))).resize(rounded_size) + if inpaint: + match inpaint_mask_src: + case 'alpha': + mask_image = PIL.ImageOps.invert(init_image.getchannel('A')) + case 'prompt': + from transformers import AutoProcessor, CLIPSegForImageSegmentation + + processor = AutoProcessor.from_pretrained("CIDAS/clipseg-rd64-refined") + clipseg = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined") + inputs = processor(text=[text_mask], images=[init_image.convert('RGB')], return_tensors="pt", padding=True) + outputs = clipseg(**inputs) + mask_image = PIL.Image.fromarray(np.uint8((1 - torch.sigmoid(outputs.logits).lt(text_mask_confidence).int().detach().numpy()) * 255), 'L').resize(init_image.size) + else: + mask_image = None # Seamless if seamless_axes == SeamlessAxes.AUTO: @@ -399,7 +504,8 @@ def __call__( prompt=prompt, image=control_image, controlnet_conditioning_scale=controlnet_conditioning_scale, - init_image=init_image, + init_image=init_image.convert('RGB') if init_image is not None else None, + mask=mask_image, strength=strength, width=rounded_size[0], height=rounded_size[1], diff --git a/operators/dream_texture.py b/operators/dream_texture.py index 2da701d6..de6a23ae 100644 --- a/operators/dream_texture.py +++ b/operators/dream_texture.py @@ -180,6 +180,7 @@ def generate_next(): if len(generated_args['control_net']) > 0: f = gen.control_net( image=init_image, + inpaint=generated_args['init_img_action'] == 'inpaint', **generated_args ) elif init_image is not None: diff --git a/operators/project.py b/operators/project.py index 948ce13c..f077e1f5 100644 --- a/operators/project.py +++ b/operators/project.py @@ -198,7 +198,7 @@ def draw_depth_map(width, height, context, matrix, projection_matrix): offscreen.free() return depth -def bake(context, mesh, src, dest, src_uv, dest_uv): +def bake(context, mesh, width, height, src, src_uv, dest_uv): def bake_shader(): vert_out = gpu.types.GPUStageInterfaceInfo("my_interface") vert_out.smooth('VEC2', "uvInterp") @@ -227,7 +227,6 @@ def bake_shader(): return gpu.shader.create_from_info(shader_info) - width, height = dest.size[0], dest.size[1] offscreen = gpu.types.GPUOffScreen(width, height) buffer = gpu.types.Buffer('FLOAT', width * height * 4, src) @@ -252,7 +251,143 @@ def bake_shader(): batch.draw(shader) projected = np.array(fb.read_color(0, 0, width, height, 4, 0, 'FLOAT').to_list()) offscreen.free() - dest.pixels[:] = projected.ravel() + return projected + +def draw_color(width, height, context, mesh, matrix, projection_matrix, texture, uvs): + """ + Generate a depth map for the given matrices. + """ + offscreen = gpu.types.GPUOffScreen(width, height) + buffer = gpu.types.Buffer('FLOAT', width * height * 4, texture) + texture = gpu.types.GPUTexture(size=(width, height), data=buffer, format='RGBA16F') + + def color_shader(): + vert_out = gpu.types.GPUStageInterfaceInfo("my_interface") + vert_out.smooth('VEC2', "uvInterp") + + shader_info = gpu.types.GPUShaderCreateInfo() + shader_info.push_constant('MAT4', 'ModelViewProjectionMatrix') + shader_info.sampler(0, 'FLOAT_2D', "image") + shader_info.vertex_in(0, 'VEC3', "pos") + shader_info.vertex_in(1, 'VEC2', "uv") + shader_info.vertex_out(vert_out) + shader_info.fragment_out(0, 'VEC4', "fragColor") + + shader_info.vertex_source(""" +void main() +{ + gl_Position = ModelViewProjectionMatrix * vec4(pos, 1.0); + uvInterp = uv; +} +""") + + shader_info.fragment_source(""" +void main() +{ + fragColor = texture(image, uvInterp); +} +""") + + return gpu.shader.create_from_info(shader_info) + + with offscreen.bind(): + fb = gpu.state.active_framebuffer_get() + fb.clear(color=(0.0, 0.0, 0.0, 1.0)) + gpu.state.depth_test_set('LESS_EQUAL') + gpu.state.depth_mask_set(True) + with gpu.matrix.push_pop(): + gpu.matrix.load_matrix(matrix) + gpu.matrix.load_projection_matrix(projection_matrix) + + mesh.calc_loop_triangles() + + vertices = np.array([vert.co for vert in mesh.verts], dtype='f') + indices = np.array([[l.vert.index for l in loop] for loop in mesh.calc_loop_triangles()], dtype='i') + + shader = color_shader() + batch = batch_for_shader( + shader, 'TRIS', + {"pos": vertices, "uv": uvs}, + indices=indices, + ) + shader.uniform_sampler("image", texture) + batch.draw(shader) + color = np.array(fb.read_color(0, 0, width, height, 4, 0, 'FLOAT').to_list()) + offscreen.free() + + return color + +def draw_mask_map(width, height, context, mesh, matrix, projection_matrix, threshold, bake): + """ + Generate a depth map for the given matrices. + """ + offscreen = gpu.types.GPUOffScreen(width, height) + + def mask_shader(): + vert_out = gpu.types.GPUStageInterfaceInfo("my_interface") + vert_out.smooth('VEC4', "color") + + shader_info = gpu.types.GPUShaderCreateInfo() + shader_info.push_constant('MAT4', 'ModelViewProjectionMatrix') + shader_info.push_constant('MAT3', 'NormalMatrix') + shader_info.push_constant('FLOAT', 'Threshold') + shader_info.push_constant('VEC3', 'CameraNormal') + shader_info.vertex_in(0, 'VEC3', "pos") + shader_info.vertex_in(1, 'VEC3', "nor") + shader_info.vertex_out(vert_out) + shader_info.fragment_out(0, 'VEC4', "fragColor") + + shader_info.vertex_source(""" +void main() +{ + gl_Position = ModelViewProjectionMatrix * vec4(pos, 1.0); + color = dot(NormalMatrix * nor, CameraNormal) > Threshold ? vec4(1, 1, 1, 1) : vec4(0, 0, 0, 1); +} +""") + + + shader_info.fragment_source(""" +void main() +{ + fragColor = color; +} +""") + + return gpu.shader.create_from_info(shader_info) + + with offscreen.bind(): + fb = gpu.state.active_framebuffer_get() + fb.clear(color=(0.0, 0.0, 0.0, 1.0)) + gpu.state.depth_test_set('LESS_EQUAL') + gpu.state.depth_mask_set(True) + with gpu.matrix.push_pop(): + if bake is not None: + gpu.matrix.load_matrix(mathutils.Matrix.Identity(4)) + gpu.matrix.load_projection_matrix(mathutils.Matrix.Identity(4)) + else: + gpu.matrix.load_matrix(matrix) + gpu.matrix.load_projection_matrix(projection_matrix) + + mesh.calc_loop_triangles() + + vertices = np.array([vert.co for vert in mesh.verts], dtype='f') + normals = np.array([vert.normal for vert in mesh.verts], dtype='f') + indices = np.array([[l.vert.index for l in loop] for loop in mesh.calc_loop_triangles()], dtype='i') + + shader = mask_shader() + shader.uniform_float('CameraNormal', context.region_data.view_rotation @ mathutils.Vector((0, 0, 1)) if bake is not None else (0, 0, 1)) + shader.uniform_float('Threshold', threshold) + + batch = batch_for_shader( + shader, 'TRIS', + {"pos": bake if bake is not None else vertices, "nor": normals}, + indices=indices, + ) + batch.draw(shader) + mask = np.array(fb.read_color(0, 0, width, height, 4, 0, 'FLOAT').to_list()) + offscreen.free() + + return mask class ProjectDreamTexture(bpy.types.Operator): bl_idname = "shade.dream_texture_project" @@ -262,12 +397,13 @@ class ProjectDreamTexture(bpy.types.Operator): @classmethod def poll(cls, context): - try: - context.scene.dream_textures_project_prompt.validate(context, task=None if context.scene.dream_textures_project_use_control_net else ModelType.DEPTH) - _validate_projection(context) - except: - return False - return Generator.shared().can_use() + return True + # try: + # context.scene.dream_textures_project_prompt.validate(context, task=None if context.scene.dream_textures_project_use_control_net else ModelType.DEPTH) + # _validate_projection(context) + # except: + # return False + # return Generator.shared().can_use() @classmethod def get_uv_layer(cls, mesh: bmesh.types.BMesh): @@ -278,7 +414,116 @@ def get_uv_layer(cls, mesh: bmesh.types.BMesh): return mesh.loops.layers.uv.new("Projected UVs"), len(mesh.loops.layers.uv) - 1 + def draw_color(self, context): + res_x, res_y = context.scene.render.resolution_x, context.scene.render.resolution_y + view3d_spaces = [] + for area in context.screen.areas: + if area.type == 'VIEW_3D': + for region in area.regions: + if region.type == 'WINDOW': + context.scene.render.resolution_x, context.scene.render.resolution_y = region.width, region.height + for space in area.spaces: + if space.type == 'VIEW_3D': + if space.overlay.show_overlays: + view3d_spaces.append(space) + space.overlay.show_overlays = False + init_img_path = tempfile.NamedTemporaryFile(suffix='.png').name + render_filepath, file_format = context.scene.render.filepath, context.scene.render.image_settings.file_format + context.scene.render.image_settings.file_format = 'PNG' + context.scene.render.filepath = init_img_path + bpy.ops.render.opengl(write_still=True, view_context=True) + for space in view3d_spaces: + space.overlay.show_overlays = True + context.scene.render.resolution_x, context.scene.render.resolution_y = res_x, res_y + context.scene.render.filepath, context.scene.render.image_settings.file_format = render_filepath, file_format + return init_img_path + + def masked_init_image(self, context, camera, texture, mesh, split_mesh, uvs, threshold=0.75): + mask = draw_mask_map(512, 512, context, mesh, camera.matrix_world.inverted(), context.space_data.region_3d.window_matrix, threshold, bake=None) + color = draw_color(512, 512, context, split_mesh, camera.matrix_world.inverted(), context.space_data.region_3d.window_matrix, texture, uvs) + bake_uvs = [(uv[0] * 2 - 1, uv[1] * 2 - 1, 0) for uv in uvs] + baked_mask = draw_mask_map(512, 512, context, split_mesh, camera.matrix_world.inverted(), context.space_data.region_3d.window_matrix, threshold, bake=bake_uvs) + + return color, mask, baked_mask + def execute(self, context): + texture = bpy.data.images['result'] + generation_result = bpy.data.images['generation_result'] + inpaint_result = bpy.data.images['inpaint_result'] + mask_result = bpy.data.images['mask'] + + mesh = bmesh.from_edit_mesh(context.object.data) + mesh.verts.ensure_lookup_table() + mesh.verts.index_update() + + split_mesh = mesh.copy() + split_mesh.select_mode = {'FACE'} + bmesh.ops.split_edges(split_mesh, edges=split_mesh.edges) + + mesh.faces.ensure_lookup_table() + + uv_layer = split_mesh.loops.layers.uv.active + projection_uv_layer, _ = ProjectDreamTexture.get_uv_layer(split_mesh) + uvs = np.empty((len(split_mesh.verts), 2), dtype=np.float32) + projection_uvs = np.empty((len(split_mesh.verts), 2), dtype=np.float32) + for face in split_mesh.faces: + for loop in face.loops: + projection_uvs[loop.vert.index] = loop[projection_uv_layer].uv + uvs[loop.vert.index] = loop[uv_layer].uv + + gen = Generator.shared() + + def step(camera, inpaint): + depth = np.flipud( + draw_depth_map(512, 512, context, camera.matrix_world.inverted(), context.space_data.region_3d.window_matrix) + ) + generated_args = context.scene.dream_textures_project_prompt.generate_args() + del generated_args['model'] + del generated_args['control'] + del generated_args['inpaint_mask_src'] + + if inpaint: + color, mask, baked_mask = self.masked_init_image(context, camera, np.array(texture.pixels, dtype=np.float32), mesh, split_mesh, uvs) + color[:, :, 3] = 1 - mask[:, :, 0] + res = gen.control_net( + model='models--runwayml--stable-diffusion-inpainting', + control=[depth], + image=np.flipud(color * 255), + inpaint=True, + inpaint_mask_src='alpha', + **generated_args + ).result() + inpaint_result.pixels[:] = res[-1].images[0].ravel() + inpaint_result.update() + mask_result.pixels[:] = color.ravel() + mask_result.update() + color = bake(context, split_mesh, 512, 512, res[-1].images[0].ravel(), projection_uvs, uvs) + color = (np.array(texture.pixels, dtype=np.float32) * (1 - baked_mask)) + (color * baked_mask) + else: + res = gen.control_net( + model='models--runwayml--stable-diffusion-v1-5', + control=[depth], + image=None, + inpaint=False, + inpaint_mask_src='alpha', + **generated_args + ).result() + generation_result.pixels[:] = res[-1].images[0].ravel() + generation_result.update() + color = bake(context, split_mesh, 512, 512, res[-1].images[0].ravel(), projection_uvs, uvs) + + texture.pixels[:] = color.ravel() + texture.update() + + started = False + for camera in [ob for ob in bpy.context.scene.objects if ob.type == 'CAMERA' and not ob.hide_viewport]: + modelview_matrix = camera.matrix_world.inverted() + context.space_data.region_3d.view_matrix = modelview_matrix + step(camera, started) + started = True + + return {'FINISHED'} + # Setup the progress indicator def step_progress_update(self, context): if hasattr(context.area, "regions"): @@ -424,7 +669,7 @@ def on_done(future): for loop in face.loops: src_uvs[loop.vert.index] = loop[src_uv_layer].uv dest_uvs[loop.vert.index] = loop[dest_uv_layer].uv - bake(context, bm, generated.images[0].ravel(), dest, src_uvs, dest_uvs) + dest.pixels[:] = bake(context, bm, dest.size[0], dest.size[1], generated.images[0].ravel(), src_uvs, dest_uvs) dest.update() dest.pack() image_texture_node.image = dest From 40df72ac6a2b3b49686f327ae55d68c2744df058 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Mon, 27 Mar 2023 18:34:30 -0400 Subject: [PATCH 27/28] Use cameras for projection --- operators/project.py | 52 +++++++++++++++++++++++++++++++------------- 1 file changed, 37 insertions(+), 15 deletions(-) diff --git a/operators/project.py b/operators/project.py index f077e1f5..f478c899 100644 --- a/operators/project.py +++ b/operators/project.py @@ -198,23 +198,24 @@ def draw_depth_map(width, height, context, matrix, projection_matrix): offscreen.free() return depth -def bake(context, mesh, width, height, src, src_uv, dest_uv): +def bake(context, mesh, width, height, src, projection_uvs, uvs): def bake_shader(): vert_out = gpu.types.GPUStageInterfaceInfo("my_interface") vert_out.smooth('VEC2', "uvInterp") shader_info = gpu.types.GPUShaderCreateInfo() + shader_info.push_constant('MAT4', 'ModelViewProjectionMatrix') shader_info.sampler(0, 'FLOAT_2D', "image") - shader_info.vertex_in(0, 'VEC2', "src_uv") - shader_info.vertex_in(1, 'VEC2', "dest_uv") + shader_info.vertex_in(0, 'VEC2', "pos") + shader_info.vertex_in(1, 'VEC2', "uv") shader_info.vertex_out(vert_out) shader_info.fragment_out(0, 'VEC4', "fragColor") shader_info.vertex_source(""" void main() { - gl_Position = vec4(dest_uv * 2 - 1, 0.0, 1.0); - uvInterp = src_uv; + gl_Position = ModelViewProjectionMatrix * vec4(pos * 2 - 1, 0.0, 1.0); + uvInterp = uv; } """) @@ -244,7 +245,7 @@ def bake_shader(): shader = bake_shader() batch = batch_for_shader( shader, 'TRIS', - {"src_uv": src_uv, "dest_uv": dest_uv}, + {"pos": uvs, "uv": projection_uvs}, indices=vertices, ) shader.uniform_sampler("image", texture) @@ -325,7 +326,7 @@ def draw_mask_map(width, height, context, mesh, matrix, projection_matrix, thres def mask_shader(): vert_out = gpu.types.GPUStageInterfaceInfo("my_interface") - vert_out.smooth('VEC4', "color") + vert_out.smooth('VEC3', "normal") shader_info = gpu.types.GPUShaderCreateInfo() shader_info.push_constant('MAT4', 'ModelViewProjectionMatrix') @@ -341,7 +342,7 @@ def mask_shader(): void main() { gl_Position = ModelViewProjectionMatrix * vec4(pos, 1.0); - color = dot(NormalMatrix * nor, CameraNormal) > Threshold ? vec4(1, 1, 1, 1) : vec4(0, 0, 0, 1); + normal = nor; } """) @@ -349,7 +350,7 @@ def mask_shader(): shader_info.fragment_source(""" void main() { - fragColor = color; + fragColor = dot(NormalMatrix * normal, CameraNormal) > Threshold ? vec4(1, 1, 1, 1) : vec4(0, 0, 0, 1); } """) @@ -357,7 +358,7 @@ def mask_shader(): with offscreen.bind(): fb = gpu.state.active_framebuffer_get() - fb.clear(color=(0.0, 0.0, 0.0, 1.0)) + fb.clear(color=(0.0, 0.0, 0.0, 0.0)) gpu.state.depth_test_set('LESS_EQUAL') gpu.state.depth_mask_set(True) with gpu.matrix.push_pop(): @@ -444,6 +445,13 @@ def masked_init_image(self, context, camera, texture, mesh, split_mesh, uvs, thr bake_uvs = [(uv[0] * 2 - 1, uv[1] * 2 - 1, 0) for uv in uvs] baked_mask = draw_mask_map(512, 512, context, split_mesh, camera.matrix_world.inverted(), context.space_data.region_3d.window_matrix, threshold, bake=bake_uvs) + kernel = np.array([0, 0.1, 0.2, 1.0, 0.2, 0.1, 0]) + mask = np.apply_along_axis(lambda x: np.convolve(x, kernel, mode='same'), 0, mask) + mask = np.apply_along_axis(lambda x: np.convolve(x, kernel, mode='same'), 1, mask) + kernel = np.array([0.1, 0.1, 0.2, 0.3, 0.2, 0.1, 0.1]) + baked_mask = np.apply_along_axis(lambda x: np.convolve(x, kernel, mode='same'), 0, baked_mask) + baked_mask = np.apply_along_axis(lambda x: np.convolve(x, kernel, mode='same'), 1, baked_mask) + return color, mask, baked_mask def execute(self, context): @@ -457,22 +465,31 @@ def execute(self, context): mesh.verts.index_update() split_mesh = mesh.copy() + split_mesh.transform(context.object.matrix_world) split_mesh.select_mode = {'FACE'} bmesh.ops.split_edges(split_mesh, edges=split_mesh.edges) - mesh.faces.ensure_lookup_table() + split_mesh.faces.ensure_lookup_table() + split_mesh.verts.ensure_lookup_table() + split_mesh.verts.index_update() + + def vert_to_uv(v): + screen_space = view3d_utils.location_3d_to_region_2d(context.region, context.space_data.region_3d, v.co) + if screen_space is None: + return None + return (screen_space[0] / context.region.width, screen_space[1] / context.region.height) uv_layer = split_mesh.loops.layers.uv.active - projection_uv_layer, _ = ProjectDreamTexture.get_uv_layer(split_mesh) uvs = np.empty((len(split_mesh.verts), 2), dtype=np.float32) projection_uvs = np.empty((len(split_mesh.verts), 2), dtype=np.float32) for face in split_mesh.faces: for loop in face.loops: - projection_uvs[loop.vert.index] = loop[projection_uv_layer].uv + projection_uvs[loop.vert.index] = vert_to_uv(split_mesh.verts[loop.vert.index]) uvs[loop.vert.index] = loop[uv_layer].uv gen = Generator.shared() + step_i = 0 def step(camera, inpaint): depth = np.flipud( draw_depth_map(512, 512, context, camera.matrix_world.inverted(), context.space_data.region_3d.window_matrix) @@ -495,9 +512,10 @@ def step(camera, inpaint): ).result() inpaint_result.pixels[:] = res[-1].images[0].ravel() inpaint_result.update() - mask_result.pixels[:] = color.ravel() + mask_result.pixels[:] = baked_mask.ravel() mask_result.update() - color = bake(context, split_mesh, 512, 512, res[-1].images[0].ravel(), projection_uvs, uvs) + color = bake(context, split_mesh, 512, 512, res[-1].images[0].ravel(), projection_uvs, uvs).ravel() + baked_mask = baked_mask.ravel() color = (np.array(texture.pixels, dtype=np.float32) * (1 - baked_mask)) + (color * baked_mask) else: res = gen.control_net( @@ -511,9 +529,13 @@ def step(camera, inpaint): generation_result.pixels[:] = res[-1].images[0].ravel() generation_result.update() color = bake(context, split_mesh, 512, 512, res[-1].images[0].ravel(), projection_uvs, uvs) + # color = bake(context, split_mesh, 512, 512, np.array(generation_result.pixels, dtype=np.float32), projection_uvs, uvs) texture.pixels[:] = color.ravel() texture.update() + nonlocal step_i + bpy.data.images.new(name=f"Step {step_i}", width=512, height=512).pixels[:] = color.ravel() + step_i += 1 started = False for camera in [ob for ob in bpy.context.scene.objects if ob.type == 'CAMERA' and not ob.hide_viewport]: From 7884149f6cc09c820d4e16dd7964af6004e149d0 Mon Sep 17 00:00:00 2001 From: Carson Katri Date: Mon, 27 Mar 2023 19:20:50 -0400 Subject: [PATCH 28/28] Fixed projection --- operators/project.py | 80 ++++++++++++++++++++++++-------------------- 1 file changed, 43 insertions(+), 37 deletions(-) diff --git a/operators/project.py b/operators/project.py index f478c899..3e2d072a 100644 --- a/operators/project.py +++ b/operators/project.py @@ -439,19 +439,12 @@ def draw_color(self, context): context.scene.render.filepath, context.scene.render.image_settings.file_format = render_filepath, file_format return init_img_path - def masked_init_image(self, context, camera, texture, mesh, split_mesh, uvs, threshold=0.75): + def masked_init_image(self, context, camera, texture, mesh, split_mesh, uvs, threshold=0.5): mask = draw_mask_map(512, 512, context, mesh, camera.matrix_world.inverted(), context.space_data.region_3d.window_matrix, threshold, bake=None) color = draw_color(512, 512, context, split_mesh, camera.matrix_world.inverted(), context.space_data.region_3d.window_matrix, texture, uvs) bake_uvs = [(uv[0] * 2 - 1, uv[1] * 2 - 1, 0) for uv in uvs] baked_mask = draw_mask_map(512, 512, context, split_mesh, camera.matrix_world.inverted(), context.space_data.region_3d.window_matrix, threshold, bake=bake_uvs) - kernel = np.array([0, 0.1, 0.2, 1.0, 0.2, 0.1, 0]) - mask = np.apply_along_axis(lambda x: np.convolve(x, kernel, mode='same'), 0, mask) - mask = np.apply_along_axis(lambda x: np.convolve(x, kernel, mode='same'), 1, mask) - kernel = np.array([0.1, 0.1, 0.2, 0.3, 0.2, 0.1, 0.1]) - baked_mask = np.apply_along_axis(lambda x: np.convolve(x, kernel, mode='same'), 0, baked_mask) - baked_mask = np.apply_along_axis(lambda x: np.convolve(x, kernel, mode='same'), 1, baked_mask) - return color, mask, baked_mask def execute(self, context): @@ -460,37 +453,51 @@ def execute(self, context): inpaint_result = bpy.data.images['inpaint_result'] mask_result = bpy.data.images['mask'] - mesh = bmesh.from_edit_mesh(context.object.data) - mesh.verts.ensure_lookup_table() - mesh.verts.index_update() - - split_mesh = mesh.copy() - split_mesh.transform(context.object.matrix_world) - split_mesh.select_mode = {'FACE'} - bmesh.ops.split_edges(split_mesh, edges=split_mesh.edges) - - split_mesh.faces.ensure_lookup_table() - split_mesh.verts.ensure_lookup_table() - split_mesh.verts.index_update() - - def vert_to_uv(v): - screen_space = view3d_utils.location_3d_to_region_2d(context.region, context.space_data.region_3d, v.co) - if screen_space is None: - return None - return (screen_space[0] / context.region.width, screen_space[1] / context.region.height) - - uv_layer = split_mesh.loops.layers.uv.active - uvs = np.empty((len(split_mesh.verts), 2), dtype=np.float32) - projection_uvs = np.empty((len(split_mesh.verts), 2), dtype=np.float32) - for face in split_mesh.faces: - for loop in face.loops: - projection_uvs[loop.vert.index] = vert_to_uv(split_mesh.verts[loop.vert.index]) - uvs[loop.vert.index] = loop[uv_layer].uv - gen = Generator.shared() step_i = 0 def step(camera, inpaint): + def location_3d_to_region_2d(coord): + perspective_matrix = context.space_data.region_3d.window_matrix @ camera.matrix_world.inverted() + prj = perspective_matrix @ mathutils.Vector((coord[0], coord[1], coord[2], 1.0)) + if prj.w > 0.0: + width_half = context.region.width / 2.0 + height_half = context.region.height / 2.0 + + return mathutils.Vector(( + width_half + width_half * (prj.x / prj.w), + height_half + height_half * (prj.y / prj.w), + )) + else: + return None + + mesh = bmesh.from_edit_mesh(context.object.data) + mesh.verts.ensure_lookup_table() + mesh.verts.index_update() + + split_mesh = mesh.copy() + split_mesh.transform(context.object.matrix_world) + split_mesh.select_mode = {'FACE'} + bmesh.ops.split_edges(split_mesh, edges=split_mesh.edges) + + split_mesh.faces.ensure_lookup_table() + split_mesh.verts.ensure_lookup_table() + split_mesh.verts.index_update() + + def vert_to_uv(v): + screen_space = location_3d_to_region_2d(v.co) + if screen_space is None: + return None + return (screen_space[0] / context.region.width, screen_space[1] / context.region.height) + + uv_layer = split_mesh.loops.layers.uv.active + uvs = np.empty((len(split_mesh.verts), 2), dtype=np.float32) + projection_uvs = np.empty((len(split_mesh.verts), 2), dtype=np.float32) + for face in split_mesh.faces: + for loop in face.loops: + projection_uvs[loop.vert.index] = vert_to_uv(split_mesh.verts[loop.vert.index]) + uvs[loop.vert.index] = loop[uv_layer].uv + depth = np.flipud( draw_depth_map(512, 512, context, camera.matrix_world.inverted(), context.space_data.region_3d.window_matrix) ) @@ -512,7 +519,7 @@ def step(camera, inpaint): ).result() inpaint_result.pixels[:] = res[-1].images[0].ravel() inpaint_result.update() - mask_result.pixels[:] = baked_mask.ravel() + mask_result.pixels[:] = mask.ravel() mask_result.update() color = bake(context, split_mesh, 512, 512, res[-1].images[0].ravel(), projection_uvs, uvs).ravel() baked_mask = baked_mask.ravel() @@ -529,7 +536,6 @@ def step(camera, inpaint): generation_result.pixels[:] = res[-1].images[0].ravel() generation_result.update() color = bake(context, split_mesh, 512, 512, res[-1].images[0].ravel(), projection_uvs, uvs) - # color = bake(context, split_mesh, 512, 512, np.array(generation_result.pixels, dtype=np.float32), projection_uvs, uvs) texture.pixels[:] = color.ravel() texture.update()