# Code is based on ReNoise https://github.com/garibida/ReNoise-Inversion
import torch
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
from diffusers import (
StableDiffusionXLImg2ImgPipeline,
)
from diffusers.utils.torch_utils import randn_tensor
from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import (
StableDiffusionXLPipelineOutput,
retrieve_timesteps,
PipelineImageInput
)
from src.eunms import Epsilon_Update_Type
def _backward_ddim(x_tm1, alpha_t, alpha_tm1, eps_xt):
"""
let a = alpha_t, b = alpha_{t - 1}
We have a > b,
x_{t} - x_{t - 1} = sqrt(a) ((sqrt(1/b) - sqrt(1/a)) * x_{t-1} + (sqrt(1/a - 1) - sqrt(1/b - 1)) * eps_{t-1})
From https://arxiv.org/pdf/2105.05233.pdf, section F.
"""
a, b = alpha_t, alpha_tm1
sa = a ** 0.5
sb = b ** 0.5
return sa * ((1 / sb) * x_tm1 + ((1 / a - 1) ** 0.5 - (1 / b - 1) ** 0.5) * eps_xt)
class SDXLDDIMPipeline(StableDiffusionXLImg2ImgPipeline):
# @torch.no_grad()
def __call__(
self,
prompt: Union[str, List[str]] = None,
prompt_2: Optional[Union[str, List[str]]] = None,
image: PipelineImageInput = None,
strength: float = 0.3,
num_inversion_steps: int = 50,
timesteps: List[int] = None,
denoising_start: Optional[float] = None,
denoising_end: Optional[float] = None,
guidance_scale: float = 1.0,
negative_prompt: Optional[Union[str, List[str]]] = None,
negative_prompt_2: Optional[Union[str, List[str]]] = None,
num_images_per_prompt: Optional[int] = 1,
eta: float = 0.0,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.FloatTensor] = None,
prompt_embeds: Optional[torch.FloatTensor] = None,
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
ip_adapter_image: Optional[PipelineImageInput] = None,
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,
negative_original_size: Optional[Tuple[int, int]] = None,
negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
negative_target_size: Optional[Tuple[int, int]] = None,
aesthetic_score: float = 6.0,
negative_aesthetic_score: float = 2.5,
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"],
num_inference_steps: int = 50,
inv_hp=None,
**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`",
)
# 1. Check inputs. Raise error if not correct
self.check_inputs(
prompt,
prompt_2,
strength,
num_inversion_steps,
callback_steps,
negative_prompt,
negative_prompt_2,
prompt_embeds,
negative_prompt_embeds,
callback_on_step_end_tensor_inputs,
)
denoising_start_fr = 1.0 - denoising_start
denoising_start = denoising_start
self._guidance_scale = guidance_scale
self._guidance_rescale = guidance_rescale
self._clip_skip = clip_skip
self._cross_attention_kwargs = cross_attention_kwargs
self._denoising_end = denoising_end
self._denoising_start = denoising_start
# 2. Define call parameters
if prompt is not None and isinstance(prompt, str):
batch_size = 1
elif prompt is not None and isinstance(prompt, list):
batch_size = len(prompt)
else:
batch_size = prompt_embeds.shape[0]
device = self._execution_device
# 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
)
(
prompt_embeds,
negative_prompt_embeds,
pooled_prompt_embeds,
negative_pooled_prompt_embeds,
) = self.encode_prompt(
prompt=prompt,
prompt_2=prompt_2,
device=device,
num_images_per_prompt=num_images_per_prompt,
do_classifier_free_guidance=self.do_classifier_free_guidance,
negative_prompt=negative_prompt,
negative_prompt_2=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,
)
# 4. Preprocess image
image = self.image_processor.preprocess(image)
# 5. Prepare timesteps
def denoising_value_valid(dnv):
return isinstance(self.denoising_end, float) and 0 < dnv < 1
timesteps, num_inversion_steps = retrieve_timesteps(self.scheduler, num_inversion_steps, device, timesteps)
timesteps_num_inference_steps, num_inference_steps = retrieve_timesteps(self.scheduler_inference,
num_inference_steps, device, None)
timesteps, num_inversion_steps = self.get_timesteps(
num_inversion_steps,
strength,
device,
denoising_start=self.denoising_start if denoising_value_valid else None,
)
# latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
# add_noise = True if self.denoising_start is None else False
# 6. Prepare latent variables
with torch.no_grad():
latents = self.prepare_latents(
image,
None,
batch_size,
num_images_per_prompt,
prompt_embeds.dtype,
device,
generator,
False,
)
# 7. Prepare extra step kwargs.
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
height, width = latents.shape[-2:]
height = height * self.vae_scale_factor
width = width * self.vae_scale_factor
original_size = original_size or (height, width)
target_size = target_size or (height, width)
# 8. Prepare added time ids & embeddings
if negative_original_size is None:
negative_original_size = original_size
if negative_target_size is None:
negative_target_size = target_size
add_text_embeds = pooled_prompt_embeds
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, add_neg_time_ids = self._get_add_time_ids(
original_size,
crops_coords_top_left,
target_size,
aesthetic_score,
negative_aesthetic_score,
negative_original_size,
negative_crops_coords_top_left,
negative_target_size,
dtype=prompt_embeds.dtype,
text_encoder_projection_dim=text_encoder_projection_dim,
)
add_time_ids = add_time_ids.repeat(batch_size * num_images_per_prompt, 1)
if self.do_classifier_free_guidance:
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)
add_neg_time_ids = add_neg_time_ids.repeat(batch_size * num_images_per_prompt, 1)
add_time_ids = torch.cat([add_neg_time_ids, add_time_ids], dim=0)
prompt_embeds = prompt_embeds.to(device)
add_text_embeds = add_text_embeds.to(device)
add_time_ids = add_time_ids.to(device)
if ip_adapter_image is not None:
image_embeds, negative_image_embeds = self.encode_image(ip_adapter_image, device, num_images_per_prompt)
if self.do_classifier_free_guidance:
image_embeds = torch.cat([negative_image_embeds, image_embeds])
image_embeds = image_embeds.to(device)
# 9. Denoising loop
num_warmup_steps = max(len(timesteps) - num_inversion_steps * self.scheduler.order, 0)
prev_timestep = None
self._num_timesteps = len(timesteps)
self.prev_z = torch.clone(latents)
self.prev_z4 = torch.clone(latents)
self.z_0 = torch.clone(latents)
g_cpu = torch.Generator().manual_seed(7865)
self.noise = randn_tensor(self.z_0.shape, generator=g_cpu, device=self.z_0.device, dtype=self.z_0.dtype)
# Friendly inversion params
timesteps_for = reversed(timesteps)
noise = randn_tensor(latents.shape, generator=g_cpu, device=latents.device, dtype=latents.dtype)
#latents = latents
z_T = latents.clone()
all_latents = [latents.clone()]
with self.progress_bar(total=num_inversion_steps) as progress_bar:
for i, t in enumerate(timesteps_for):
added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
if ip_adapter_image is not None:
added_cond_kwargs["image_embeds"] = image_embeds
z_tp1 = self.inversion_step(latents,
t,
prompt_embeds,
added_cond_kwargs,
prev_timestep=prev_timestep,
inv_hp=inv_hp,
z_0=self.z_0)
prev_timestep = t
latents = z_tp1
all_latents.append(latents.clone())
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)
# call the callback, if provided
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)
image = latents
# Offload all models
self.maybe_free_model_hooks()
return StableDiffusionXLPipelineOutput(images=image), all_latents
def get_timestamp_dist(self, z_0, timesteps):
timesteps = timesteps.to(z_0.device)
if "cuda" in str(z_0.device):
sigma = self.scheduler.sigmas.cuda()[:-1][self.scheduler.timesteps == timesteps]
else:
sigma = self.scheduler.sigmas[:-1][self.scheduler.timesteps == timesteps]
z_0 = z_0.reshape(-1, 1)
def gaussian_pdf(x):
shape = x.shape
x = x.reshape(-1, 1)
all_probs = - 0.5 * torch.pow(((x - z_0) / sigma), 2)
return all_probs.reshape(shape)
return gaussian_pdf
# @torch.no_grad()
def inversion_step(
self,
z_t: torch.tensor,
t: torch.tensor,
prompt_embeds,
added_cond_kwargs,
prev_timestep: Optional[torch.tensor] = None,
inv_hp=None,
z_0=None,
) -> torch.tensor:
n_iters, alpha, lr = inv_hp
latent = z_t
best_latent = None
best_score = torch.inf
curr_dist = self.get_timestamp_dist(z_0, t)
for i in range(n_iters):
latent.requires_grad = True
noise_pred = self.unet_pass(latent, t, prompt_embeds, added_cond_kwargs)
next_latent = self.backward_step(noise_pred, t, z_t, prev_timestep)
f_x = (next_latent - latent).abs() - alpha * curr_dist(next_latent)
score = f_x.mean()
if score < best_score:
best_score = score
best_latent = next_latent.detach()
f_x.sum().backward()
latent = latent - lr * (f_x / latent.grad)
latent.grad = None
latent._grad_fn = None
# if self.cfg.update_epsilon_type != Epsilon_Update_Type.NONE:
# noise_pred = self.unet_pass(best_latent, t, prompt_embeds, added_cond_kwargs)
# self.scheduler.step_and_update_noise(noise_pred, t, best_latent, z_t, return_dict=False,
# update_epsilon_type=self.cfg.update_epsilon_type)
return best_latent
@torch.no_grad()
def unet_pass(self, z_t, t, prompt_embeds, added_cond_kwargs):
latent_model_input = torch.cat([z_t] * 2) if self.do_classifier_free_guidance else z_t
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
return self.unet(
latent_model_input,
t,
encoder_hidden_states=prompt_embeds,
timestep_cond=None,
cross_attention_kwargs=self.cross_attention_kwargs,
added_cond_kwargs=added_cond_kwargs,
return_dict=False,
)[0]
@torch.no_grad()
def backward_step(self, nosie_pred, t, z_t, prev_timestep):
extra_step_kwargs = {}
return self.scheduler.inv_step(nosie_pred, t, z_t, **extra_step_kwargs, return_dict=False)[0].detach()