Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-11 07:18:44 +02:00
Code Issues Packages Projects Releases Wiki Activity

major refactoring of modules

Browse Source 
Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
2025-07-03 09:18:38 -04:00
parent 772a5c9ad3
commit c4d9338d2e
214 changed files with 1154 additions and 1153 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/ctrlx/__init__.py
+656
View File
@@ -0,0 +1,656 @@
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,
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,
):
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(
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`: {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:
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)
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
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)
scripts/ctrlx/features.py
+70
View File
@@ -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
scripts/ctrlx/media.py
+21
View File
@@ -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
scripts/ctrlx/sdxl.py
+298
View File
@@ -0,0 +1,298 @@
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 <class 'diffusers.models.resnet.ResnetBlock2D'>, forward (diffusers==0.28.0)
self,
input_tensor: torch.Tensor,
temb: torch.Tensor,
*args, # pylint: disable=unused-argument
**kwargs, # pylint: disable=unused-argument
) -> 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 <class 'diffusers.models.attention_processor.AttnProcessor2_0'> (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__( # pylint: disable=keyword-arg-before-vararg
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)
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
scripts/ctrlx/utils.py
+100
View File
@@ -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