diffusers-fast-inpaint / pipeline_fill_sd_xl.py
OzzyGT's picture
OzzyGT HF staff
initial commit
7890545
raw
history blame
21.2 kB
# Copyright 2024 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import List, Optional, Union
import cv2
import PIL.Image
import torch
import torch.nn.functional as F
from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
from diffusers.models import AutoencoderKL, UNet2DConditionModel
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusionMixin
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.utils.torch_utils import randn_tensor
from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
from controlnet_union import ControlNetModel_Union
def latents_to_rgb(latents):
weights = ((60, -60, 25, -70), (60, -5, 15, -50), (60, 10, -5, -35))
weights_tensor = torch.t(
torch.tensor(weights, dtype=latents.dtype).to(latents.device)
)
biases_tensor = torch.tensor((150, 140, 130), dtype=latents.dtype).to(
latents.device
)
rgb_tensor = torch.einsum(
"...lxy,lr -> ...rxy", latents, weights_tensor
) + biases_tensor.unsqueeze(-1).unsqueeze(-1)
image_array = rgb_tensor.clamp(0, 255)[0].byte().cpu().numpy()
image_array = image_array.transpose(1, 2, 0) # Change the order of dimensions
denoised_image = cv2.fastNlMeansDenoisingColored(image_array, None, 10, 10, 7, 21)
blurred_image = cv2.GaussianBlur(denoised_image, (5, 5), 0)
final_image = PIL.Image.fromarray(blurred_image)
width, height = final_image.size
final_image = final_image.resize(
(width * 8, height * 8), PIL.Image.Resampling.LANCZOS
)
return final_image
def retrieve_timesteps(
scheduler,
num_inference_steps: Optional[int] = None,
device: Optional[Union[str, torch.device]] = None,
**kwargs,
):
scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
timesteps = scheduler.timesteps
return timesteps, num_inference_steps
class StableDiffusionXLFillPipeline(DiffusionPipeline, StableDiffusionMixin):
model_cpu_offload_seq = "text_encoder->text_encoder_2->unet->vae"
_optional_components = [
"tokenizer",
"tokenizer_2",
"text_encoder",
"text_encoder_2",
]
def __init__(
self,
vae: AutoencoderKL,
text_encoder: CLIPTextModel,
text_encoder_2: CLIPTextModelWithProjection,
tokenizer: CLIPTokenizer,
tokenizer_2: CLIPTokenizer,
unet: UNet2DConditionModel,
controlnet: ControlNetModel_Union,
scheduler: KarrasDiffusionSchedulers,
force_zeros_for_empty_prompt: bool = True,
):
super().__init__()
self.register_modules(
vae=vae,
text_encoder=text_encoder,
text_encoder_2=text_encoder_2,
tokenizer=tokenizer,
tokenizer_2=tokenizer_2,
unet=unet,
controlnet=controlnet,
scheduler=scheduler,
)
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
self.image_processor = VaeImageProcessor(
vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True
)
self.control_image_processor = VaeImageProcessor(
vae_scale_factor=self.vae_scale_factor,
do_convert_rgb=True,
do_normalize=False,
)
self.register_to_config(
force_zeros_for_empty_prompt=force_zeros_for_empty_prompt
)
def encode_prompt(
self,
prompt: str,
device: Optional[torch.device] = None,
do_classifier_free_guidance: bool = True,
):
device = device or self._execution_device
prompt = [prompt] if isinstance(prompt, str) else prompt
if prompt is not None:
batch_size = len(prompt)
# Define tokenizers and text encoders
tokenizers = (
[self.tokenizer, self.tokenizer_2]
if self.tokenizer is not None
else [self.tokenizer_2]
)
text_encoders = (
[self.text_encoder, self.text_encoder_2]
if self.text_encoder is not None
else [self.text_encoder_2]
)
prompt_2 = prompt
prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
# textual inversion: process multi-vector tokens if necessary
prompt_embeds_list = []
prompts = [prompt, prompt_2]
for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders):
text_inputs = tokenizer(
prompt,
padding="max_length",
max_length=tokenizer.model_max_length,
truncation=True,
return_tensors="pt",
)
text_input_ids = text_inputs.input_ids
prompt_embeds = text_encoder(
text_input_ids.to(device), output_hidden_states=True
)
# We are only ALWAYS interested in the pooled output of the final text encoder
pooled_prompt_embeds = prompt_embeds[0]
prompt_embeds = prompt_embeds.hidden_states[-2]
prompt_embeds_list.append(prompt_embeds)
prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
# get unconditional embeddings for classifier free guidance
zero_out_negative_prompt = True
negative_prompt_embeds = None
negative_pooled_prompt_embeds = None
if do_classifier_free_guidance and zero_out_negative_prompt:
negative_prompt_embeds = torch.zeros_like(prompt_embeds)
negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds)
elif do_classifier_free_guidance and negative_prompt_embeds is None:
negative_prompt = ""
negative_prompt_2 = negative_prompt
# normalize str to list
negative_prompt = (
batch_size * [negative_prompt]
if isinstance(negative_prompt, str)
else negative_prompt
)
negative_prompt_2 = (
batch_size * [negative_prompt_2]
if isinstance(negative_prompt_2, str)
else negative_prompt_2
)
uncond_tokens: List[str]
if prompt is not None and type(prompt) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
f" {type(prompt)}."
)
elif batch_size != len(negative_prompt):
raise ValueError(
f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
" the batch size of `prompt`."
)
else:
uncond_tokens = [negative_prompt, negative_prompt_2]
negative_prompt_embeds_list = []
for negative_prompt, tokenizer, text_encoder in zip(
uncond_tokens, tokenizers, text_encoders
):
max_length = prompt_embeds.shape[1]
uncond_input = tokenizer(
negative_prompt,
padding="max_length",
max_length=max_length,
truncation=True,
return_tensors="pt",
)
negative_prompt_embeds = text_encoder(
uncond_input.input_ids.to(device),
output_hidden_states=True,
)
# We are only ALWAYS interested in the pooled output of the final text encoder
negative_pooled_prompt_embeds = negative_prompt_embeds[0]
negative_prompt_embeds = negative_prompt_embeds.hidden_states[-2]
negative_prompt_embeds_list.append(negative_prompt_embeds)
negative_prompt_embeds = torch.concat(negative_prompt_embeds_list, dim=-1)
prompt_embeds = prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device)
bs_embed, seq_len, _ = prompt_embeds.shape
# duplicate text embeddings for each generation per prompt, using mps friendly method
prompt_embeds = prompt_embeds.repeat(1, 1, 1)
prompt_embeds = prompt_embeds.view(bs_embed * 1, seq_len, -1)
if do_classifier_free_guidance:
# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
seq_len = negative_prompt_embeds.shape[1]
if self.text_encoder_2 is not None:
negative_prompt_embeds = negative_prompt_embeds.to(
dtype=self.text_encoder_2.dtype, device=device
)
else:
negative_prompt_embeds = negative_prompt_embeds.to(
dtype=self.unet.dtype, device=device
)
negative_prompt_embeds = negative_prompt_embeds.repeat(1, 1, 1)
negative_prompt_embeds = negative_prompt_embeds.view(
batch_size * 1, seq_len, -1
)
pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, 1).view(bs_embed * 1, -1)
if do_classifier_free_guidance:
negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(
1, 1
).view(bs_embed * 1, -1)
return (
prompt_embeds,
negative_prompt_embeds,
pooled_prompt_embeds,
negative_pooled_prompt_embeds,
)
def check_inputs(
self,
prompt_embeds,
negative_prompt_embeds,
pooled_prompt_embeds,
negative_pooled_prompt_embeds,
image,
controlnet_conditioning_scale=1.0,
):
if prompt_embeds is None:
raise ValueError(
"Provide `prompt_embeds`. Cannot leave `prompt_embeds` undefined."
)
if negative_prompt_embeds is None:
raise ValueError(
"Provide `negative_prompt_embeds`. Cannot leave `negative_prompt_embeds` undefined."
)
if prompt_embeds.shape != negative_prompt_embeds.shape:
raise ValueError(
"`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
f" {negative_prompt_embeds.shape}."
)
if prompt_embeds is not None and pooled_prompt_embeds is None:
raise ValueError(
"If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`."
)
if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is None:
raise ValueError(
"If `negative_prompt_embeds` are provided, `negative_pooled_prompt_embeds` also have to be passed. Make sure to generate `negative_pooled_prompt_embeds` from the same text encoder that was used to generate `negative_prompt_embeds`."
)
# Check `image`
is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance(
self.controlnet, torch._dynamo.eval_frame.OptimizedModule
)
if (
isinstance(self.controlnet, ControlNetModel_Union)
or is_compiled
and isinstance(self.controlnet._orig_mod, ControlNetModel_Union)
):
if not isinstance(image, PIL.Image.Image):
raise TypeError(
f"image must be passed and has to be a PIL image, but is {type(image)}"
)
else:
assert False
# Check `controlnet_conditioning_scale`
if (
isinstance(self.controlnet, ControlNetModel_Union)
or is_compiled
and isinstance(self.controlnet._orig_mod, ControlNetModel_Union)
):
if not isinstance(controlnet_conditioning_scale, float):
raise TypeError(
"For single controlnet: `controlnet_conditioning_scale` must be type `float`."
)
else:
assert False
def prepare_image(self, image, device, dtype, do_classifier_free_guidance=False):
image = self.control_image_processor.preprocess(image).to(dtype=torch.float32)
image_batch_size = image.shape[0]
image = image.repeat_interleave(image_batch_size, dim=0)
image = image.to(device=device, dtype=dtype)
if do_classifier_free_guidance:
image = torch.cat([image] * 2)
return image
def prepare_latents(
self, batch_size, num_channels_latents, height, width, dtype, device
):
shape = (
batch_size,
num_channels_latents,
int(height) // self.vae_scale_factor,
int(width) // self.vae_scale_factor,
)
latents = randn_tensor(shape, device=device, dtype=dtype)
# scale the initial noise by the standard deviation required by the scheduler
latents = latents * self.scheduler.init_noise_sigma
return latents
@property
def guidance_scale(self):
return self._guidance_scale
# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
# corresponds to doing no classifier free guidance.
@property
def do_classifier_free_guidance(self):
return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None
@property
def num_timesteps(self):
return self._num_timesteps
@torch.no_grad()
def __call__(
self,
prompt_embeds: torch.Tensor,
negative_prompt_embeds: torch.Tensor,
pooled_prompt_embeds: torch.Tensor,
negative_pooled_prompt_embeds: torch.Tensor,
image: PipelineImageInput = None,
num_inference_steps: int = 8,
guidance_scale: float = 1.5,
controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
):
# 1. Check inputs. Raise error if not correct
self.check_inputs(
prompt_embeds,
negative_prompt_embeds,
pooled_prompt_embeds,
negative_pooled_prompt_embeds,
image,
controlnet_conditioning_scale,
)
self._guidance_scale = guidance_scale
# 2. Define call parameters
batch_size = 1
device = self._execution_device
# 4. Prepare image
if isinstance(self.controlnet, ControlNetModel_Union):
image = self.prepare_image(
image=image,
device=device,
dtype=self.controlnet.dtype,
do_classifier_free_guidance=self.do_classifier_free_guidance,
)
height, width = image.shape[-2:]
else:
assert False
# 5. Prepare timesteps
timesteps, num_inference_steps = retrieve_timesteps(
self.scheduler, num_inference_steps, device
)
self._num_timesteps = len(timesteps)
# 6. Prepare latent variables
num_channels_latents = self.unet.config.in_channels
latents = self.prepare_latents(
batch_size,
num_channels_latents,
height,
width,
prompt_embeds.dtype,
device,
)
# 7 Prepare added time ids & embeddings
add_text_embeds = pooled_prompt_embeds
add_time_ids = negative_add_time_ids = torch.tensor(
image.shape[-2:] + torch.Size([0, 0]) + image.shape[-2:]
).unsqueeze(0)
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_time_ids = torch.cat([negative_add_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).repeat(batch_size, 1)
controlnet_image_list = [0, 0, 0, 0, 0, 0, image, 0]
union_control_type = (
torch.Tensor([0, 0, 0, 0, 0, 0, 1, 0])
.to(device, dtype=prompt_embeds.dtype)
.repeat(batch_size * 2, 1)
)
added_cond_kwargs = {
"text_embeds": add_text_embeds,
"time_ids": add_time_ids,
"control_type": union_control_type,
}
controlnet_prompt_embeds = prompt_embeds
controlnet_added_cond_kwargs = added_cond_kwargs
# 8. Denoising loop
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
# expand the latents if we are doing classifier free guidance
latent_model_input = (
torch.cat([latents] * 2)
if self.do_classifier_free_guidance
else latents
)
latent_model_input = self.scheduler.scale_model_input(
latent_model_input, t
)
# controlnet(s) inference
control_model_input = latent_model_input
down_block_res_samples, mid_block_res_sample = self.controlnet(
control_model_input,
t,
encoder_hidden_states=controlnet_prompt_embeds,
controlnet_cond_list=controlnet_image_list,
conditioning_scale=controlnet_conditioning_scale,
guess_mode=False,
added_cond_kwargs=controlnet_added_cond_kwargs,
return_dict=False,
)
# predict the noise residual
noise_pred = self.unet(
latent_model_input,
t,
encoder_hidden_states=prompt_embeds,
timestep_cond=None,
cross_attention_kwargs={},
down_block_additional_residuals=down_block_res_samples,
mid_block_additional_residual=mid_block_res_sample,
added_cond_kwargs=added_cond_kwargs,
return_dict=False,
)[0]
# perform guidance
if self.do_classifier_free_guidance:
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + guidance_scale * (
noise_pred_text - noise_pred_uncond
)
# compute the previous noisy sample x_t -> x_t-1
latents = self.scheduler.step(
noise_pred, t, latents, return_dict=False
)[0]
if i == 2:
prompt_embeds = prompt_embeds[-1:]
add_text_embeds = add_text_embeds[-1:]
add_time_ids = add_time_ids[-1:]
union_control_type = union_control_type[-1:]
added_cond_kwargs = {
"text_embeds": add_text_embeds,
"time_ids": add_time_ids,
"control_type": union_control_type,
}
controlnet_prompt_embeds = prompt_embeds
controlnet_added_cond_kwargs = added_cond_kwargs
image = image[-1:]
controlnet_image_list = [0, 0, 0, 0, 0, 0, image, 0]
self._guidance_scale = 0.0
if i == len(timesteps) - 1 or (
(i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
):
progress_bar.update()
yield latents_to_rgb(latents)
latents = latents / self.vae.config.scaling_factor
image = self.vae.decode(latents, return_dict=False)[0]
image = self.image_processor.postprocess(image)[0]
# Offload all models
self.maybe_free_model_hooks()
yield image