From 2e3a3a3ec725e324464d2cef58372d844f3d14bf Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sun, 29 Sep 2024 12:35:48 -0400 Subject: [PATCH] add ctrl+x Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 10 +- modules/ctrlx/__init__.py | 662 +++++++++++++++++++++++++++++++ modules/ctrlx/features.py | 70 ++++ modules/ctrlx/media.py | 21 + modules/ctrlx/sdxl.py | 299 ++++++++++++++ modules/ctrlx/utils.py | 100 +++++ modules/sd_models.py | 4 +- modules/sd_samplers_diffusers.py | 6 +- scripts/ctrlx.py | 92 +++++ 9 files changed, 1258 insertions(+), 6 deletions(-) create mode 100644 modules/ctrlx/__init__.py create mode 100644 modules/ctrlx/features.py create mode 100644 modules/ctrlx/media.py create mode 100644 modules/ctrlx/sdxl.py create mode 100644 modules/ctrlx/utils.py create mode 100644 scripts/ctrlx.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 778c5288a..d6040955c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,6 @@ # Change Log for SD.Next -## Update for 2024-09-27 +## Update for 2024-09-29 - **reprocess** - new top-level button: reprocess your last generated image(s) @@ -18,6 +18,14 @@ *note* sd/sdxl contain heavily distilled versions of reference models, so switching to reference model produces vastly different results - xyz grid support for text encoder - full prompt parser now correctly works with different prompts in batch +- [Ctrl+X](https://github.com/genforce/ctrl-x): + - control **structure** (*similar to controlnet*) and **appearance** (*similar to ipadapter*) + without the need for extra models, all via code feed-forwards! + - can run in structure-only or appearance-only or both modes + - when providing structure and appearance input images, its best to provide a short prompts describing them + - structure image can be *almost anything*: *actual photo, openpose-style stick man, 3d render, sketch, depth-map, etc.* + just describe what it is in a structure prompt so it can be de-structured and correctly applied + - supports sdxl in both txt2img and img2img, simply select from scripts - **flux** - avoid unet load if unchanged - mark specific unet as unavailable if load failed diff --git a/modules/ctrlx/__init__.py b/modules/ctrlx/__init__.py new file mode 100644 index 000000000..c5afc2926 --- /dev/null +++ b/modules/ctrlx/__init__.py @@ -0,0 +1,662 @@ +from copy import deepcopy +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from diffusers import StableDiffusionXLPipeline +from diffusers.image_processor import PipelineImageInput +from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img import rescale_noise_cfg, retrieve_latents, retrieve_timesteps +from diffusers.utils import BaseOutput, deprecate +from diffusers.utils.torch_utils import randn_tensor +import numpy as np +import PIL +import torch +from .sdxl import register_attr +from .media import preprocess +from .utils import batch_dict_to_tensor, batch_tensor_to_dict, noise_prev, noise_t2t + + +BATCH_ORDER = [ + "structure_uncond", "appearance_uncond", "uncond", "structure_cond", "appearance_cond", "cond", +] + + +def get_last_control_i(control_schedule, num_inference_steps): + if control_schedule is None: + return num_inference_steps, num_inference_steps + + def max_(l): + if len(l) == 0: + return 0.0 + return max(l) + + structure_max = 0.0 + appearance_max = 0.0 + for block in control_schedule.values(): + if isinstance(block, list): # Handling mid_block + block = {0: block} + for layer in block.values(): + structure_max = max(structure_max, max_(layer[0] + layer[1])) + appearance_max = max(appearance_max, max_(layer[2])) + + structure_i = round(num_inference_steps * structure_max) + appearance_i = round(num_inference_steps * appearance_max) + return structure_i, appearance_i + + +@dataclass +class CtrlXStableDiffusionXLPipelineOutput(BaseOutput): + images: Union[List[PIL.Image.Image], np.ndarray] = None + structures: Union[List[PIL.Image.Image], np.ndarray] = None + appearances: Union[List[PIL.Image.Image], np.ndarray] = None + + +class CtrlXStableDiffusionXLPipeline(StableDiffusionXLPipeline): # diffusers==0.28.0 + + def prepare_latents( + self, image, batch_size, num_images_per_prompt, num_channels_latents, height, width, + dtype, device, generator=None, noise=None, + ): + batch_size = batch_size * num_images_per_prompt + + if noise is None: + shape = ( + batch_size, + num_channels_latents, + height // self.vae_scale_factor, + width // self.vae_scale_factor + ) + noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + noise = noise * self.scheduler.init_noise_sigma # Starting noise, need to scale + else: + noise = noise.to(device) + + if image is None: + return noise, 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)}" + ) + + # Offload text encoder if `enable_model_cpu_offload` was enabled + if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: + self.text_encoder_2.to("cpu") + torch.cuda.empty_cache() + + image = image.to(device=device, dtype=dtype) + + if image.shape[1] == 4: # Image already in latents form + init_latents = image + + else: + # Make sure the VAE is in float32 mode, as it overflows in float16 + if self.vae.config.force_upcast: + image = image.to(torch.float32) + self.vae.to(torch.float32) + + 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." + ) + elif isinstance(generator, list): + init_latents = [ + retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i]) + for i in range(batch_size) + ] + init_latents = torch.cat(init_latents, dim=0) + else: + init_latents = retrieve_latents(self.vae.encode(image), generator=generator) + + if self.vae.config.force_upcast: + self.vae.to(dtype) + + init_latents = init_latents.to(dtype) + 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 + 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) + + return noise, init_latents + + @property + def structure_guidance_scale(self): + return self._guidance_scale if self._structure_guidance_scale is None else self._structure_guidance_scale + + @property + def appearance_guidance_scale(self): + return self._guidance_scale if self._appearance_guidance_scale is None else self._appearance_guidance_scale + + @torch.no_grad() + def __call__( + self, + prompt: Union[str, List[str]] = None, # TODO: Support prompt_2 and negative_prompt_2 + structure_prompt: Optional[Union[str, List[str]]] = None, + appearance_prompt: Optional[Union[str, List[str]]] = None, + structure_image: Optional[PipelineImageInput] = None, + appearance_image: Optional[PipelineImageInput] = None, + num_inference_steps: int = 50, + timesteps: List[int] = None, + negative_prompt: Optional[Union[str, List[str]]] = None, + positive_prompt: Optional[Union[str, List[str]]] = None, + height: Optional[int] = None, + width: Optional[int] = None, + guidance_scale: float = 5.0, + structure_guidance_scale: Optional[float] = None, + appearance_guidance_scale: Optional[float] = None, + num_images_per_prompt: Optional[int] = 1, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + structure_latents: Optional[torch.Tensor] = None, + appearance_latents: Optional[torch.Tensor] = None, + prompt_embeds: Optional[torch.Tensor] = None, # Positive prompt is concatenated with prompt, so no embeddings + structure_prompt_embeds: Optional[torch.Tensor] = None, + appearance_prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + pooled_prompt_embeds: Optional[torch.Tensor] = None, + structure_pooled_prompt_embeds: Optional[torch.Tensor] = None, + appearance_pooled_prompt_embeds: Optional[torch.Tensor] = None, + negative_pooled_prompt_embeds: Optional[torch.Tensor] = None, + control_schedule: Optional[Dict] = None, + self_recurrence_schedule: Optional[List[int]] = [], # Format: [(start, end, num_repeat)] + decode_structure: Optional[bool] = True, + decode_appearance: Optional[bool] = True, + output_type: Optional[str] = "pil", + return_dict: bool = True, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + guidance_rescale: float = 0.0, + original_size: Tuple[int, int] = None, + crops_coords_top_left: Tuple[int, int] = (0, 0), + target_size: Tuple[int, int] = None, + clip_skip: Optional[int] = None, + callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + **kwargs, + ): + # TODO: Add function argument documentation + + callback = kwargs.pop("callback", None) + callback_steps = kwargs.pop("callback_steps", None) + + if callback is not None: + deprecate( + "callback", + "1.0.0", + "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", + ) + if callback_steps is not None: + deprecate( + "callback_steps", + "1.0.0", + "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", + ) + + # 0. Default height and width to U-Net + height = height or self.default_sample_size * self.vae_scale_factor + width = width or self.default_sample_size * self.vae_scale_factor + original_size = original_size or (height, width) + target_size = target_size or (height, width) + + # 1. Check inputs. Raise error if not correct + self.check_inputs( # TODO: Custom check_inputs for our method + prompt, + None, # prompt_2 + height, + width, + callback_steps, + negative_prompt = negative_prompt, + negative_prompt_2 = None, # negative_prompt_2 + prompt_embeds = prompt_embeds, + negative_prompt_embeds = negative_prompt_embeds, + pooled_prompt_embeds = pooled_prompt_embeds, + negative_pooled_prompt_embeds = negative_pooled_prompt_embeds, + callback_on_step_end_tensor_inputs = callback_on_step_end_tensor_inputs, + ) + + self._guidance_scale = guidance_scale + self._structure_guidance_scale = structure_guidance_scale + self._appearance_guidance_scale = appearance_guidance_scale + self._guidance_rescale = guidance_rescale + self._clip_skip = clip_skip + self._cross_attention_kwargs = cross_attention_kwargs + self._denoising_end = None # denoising_end + self._denoising_start = None # denoising_start + self._interrupt = False + + # 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] + + if batch_size * num_images_per_prompt != 1: + raise ValueError( + f"Pipeline currently does not support batch_size={batch_size} and num_images_per_prompt=1. " + "Effective batch size (batch_size * num_images_per_prompt) must be 1." + ) + + device = self._execution_device + + # 3. Encode input prompt + text_encoder_lora_scale = ( + self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None + ) + + if positive_prompt is not None and positive_prompt != "": + prompt = prompt + ", " + positive_prompt # Add positive prompt with comma + # By default, only add positive prompt to the appearance prompt and not the structure prompt + if appearance_prompt is not None and appearance_prompt != "": + appearance_prompt = appearance_prompt + ", " + positive_prompt + + ( + prompt_embeds_, + negative_prompt_embeds, + pooled_prompt_embeds_, + negative_pooled_prompt_embeds, + ) = self.encode_prompt( + prompt = prompt, + prompt_2 = None, # prompt_2 + device = device, + num_images_per_prompt = num_images_per_prompt, + do_classifier_free_guidance = True, # self.do_classifier_free_guidance, TODO: Support no CFG + negative_prompt = negative_prompt, + negative_prompt_2 = None, # negative_prompt_2 + prompt_embeds = prompt_embeds, + negative_prompt_embeds = negative_prompt_embeds, + pooled_prompt_embeds = pooled_prompt_embeds, + negative_pooled_prompt_embeds = negative_pooled_prompt_embeds, + lora_scale = text_encoder_lora_scale, + clip_skip = self.clip_skip, + ) + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds_], dim=0).to(device) + add_text_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds_], dim=0).to(device) + + # 3.1. Structure prompt embeddings + if structure_prompt is not None and structure_prompt != "": + ( + structure_prompt_embeds, + negative_structure_prompt_embeds, + structure_pooled_prompt_embeds, + negative_structure_pooled_prompt_embeds, + ) = self.encode_prompt( + prompt = structure_prompt, + prompt_2 = None, # prompt_2 + device = device, + num_images_per_prompt = num_images_per_prompt, + do_classifier_free_guidance = True, # self.do_classifier_free_guidance, TODO: Support no CFG + negative_prompt = negative_prompt if structure_image is None else "", + negative_prompt_2 = None, # negative_prompt_2 + prompt_embeds = structure_prompt_embeds, + negative_prompt_embeds = None, # negative_prompt_embeds + pooled_prompt_embeds = structure_pooled_prompt_embeds, + negative_pooled_prompt_embeds = None, # negative_pooled_prompt_embeds + lora_scale = text_encoder_lora_scale, + clip_skip = self.clip_skip, + ) + structure_prompt_embeds = torch.cat( + [negative_structure_prompt_embeds, structure_prompt_embeds], dim=0 + ).to(device) + structure_add_text_embeds = torch.cat( + [negative_structure_pooled_prompt_embeds, structure_pooled_prompt_embeds], dim=0 + ).to(device) + else: + structure_prompt_embeds = prompt_embeds + structure_add_text_embeds = add_text_embeds + + # 3.2. Appearance prompt embeddings + if appearance_prompt is not None and appearance_prompt != "": + ( + appearance_prompt_embeds, + negative_appearance_prompt_embeds, + appearance_pooled_prompt_embeds, + negative_appearance_pooled_prompt_embeds, + ) = self.encode_prompt( + prompt = appearance_prompt, + prompt_2 = None, # prompt_2 + device = device, + num_images_per_prompt = num_images_per_prompt, + do_classifier_free_guidance = True, # self.do_classifier_free_guidance, TODO: Support no CFG + negative_prompt = negative_prompt if appearance_image is None else "", + negative_prompt_2 = None, # negative_prompt_2 + prompt_embeds = appearance_prompt_embeds, + negative_prompt_embeds = None, # negative_prompt_embeds + pooled_prompt_embeds = appearance_pooled_prompt_embeds, # pooled_prompt_embeds + negative_pooled_prompt_embeds = None, # negative_pooled_prompt_embeds + lora_scale = text_encoder_lora_scale, + clip_skip = self.clip_skip, + ) + appearance_prompt_embeds = torch.cat( + [negative_appearance_prompt_embeds, appearance_prompt_embeds], dim=0 + ).to(device) + appearance_add_text_embeds = torch.cat( + [negative_appearance_pooled_prompt_embeds, appearance_pooled_prompt_embeds], dim=0 + ).to(device) + else: + appearance_prompt_embeds = prompt_embeds + appearance_add_text_embeds = add_text_embeds + + # 3.3. Prepare added time ids & embeddings, TODO: Support no CFG + if self.text_encoder_2 is None: + text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1]) + else: + text_encoder_projection_dim = self.text_encoder_2.config.projection_dim + + add_time_ids = self._get_add_time_ids( + original_size, + crops_coords_top_left, + target_size, + dtype = prompt_embeds.dtype, + text_encoder_projection_dim = text_encoder_projection_dim, + ) + negative_add_time_ids = add_time_ids + add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0).to(device) + + # 4. Prepare timesteps + timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) + + # 5. Prepare latent variables + num_channels_latents = self.unet.config.in_channels + + latents, _ = self.prepare_latents( + None, batch_size, num_images_per_prompt, num_channels_latents, height, width, + prompt_embeds.dtype, device, generator, latents + ) + + if structure_image is not None: + structure_image = preprocess( # Center crop + resize + structure_image, self.image_processor, height=height, width=width, resize_mode="crop" + ) + _, clean_structure_latents = self.prepare_latents( + structure_image, batch_size, num_images_per_prompt, num_channels_latents, height, width, + prompt_embeds.dtype, device, generator, structure_latents, + ) + else: + clean_structure_latents = None + structure_latents = latents if structure_latents is None else structure_latents + + if appearance_image is not None: + appearance_image = preprocess( # Center crop + resize + appearance_image, self.image_processor, height=height, width=width, resize_mode="crop" + ) + _, clean_appearance_latents = self.prepare_latents( + appearance_image, batch_size, num_images_per_prompt, num_channels_latents, height, width, + prompt_embeds.dtype, device, generator, appearance_latents, + ) + else: + clean_appearance_latents = None + appearance_latents = latents if appearance_latents is None else appearance_latents + + # 6. Prepare extra step kwargs + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + + # 7. Denoising loop + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) + + # 7.1 Apply denoising_end + def denoising_value_valid(dnv): + return isinstance(self.denoising_end, float) and 0 < dnv < 1 + + if ( + self.denoising_end is not None + and self.denoising_start is not None + and denoising_value_valid(self.denoising_end) + and denoising_value_valid(self.denoising_start) + and self.denoising_start >= self.denoising_end + ): + raise ValueError( + f"`denoising_start`: {self.denoising_start} cannot be larger than or equal to `denoising_end`: " + + f" {self.denoising_end} when using type float." + ) + elif self.denoising_end is not None and denoising_value_valid(self.denoising_end): + discrete_timestep_cutoff = int( + round( + self.scheduler.config.num_train_timesteps + - (self.denoising_end * self.scheduler.config.num_train_timesteps) + ) + ) + num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))) + timesteps = timesteps[:num_inference_steps] + + # 7.2 Optionally get guidance scale embedding + timestep_cond = None + if self.unet.config.time_cond_proj_dim is not None: # TODO: Make guidance scale embedding work with batch_order + guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(batch_size * num_images_per_prompt) + timestep_cond = self.get_guidance_scale_embedding( + guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim + ).to(device=device, dtype=latents.dtype) + + # 7.3 Get batch order + batch_order = deepcopy(BATCH_ORDER) + if structure_image is not None: # If image is provided, not generating, so no CFG needed + batch_order.remove("structure_uncond") + if appearance_image is not None: + batch_order.remove("appearance_uncond") + + structure_control_stop_i, appearance_control_stop_i = get_last_control_i(control_schedule, num_inference_steps) + if self_recurrence_schedule is None or len(self_recurrence_schedule) == 0: + self_recurrence_schedule = [0] * num_inference_steps + + self._num_timesteps = len(timesteps) + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + if i == structure_control_stop_i: # If not generating structure/appearance, drop after last control + if "structure_uncond" not in batch_order: + batch_order.remove("structure_cond") + if i == appearance_control_stop_i: + if "appearance_uncond" not in batch_order: + batch_order.remove("appearance_cond") + + register_attr(self, t=t.item(), do_control=True, batch_order=batch_order) + + # TODO: For now, assume we are doing classifier-free guidance, support no CF-guidance later + latent_model_input = self.scheduler.scale_model_input(latents, t) + structure_latent_model_input = self.scheduler.scale_model_input(structure_latents, t) + appearance_latent_model_input = self.scheduler.scale_model_input(appearance_latents, t) + + all_latent_model_input = { + "structure_uncond": structure_latent_model_input[0:1], + "appearance_uncond": appearance_latent_model_input[0:1], + "uncond": latent_model_input[0:1], + "structure_cond": structure_latent_model_input[0:1], + "appearance_cond": appearance_latent_model_input[0:1], + "cond": latent_model_input[0:1], + } + all_prompt_embeds = { + "structure_uncond": structure_prompt_embeds[0:1], + "appearance_uncond": appearance_prompt_embeds[0:1], + "uncond": prompt_embeds[0:1], + "structure_cond": structure_prompt_embeds[1:2], + "appearance_cond": appearance_prompt_embeds[1:2], + "cond": prompt_embeds[1:2], + } + all_add_text_embeds = { + "structure_uncond": structure_add_text_embeds[0:1], + "appearance_uncond": appearance_add_text_embeds[0:1], + "uncond": add_text_embeds[0:1], + "structure_cond": structure_add_text_embeds[1:2], + "appearance_cond": appearance_add_text_embeds[1:2], + "cond": add_text_embeds[1:2], + } + all_time_ids = { + "structure_uncond": add_time_ids[0:1], + "appearance_uncond": add_time_ids[0:1], + "uncond": add_time_ids[0:1], + "structure_cond": add_time_ids[1:2], + "appearance_cond": add_time_ids[1:2], + "cond": add_time_ids[1:2], + } + + concat_latent_model_input = batch_dict_to_tensor(all_latent_model_input, batch_order) + concat_prompt_embeds = batch_dict_to_tensor(all_prompt_embeds, batch_order) + concat_add_text_embeds = batch_dict_to_tensor(all_add_text_embeds, batch_order) + concat_add_time_ids = batch_dict_to_tensor(all_time_ids, batch_order) + + # Predict the noise residual + added_cond_kwargs = {"text_embeds": concat_add_text_embeds, "time_ids": concat_add_time_ids} + + concat_noise_pred = self.unet( + concat_latent_model_input, + t, + encoder_hidden_states = concat_prompt_embeds, + timestep_cond = timestep_cond, + cross_attention_kwargs = self.cross_attention_kwargs, + added_cond_kwargs = added_cond_kwargs, + ).sample + all_noise_pred = batch_tensor_to_dict(concat_noise_pred, batch_order) + + # Classifier-free guidance, TODO: Support no CFG + noise_pred = all_noise_pred["uncond"] +\ + self.guidance_scale * (all_noise_pred["cond"] - all_noise_pred["uncond"]) + + structure_noise_pred = all_noise_pred["structure_cond"]\ + if "structure_cond" in batch_order else noise_pred + if "structure_uncond" in all_noise_pred: + structure_noise_pred = all_noise_pred["structure_uncond"] +\ + self.structure_guidance_scale * (structure_noise_pred - all_noise_pred["structure_uncond"]) + + appearance_noise_pred = all_noise_pred["appearance_cond"]\ + if "appearance_cond" in batch_order else noise_pred + if "appearance_uncond" in all_noise_pred: + appearance_noise_pred = all_noise_pred["appearance_uncond"] +\ + self.appearance_guidance_scale * (appearance_noise_pred - all_noise_pred["appearance_uncond"]) + + if self.guidance_rescale > 0.0: + noise_pred = rescale_noise_cfg( + noise_pred, all_noise_pred["cond"], guidance_rescale=self.guidance_rescale + ) + if "structure_uncond" in all_noise_pred: + structure_noise_pred = rescale_noise_cfg( + structure_noise_pred, all_noise_pred["structure_cond"], + guidance_rescale=self.guidance_rescale + ) + if "appearance_uncond" in all_noise_pred: + appearance_noise_pred = rescale_noise_cfg( + appearance_noise_pred, all_noise_pred["appearance_cond"], + guidance_rescale=self.guidance_rescale + ) + + # Compute the previous noisy sample x_t -> x_t-1 + concat_noise_pred = torch.cat( + [structure_noise_pred, appearance_noise_pred, noise_pred], dim=0, + ) + concat_latents = torch.cat( + [structure_latents, appearance_latents, latents], dim=0, + ) + structure_latents, appearance_latents, latents = self.scheduler.step( + concat_noise_pred, t, concat_latents, **extra_step_kwargs, + ).prev_sample.chunk(3) + + if clean_structure_latents is not None: + structure_latents = noise_prev(self.scheduler, t, clean_structure_latents) + if clean_appearance_latents is not None: + appearance_latents = noise_prev(self.scheduler, t, clean_appearance_latents) + + # Self-recurrence + for _ in range(self_recurrence_schedule[i]): + if hasattr(self.scheduler, "_step_index"): # For fancier schedulers + self.scheduler._step_index -= 1 # TODO: Does this actually work? + + t_prev = 0 if i + 1 >= num_inference_steps else timesteps[i + 1] + latents = noise_t2t(self.scheduler, t_prev, t, latents) + latent_model_input = torch.cat([latents] * 2) + + register_attr(self, t=t.item(), do_control=False, batch_order=["uncond", "cond"]) + + # Predict the noise residual + added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} + noise_pred_uncond, noise_pred_ = self.unet( + latent_model_input, + t, + encoder_hidden_states = prompt_embeds, + timestep_cond = timestep_cond, + cross_attention_kwargs = self.cross_attention_kwargs, + added_cond_kwargs = added_cond_kwargs, + ).sample.chunk(2) + noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_ - noise_pred_uncond) + + if self.guidance_rescale > 0.0: + noise_pred = rescale_noise_cfg(noise_pred, noise_pred_, guidance_rescale=self.guidance_rescale) + + latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample + + # Callbacks + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) + add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds) + negative_pooled_prompt_embeds = callback_outputs.pop("negative_pooled_prompt_embeds", negative_pooled_prompt_embeds) + add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids) + # add_neg_time_ids = callback_outputs.pop("add_neg_time_ids", add_neg_time_ids) + + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + if callback is not None and i % callback_steps == 0: + step_idx = i // getattr(self.scheduler, "order", 1) + callback(step_idx, t, latents) + + # "Reconstruction" + if clean_structure_latents is not None: + structure_latents = clean_structure_latents + if clean_appearance_latents is not None: + appearance_latents = clean_appearance_latents + + # For passing important information onto the refiner + self.refiner_args = {"latents": latents.detach(), "prompt": prompt, "negative_prompt": negative_prompt} + + if output_type != "latent": + # Make sure the VAE is in float32 mode, as it overflows in float16 + if self.vae.config.force_upcast: + self.upcast_vae() + vae_dtype = next(iter(self.vae.post_quant_conv.parameters())).dtype + latents = latents.to(vae_dtype) + structure_latents = structure_latents.to(vae_dtype) + appearance_latents = appearance_latents.to(vae_dtype) + + image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] + image = self.image_processor.postprocess(image, output_type=output_type) + if decode_structure: + structure = self.vae.decode(structure_latents / self.vae.config.scaling_factor, return_dict=False)[0] + structure = self.image_processor.postprocess(structure, output_type=output_type) + else: + structure = structure_latents + if decode_appearance: + appearance = self.vae.decode(appearance_latents / self.vae.config.scaling_factor, return_dict=False)[0] + appearance = self.image_processor.postprocess(appearance, output_type=output_type) + else: + appearance = appearance_latents + + # Cast back to fp16 if needed + if self.vae.config.force_upcast: + self.vae.to(dtype=torch.float16) + + else: + # combined = torch.cat([latents, structure_latents, appearance_latents], dim=0) + # return CtrlXStableDiffusionXLPipelineOutput(images=combined) + return CtrlXStableDiffusionXLPipelineOutput(images=latents, structures=structure_latents, appearances=appearance_latents) + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (image, structure, appearance) + + return CtrlXStableDiffusionXLPipelineOutput(images=image, structures=structure, appearances=appearance) diff --git a/modules/ctrlx/features.py b/modules/ctrlx/features.py new file mode 100644 index 000000000..0fba30f49 --- /dev/null +++ b/modules/ctrlx/features.py @@ -0,0 +1,70 @@ +import torch.nn.functional as F +from .utils import batch_dict_to_tensor, batch_tensor_to_dict + + +def get_schedule(timesteps, schedule): + end = round(len(timesteps) * schedule) + timesteps = timesteps[:end] + return timesteps + + +def get_elem(l, i, default=0.0): + if i >= len(l): + return default + return l[i] + + +def pad_list(l_1, l_2, pad=0.0): + max_len = max(len(l_1), len(l_2)) + l_1 = l_1 + [pad] * (max_len - len(l_1)) + l_2 = l_2 + [pad] * (max_len - len(l_2)) + return l_1, l_2 + + +def normalize(x, dim): + x_mean = x.mean(dim=dim, keepdim=True) + x_std = x.std(dim=dim, keepdim=True) + x_normalized = (x - x_mean) / x_std + return x_normalized + + +# https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html +def appearance_mean_std(q_c_normed, k_s_normed, v_s): # c: content, s: style + q_c = q_c_normed # q_c and k_s must be projected from normalized features + k_s = k_s_normed + mean = F.scaled_dot_product_attention(q_c, k_s, v_s) # Use scaled_dot_product_attention for efficiency + std = (F.scaled_dot_product_attention(q_c, k_s, v_s.square()) - mean.square()).relu().sqrt() + + return mean, std + + +def feature_injection(features, batch_order): + assert features.shape[0] % len(batch_order) == 0 + features_dict = batch_tensor_to_dict(features, batch_order) + features_dict["cond"] = features_dict["structure_cond"] + features = batch_dict_to_tensor(features_dict, batch_order) + return features + + +def appearance_transfer(features, q_normed, k_normed, batch_order, v=None, reshape_fn=None): + assert features.shape[0] % len(batch_order) == 0 + + features_dict = batch_tensor_to_dict(features, batch_order) + q_normed_dict = batch_tensor_to_dict(q_normed, batch_order) + k_normed_dict = batch_tensor_to_dict(k_normed, batch_order) + v_dict = features_dict + if v is not None: + v_dict = batch_tensor_to_dict(v, batch_order) + + mean_cond, std_cond = appearance_mean_std( + q_normed_dict["cond"], k_normed_dict["appearance_cond"], v_dict["appearance_cond"], + ) + + if reshape_fn is not None: + mean_cond = reshape_fn(mean_cond) + std_cond = reshape_fn(std_cond) + + features_dict["cond"] = std_cond * normalize(features_dict["cond"], dim=-2) + mean_cond + + features = batch_dict_to_tensor(features_dict, batch_order) + return features diff --git a/modules/ctrlx/media.py b/modules/ctrlx/media.py new file mode 100644 index 000000000..39086d5d5 --- /dev/null +++ b/modules/ctrlx/media.py @@ -0,0 +1,21 @@ +import numpy as np +import torch +import torchvision.transforms.functional as vF +import PIL + + +JPEG_QUALITY = 95 + + +def preprocess(image, processor, **kwargs): + if isinstance(image, PIL.Image.Image): + pass + elif isinstance(image, np.ndarray): + image = PIL.Image.fromarray(image) + elif isinstance(image, torch.Tensor): + image = vF.to_pil_image(image) + else: + raise TypeError(f"Image must be of type PIL.Image, np.ndarray, or torch.Tensor, got {type(image)} instead.") + + image = processor.preprocess(image, **kwargs) + return image diff --git a/modules/ctrlx/sdxl.py b/modules/ctrlx/sdxl.py new file mode 100644 index 000000000..4b964e7e8 --- /dev/null +++ b/modules/ctrlx/sdxl.py @@ -0,0 +1,299 @@ +from types import MethodType +from typing import Optional +from diffusers.models.attention_processor import Attention +import torch +import torch.nn.functional as F +from .features import feature_injection, normalize, appearance_transfer, get_elem, get_schedule + + +def get_control_config(structure_schedule, appearance_schedule): + s = structure_schedule + a = appearance_schedule + + control_config =\ +f"""control_schedule: + # structure_conv structure_attn appearance_attn conv/attn + encoder: # (num layers) + 0: [[ ], [ ], [ ]] # 2/0 + 1: [[ ], [ ], [{a}, {a} ]] # 2/2 + 2: [[ ], [ ], [{a}, {a} ]] # 2/2 + middle: [[ ], [ ], [ ]] # 2/1 + decoder: + 0: [[{s} ], [{s}, {s}, {s}], [0.0, {a}, {a}]] # 3/3 + 1: [[ ], [ ], [{a}, {a} ]] # 3/3 + 2: [[ ], [ ], [ ]] # 3/0 + +control_target: + - [output_tensor] # structure_conv choices: {{hidden_states, output_tensor}} + - [query, key] # structure_attn choices: {{query, key, value}} + - [before] # appearance_attn choices: {{before, value, after}} + +self_recurrence_schedule: + - [0.1, 0.5, 2] # format: [start, end, num_recurrence]""" + + return control_config + + +def convolution_forward( # From , forward (diffusers==0.28.0) + self, + input_tensor: torch.Tensor, + temb: torch.Tensor, + *args, + **kwargs, +) -> torch.Tensor: + do_structure_control = self.do_control and self.t in self.structure_schedule + + hidden_states = input_tensor + + hidden_states = self.norm1(hidden_states) + hidden_states = self.nonlinearity(hidden_states) + + if self.upsample is not None: + # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984 + if hidden_states.shape[0] >= 64: + input_tensor = input_tensor.contiguous() + hidden_states = hidden_states.contiguous() + input_tensor = self.upsample(input_tensor) + hidden_states = self.upsample(hidden_states) + elif self.downsample is not None: + input_tensor = self.downsample(input_tensor) + hidden_states = self.downsample(hidden_states) + + hidden_states = self.conv1(hidden_states) + + if self.time_emb_proj is not None: + if not self.skip_time_act: + temb = self.nonlinearity(temb) + temb = self.time_emb_proj(temb)[:, :, None, None] + + if self.time_embedding_norm == "default": + if temb is not None: + hidden_states = hidden_states + temb + hidden_states = self.norm2(hidden_states) + elif self.time_embedding_norm == "scale_shift": + if temb is None: + raise ValueError( + f" `temb` should not be None when `time_embedding_norm` is {self.time_embedding_norm}" + ) + time_scale, time_shift = torch.chunk(temb, 2, dim=1) + hidden_states = self.norm2(hidden_states) + hidden_states = hidden_states * (1 + time_scale) + time_shift + else: + hidden_states = self.norm2(hidden_states) + + hidden_states = self.nonlinearity(hidden_states) + + hidden_states = self.dropout(hidden_states) + hidden_states = self.conv2(hidden_states) + + # Feature injection and AdaIN (hidden_states) + if do_structure_control and "hidden_states" in self.structure_target: + hidden_states = feature_injection(hidden_states, batch_order=self.batch_order) + + if self.conv_shortcut is not None: + input_tensor = self.conv_shortcut(input_tensor) + + output_tensor = (input_tensor + hidden_states) / self.output_scale_factor + + # Feature injection and AdaIN (output_tensor) + if do_structure_control and "output_tensor" in self.structure_target: + output_tensor = feature_injection(output_tensor, batch_order=self.batch_order) + + return output_tensor + + +class AttnProcessor2_0: # From (diffusers==0.28.0) + + def __init__(self): + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") + + def __call__( + self, + attn: Attention, + hidden_states: torch.FloatTensor, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + temb: Optional[torch.FloatTensor] = None, + *args, + **kwargs, + ) -> torch.FloatTensor: + do_structure_control = attn.do_control and attn.t in attn.structure_schedule + do_appearance_control = attn.do_control and attn.t in attn.appearance_schedule + + residual = hidden_states + if attn.spatial_norm is not None: + hidden_states = attn.spatial_norm(hidden_states, temb) + + input_ndim = hidden_states.ndim + + if input_ndim == 4: + batch_size, channel, height, width = hidden_states.shape + hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) + + batch_size, sequence_length, _ = ( + hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape + ) + + if attention_mask is not None: + attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) + # scaled_dot_product_attention expects attention_mask shape to be + # (batch, heads, source_length, target_length) + attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) + + if attn.group_norm is not None: + hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) + + no_encoder_hidden_states = encoder_hidden_states is None + if no_encoder_hidden_states: + encoder_hidden_states = hidden_states + elif attn.norm_cross: + encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) + + if do_appearance_control: # Assume we only have this for self attention + hidden_states_normed = normalize(hidden_states, dim=-2) # B H D C + encoder_hidden_states_normed = normalize(encoder_hidden_states, dim=-2) + + query_normed = attn.to_q(hidden_states_normed) + key_normed = attn.to_k(encoder_hidden_states_normed) + + inner_dim = key_normed.shape[-1] + head_dim = inner_dim // attn.heads + query_normed = query_normed.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + key_normed = key_normed.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + # Match query and key injection with structure injection (if injection is happening this layer) + if do_structure_control: + if "query" in attn.structure_target: + query_normed = feature_injection(query_normed, batch_order=attn.batch_order) + if "key" in attn.structure_target: + key_normed = feature_injection(key_normed, batch_order=attn.batch_order) + + # Appearance transfer (before) + if do_appearance_control and "before" in attn.appearance_target: + hidden_states = hidden_states.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + hidden_states = appearance_transfer(hidden_states, query_normed, key_normed, batch_order=attn.batch_order) + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + + if no_encoder_hidden_states: + encoder_hidden_states = hidden_states + elif attn.norm_cross: + encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) + + query = attn.to_q(hidden_states) + + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) + + inner_dim = key.shape[-1] + head_dim = inner_dim // attn.heads + + query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) + + # Feature injection (query, key, and/or value) + if do_structure_control: + if "query" in attn.structure_target: + query = feature_injection(query, batch_order=attn.batch_order) + if "key" in attn.structure_target: + key = feature_injection(key, batch_order=attn.batch_order) + if "value" in attn.structure_target: + value = feature_injection(value, batch_order=attn.batch_order) + + # Appearance transfer (value) + if do_appearance_control and "value" in attn.appearance_target: + value = appearance_transfer(value, query_normed, key_normed, batch_order=attn.batch_order) + + # The output of sdp = (batch, num_heads, seq_len, head_dim) + # TODO: add support for attn.scale when we move to Torch 2.1 + hidden_states = F.scaled_dot_product_attention( + query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False + ) + + # Appearance transfer (after) + if do_appearance_control and "after" in attn.appearance_target: + hidden_states = appearance_transfer(hidden_states, query_normed, key_normed, batch_order=attn.batch_order) + + hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + hidden_states = hidden_states.to(query.dtype) + + # Linear projection + hidden_states = attn.to_out[0](hidden_states, *args) + # Dropout + hidden_states = attn.to_out[1](hidden_states) + + if input_ndim == 4: + hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) + + if attn.residual_connection: + hidden_states = hidden_states + residual + + hidden_states = hidden_states / attn.rescale_output_factor + + return hidden_states + + +def register_control( + model, + timesteps, + control_schedule, # structure_conv, structure_attn, appearance_attn + control_target = [["output_tensor"], ["query", "key"], ["before"]], +): + # Assume timesteps in reverse order (T -> 0) + for block_type in ["encoder", "decoder", "middle"]: + blocks = { + "encoder": model.unet.down_blocks, + "decoder": model.unet.up_blocks, + "middle": [model.unet.mid_block], + }[block_type] + + control_schedule_block = control_schedule[block_type] + if block_type == "middle": + control_schedule_block = [control_schedule_block] + + for layer in range(len(control_schedule_block)): + # Convolution + num_blocks = len(blocks[layer].resnets) if hasattr(blocks[layer], "resnets") else 0 + for block in range(num_blocks): + convolution = blocks[layer].resnets[block] + convolution.structure_target = control_target[0] + convolution.structure_schedule = get_schedule( + timesteps, get_elem(control_schedule_block[layer][0], block) + ) + convolution.forward = MethodType(convolution_forward, convolution) + + # Self-attention + num_blocks = len(blocks[layer].attentions) if hasattr(blocks[layer], "attentions") else 0 + for block in range(num_blocks): + for transformer_block in blocks[layer].attentions[block].transformer_blocks: + attention = transformer_block.attn1 + attention.structure_target = control_target[1] + attention.structure_schedule = get_schedule( + timesteps, get_elem(control_schedule_block[layer][1], block) + ) + attention.appearance_target = control_target[2] + attention.appearance_schedule = get_schedule( + timesteps, get_elem(control_schedule_block[layer][2], block) + ) + attention.processor = AttnProcessor2_0() + + +def register_attr(model, t, do_control, batch_order): + for layer_type in ["encoder", "decoder", "middle"]: + blocks = {"encoder": model.unet.down_blocks, "decoder": model.unet.up_blocks, + "middle": [model.unet.mid_block]}[layer_type] + for layer in blocks: + # Convolution + for module in layer.resnets: + module.t = t + module.do_control = do_control + module.batch_order = batch_order + # Self-attention + if hasattr(layer, "attentions"): + for block in layer.attentions: + for module in block.transformer_blocks: + module.attn1.t = t + module.attn1.do_control = do_control + module.attn1.batch_order = batch_order diff --git a/modules/ctrlx/utils.py b/modules/ctrlx/utils.py new file mode 100644 index 000000000..660dcd2f2 --- /dev/null +++ b/modules/ctrlx/utils.py @@ -0,0 +1,100 @@ +import random +from os import environ +import numpy as np +import torch + + +JPEG_QUALITY = 100 + + +def seed_everything(seed): + random.seed(seed) + environ["PYTHONHASHSEED"] = str(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + + +def exists(x): + return x is not None + + +def get(x, default): + if exists(x): + return x + return default + + +def get_self_recurrence_schedule(schedule, num_inference_steps): + self_recurrence_schedule = [0] * num_inference_steps + for schedule_current in reversed(schedule): + if schedule_current is None or len(schedule_current) == 0: + continue + [start, end, repeat] = schedule_current + start_i = round(num_inference_steps * start) + end_i = round(num_inference_steps * end) + for i in range(start_i, end_i): + self_recurrence_schedule[i] = repeat + return self_recurrence_schedule + + +def batch_dict_to_tensor(batch_dict, batch_order): + batch_tensor = [] + for batch_type in batch_order: + batch_tensor.append(batch_dict[batch_type]) + batch_tensor = torch.cat(batch_tensor, dim=0) + return batch_tensor + + +def batch_tensor_to_dict(batch_tensor, batch_order): + batch_tensor_chunk = batch_tensor.chunk(len(batch_order)) + batch_dict = {} + for i, batch_type in enumerate(batch_order): + batch_dict[batch_type] = batch_tensor_chunk[i] + return batch_dict + + +def noise_prev(scheduler, timestep, x_0, noise=None): + if scheduler.num_inference_steps is None: + raise ValueError( + "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" + ) + + if noise is None: + noise = torch.randn_like(x_0).to(x_0) + + # From DDIMScheduler step function (hopefully this works) + timestep_i = (scheduler.timesteps == timestep).nonzero(as_tuple=True)[0][0].item() + if timestep_i + 1 >= scheduler.timesteps.shape[0]: # We are at t = 0 (ish) + return x_0 + prev_timestep = scheduler.timesteps[timestep_i + 1:timestep_i + 2] # Make sure t is not 0-dim + + x_t_prev = scheduler.add_noise(x_0, noise, prev_timestep) + return x_t_prev + + +def noise_t2t(scheduler, timestep, timestep_target, x_t, noise=None): + assert timestep_target >= timestep + if noise is None: + noise = torch.randn_like(x_t).to(x_t) + + alphas_cumprod = scheduler.alphas_cumprod.to(device=x_t.device, dtype=x_t.dtype) + + timestep = timestep.to(torch.long) + timestep_target = timestep_target.to(torch.long) + + alpha_prod_t = alphas_cumprod[timestep] + alpha_prod_tt = alphas_cumprod[timestep_target] + alpha_prod = alpha_prod_tt / alpha_prod_t + + sqrt_alpha_prod = (alpha_prod ** 0.5).flatten() + while len(sqrt_alpha_prod.shape) < len(x_t.shape): + sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) + + sqrt_one_minus_alpha_prod = ((1 - alpha_prod) ** 0.5).flatten() + while len(sqrt_one_minus_alpha_prod.shape) < len(x_t.shape): + sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) + + x_tt = sqrt_alpha_prod * x_t + sqrt_one_minus_alpha_prod * noise + return x_tt diff --git a/modules/sd_models.py b/modules/sd_models.py index 702223498..0eef568dd 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -1382,7 +1382,7 @@ def get_signature(cls): return signature.parameters -def switch_pipe(cls: diffusers.DiffusionPipeline, pipeline: diffusers.DiffusionPipeline = None, args = {}): +def switch_pipe(cls: diffusers.DiffusionPipeline, pipeline: diffusers.DiffusionPipeline = None, force = False, args = {}): """ args: - cls: can be pipeline class or a string from custom pipelines @@ -1400,7 +1400,7 @@ def switch_pipe(cls: diffusers.DiffusionPipeline, pipeline: diffusers.DiffusionP new_pipe = None signature = get_signature(cls) possible = signature.keys() - if isinstance(pipeline, cls) and args == {}: + if not force and isinstance(pipeline, cls) and args == {}: return pipeline pipe_dict = {} components_used = [] diff --git a/modules/sd_samplers_diffusers.py b/modules/sd_samplers_diffusers.py index 6a4c3cda6..4de463b07 100644 --- a/modules/sd_samplers_diffusers.py +++ b/modules/sd_samplers_diffusers.py @@ -54,8 +54,9 @@ config = { 'DPM++ 2M': { 'thresholding': False, 'sample_max_value': 1.0, 'algorithm_type': "dpmsolver++", 'solver_type': "midpoint", 'lower_order_final': True, 'use_karras_sigmas': False, 'final_sigmas_type': 'zero', 'timestep_spacing': 'linspace', 'solver_order': 2 }, 'DPM++ 3M': { 'thresholding': False, 'sample_max_value': 1.0, 'algorithm_type': "dpmsolver++", 'solver_type': "midpoint", 'lower_order_final': True, 'use_karras_sigmas': False, 'final_sigmas_type': 'zero', 'timestep_spacing': 'linspace', 'solver_order': 3 }, 'DPM SDE': { 'use_karras_sigmas': False, 'noise_sampler_seed': None, 'timestep_spacing': 'linspace', 'steps_offset': 0 }, - 'Euler a': { 'steps_offset': 0, 'rescale_betas_zero_snr': False, 'timestep_spacing': 'linspace' }, 'Euler': { 'steps_offset': 0, 'interpolation_type': "linear", 'use_karras_sigmas': False, 'rescale_betas_zero_snr': False, 'final_sigmas_type': 'zero', 'timestep_spacing': 'linspace' }, + 'Euler a': { 'steps_offset': 0, 'rescale_betas_zero_snr': False, 'timestep_spacing': 'linspace' }, + 'Euler SGM': { 'timestep_spacing': "trailing", 'prediction_type': "sample" }, 'Heun': { 'use_karras_sigmas': False, 'timestep_spacing': 'linspace' }, 'DDPM': { 'variance_type': "fixed_small", 'clip_sample': False, 'thresholding': False, 'clip_sample_range': 1.0, 'sample_max_value': 1.0, 'timestep_spacing': 'linspace', 'rescale_betas_zero_snr': False }, 'KDPM2': { 'steps_offset': 0, 'timestep_spacing': 'linspace' }, @@ -66,7 +67,6 @@ config = { 'DC Solver': { 'beta_start': 0.0001, 'beta_end': 0.02, 'solver_order': 2, 'prediction_type': "epsilon", 'thresholding': False, 'solver_type': 'bh2', 'lower_order_final': True, 'dc_order': 2, 'disable_corrector': [0] }, 'LCM': { 'beta_start': 0.00085, 'beta_end': 0.012, 'beta_schedule': "scaled_linear", 'set_alpha_to_one': True, 'rescale_betas_zero_snr': False, 'thresholding': False, 'timestep_spacing': 'linspace' }, 'TCD': { 'set_alpha_to_one': True, 'rescale_betas_zero_snr': False, 'beta_schedule': 'scaled_linear' }, - 'Euler SGM': { 'timestep_spacing': "trailing", 'prediction_type': "sample" }, 'Euler EDM': { }, 'Variational VDM': { 'clip_sample_range': 2.0, }, 'DPM++ 2M EDM': { 'solver_order': 2, 'solver_type': 'midpoint', 'final_sigmas_type': 'zero', 'algorithm_type': 'dpmsolver++' }, @@ -87,6 +87,7 @@ samplers_data_diffusers = [ sd_samplers_common.SamplerData('Heun', lambda model: DiffusionSampler('Heun', HeunDiscreteScheduler, model), [], {}), sd_samplers_common.SamplerData('Euler', lambda model: DiffusionSampler('Euler', EulerDiscreteScheduler, model), [], {}), sd_samplers_common.SamplerData('Euler a', lambda model: DiffusionSampler('Euler a', EulerAncestralDiscreteScheduler, model), [], {}), + sd_samplers_common.SamplerData('Euler SGM', lambda model: DiffusionSampler('Euler SGM', EulerDiscreteScheduler, model), [], {}), sd_samplers_common.SamplerData('DPM++', lambda model: DiffusionSampler('DPM++', DPMSolverSinglestepScheduler, model), [], {}), sd_samplers_common.SamplerData('DPM++ 1S', lambda model: DiffusionSampler('DPM++ 1S', DPMSolverMultistepScheduler, model), [], {}), sd_samplers_common.SamplerData('DPM++ 2M', lambda model: DiffusionSampler('DPM++ 2M', DPMSolverMultistepScheduler, model), [], {}), @@ -101,7 +102,6 @@ samplers_data_diffusers = [ sd_samplers_common.SamplerData('KDPM2 a', lambda model: DiffusionSampler('KDPM2 a', KDPM2AncestralDiscreteScheduler, model), [], {}), sd_samplers_common.SamplerData('DPM++ 2M EDM', lambda model: DiffusionSampler('DPM++ 2M EDM', EDMDPMSolverMultistepScheduler, model), [], {}), sd_samplers_common.SamplerData('Euler EDM', lambda model: DiffusionSampler('Euler EDM', EDMEulerScheduler, model), [], {}), - sd_samplers_common.SamplerData('Euler SGM', lambda model: DiffusionSampler('Euler SGM', EulerDiscreteScheduler, model), [], {}), sd_samplers_common.SamplerData('LCM', lambda model: DiffusionSampler('LCM', LCMScheduler, model), [], {}), sd_samplers_common.SamplerData('TCD', lambda model: DiffusionSampler('TCD', TCDScheduler, model), [], {}), sd_samplers_common.SamplerData('CMSI', lambda model: DiffusionSampler('CMSI', CMStochasticIterativeScheduler, model), [], {}), diff --git a/scripts/ctrlx.py b/scripts/ctrlx.py new file mode 100644 index 000000000..4a24c9560 --- /dev/null +++ b/scripts/ctrlx.py @@ -0,0 +1,92 @@ +# https://github.com/genforce/ctrl-x + +import gradio as gr +from modules import shared, scripts, processing, processing_helpers, sd_models, devices + + +class Script(scripts.Script): + def title(self): + return 'Ctrl-X' + + def show(self, is_img2img): + return shared.native + + def ui(self, _is_img2img): + with gr.Row(): + gr.HTML('  Ctrl-X
') + with gr.Accordion(label='Structure', open=True): + with gr.Row(): + struct_prompt = gr.Textbox(label='Prompt', value='', rows=1) + with gr.Row(): + struct_strength = gr.Slider(label='Strength', value=0.5, minimum=0.0, maximum=1.0, step=0.05) + struct_guidance = gr.Slider(label='Guidance', value=5.0, minimum=0.0, maximum=14.0, step=0.05) + with gr.Row(): + struct_image = gr.Image(label='Image', source='upload', type='pil') + with gr.Accordion(label='Appearance', open=True): + with gr.Row(): + appear_prompt = gr.Textbox(label='Prompt', value='', rows=1) + with gr.Row(): + appear_strength = gr.Slider(label='Strength', value=0.5, minimum=0.0, maximum=1.0, step=0.05) + appear_guidance = gr.Slider(label='Guidance', value=5.0, minimum=0.0, maximum=14.0, step=0.05) + with gr.Row(): + appear_image = gr.Image(label='Image', source='upload', type='pil') + return struct_prompt, struct_strength, struct_guidance, struct_image, appear_prompt, appear_strength, appear_guidance, appear_image + + def run(self, p: processing.StableDiffusionProcessing, struct_prompt, struct_strength, struct_guidance, struct_image, appear_prompt, appear_strength, appear_guidance, appear_image): # pylint: disable=arguments-differ + c = shared.sd_model.__class__.__name__ if shared.sd_loaded else '' + if shared.sd_model_type != 'sdxl': + shared.log.warning(f'Ctrl-X: pipeline={c} required=StableDiffusionXLPipeline') + return None + + import yaml + from diffusers import StableDiffusionXLPipeline + from modules.ctrlx import CtrlXStableDiffusionXLPipeline + from modules.ctrlx.sdxl import get_control_config, register_control + from modules.ctrlx.utils import get_self_recurrence_schedule + + orig_prompt_attention = shared.opts.prompt_attention + shared.opts.data['prompt_attention'] = 'Fixed attention' + shared.sd_model = sd_models.switch_pipe(CtrlXStableDiffusionXLPipeline, shared.sd_model) + + # calculate ctrx+x schedule + if p.sampler_name not in ['DDIM', 'Euler', 'Euler a', 'DPM++ 1S', 'DDPM', 'Euler SGM', 'LCM', 'TCD']: + shared.log.warning(f'Ctrl-X: sampler={p.sampler_name} override="Euler a" supported=[Euler, Euler a, Euler SGM, DDIM, DDPM, , LCM, TCD]') + p.sampler_name = 'Euler a' + processing_helpers.update_sampler(p, shared.sd_model) + shared.sd_model.scheduler.set_timesteps(p.steps, device=devices.device) + timesteps = shared.sd_model.scheduler.timesteps + control_config = get_control_config(structure_schedule=struct_strength, appearance_schedule=appear_strength) + config = yaml.safe_load(control_config) + register_control( + model=shared.sd_model, + timesteps=timesteps, + control_schedule=config['control_schedule'], + control_target=config['control_target'], + ) + + # set args + if struct_image is not None: + p.task_args['structure_prompt'] = struct_prompt + p.task_args['structure_image'] = struct_image + p.task_args['structure_guidance_scale'] = struct_guidance + if appear_image is not None: + p.task_args['appearance_prompt'] = appear_prompt + p.task_args['appearance_image'] = appear_image + p.task_args['appearance_guidance_scale'] = appear_guidance + elif hasattr(p, 'init_images') and p.init_images is not None and len(p.init_images) > 0: + p.task_args['appearance_image'] = p.init_images[0] + p.init_images = None + p.task_args['control_schedule'] = config['control_schedule'] + p.task_args['self_recurrence_schedule'] = get_self_recurrence_schedule(config['self_recurrence_schedule'], p.steps) + is_struct = p.task_args.get('structure_image') is not None + is_appear = p.task_args.get('appearance_image') is not None + shared.log.info(f'Ctrl-X: structure={struct_strength if is_struct else None} appearance={appear_strength if is_appear else None}') + shared.log.debug(f'Ctrl-X: config={control_config} args={p.task_args}') + + # process + processed: processing.Processed = processing.process_images(p) + + # restore and return + shared.opts.data['prompt_attention'] = orig_prompt_attention + shared.sd_model = sd_models.switch_pipe(StableDiffusionXLPipeline, shared.sd_model, force=True) + return processed