Spaces:
Running
Running
import importlib | |
import inspect | |
import math | |
from pathlib import Path | |
import re | |
from collections import defaultdict | |
import cv2 | |
import time | |
import k_diffusion | |
import numpy as np | |
import PIL | |
import torch | |
import torch.nn as nn | |
import torch.nn.functional as F | |
from einops import rearrange | |
from .external_k_diffusion import CompVisDenoiser, CompVisVDenoiser | |
from torch import einsum | |
from torch.autograd.function import Function | |
from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback | |
from diffusers import DiffusionPipeline | |
from diffusers.utils import PIL_INTERPOLATION, is_accelerate_available, logging | |
from diffusers.utils.torch_utils import randn_tensor,is_compiled_module,is_torch_version | |
from diffusers.image_processor import VaeImageProcessor,PipelineImageInput | |
from safetensors.torch import load_file | |
from diffusers import ControlNetModel | |
from PIL import Image | |
import torchvision.transforms as transforms | |
from diffusers import StableDiffusionPipeline,StableDiffusionControlNetPipeline,StableDiffusionControlNetImg2ImgPipeline,StableDiffusionImg2ImgPipeline,StableDiffusionInpaintPipeline,StableDiffusionControlNetInpaintPipeline | |
from typing import Any, Callable, Dict, List, Optional, Union | |
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput | |
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer | |
from diffusers import AutoencoderKL, LMSDiscreteScheduler | |
from .u_net_condition_modify import UNet2DConditionModel | |
from diffusers.models.lora import adjust_lora_scale_text_encoder | |
from diffusers.models import AutoencoderKL, ImageProjection,AsymmetricAutoencoderKL | |
from diffusers.schedulers import KarrasDiffusionSchedulers | |
from diffusers.utils import ( | |
USE_PEFT_BACKEND, | |
deprecate, | |
logging, | |
replace_example_docstring, | |
scale_lora_layers, | |
unscale_lora_layers, | |
) | |
from diffusers.loaders import FromSingleFileMixin, LoraLoaderMixin, TextualInversionLoaderMixin | |
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker | |
from diffusers.pipelines.pipeline_utils import DiffusionPipeline | |
from packaging import version | |
from diffusers.configuration_utils import FrozenDict | |
from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel | |
from .ip_adapter import IPAdapterMixin | |
from .t2i_adapter import preprocessing_t2i_adapter,default_height_width | |
from .encoder_prompt_modify import encode_prompt_function | |
from .encode_region_map_function import encode_region_map | |
def get_image_size(image): | |
height, width = None, None | |
if isinstance(image, Image.Image): | |
return image.size | |
elif isinstance(image, np.ndarray): | |
height, width = image.shape[:2] | |
return (width, height) | |
elif torch.is_tensor(image): | |
#RGB image | |
if len(image.shape) == 3: | |
_, height, width = image.shape | |
else: | |
height, width = image.shape | |
return (width, height) | |
else: | |
raise TypeError("The image must be an instance of PIL.Image, numpy.ndarray, or torch.Tensor.") | |
#Get id token of text at present only support for batch_size = 1 because prompt is a string ("For easy to handle") | |
#Class_name is the name of the class for example StableDiffusion | |
def get_id_text(class_name,prompt,max_length,negative_prompt = None,prompt_embeds: Optional[torch.Tensor] = None,negative_prompt_embeds: Optional[torch.Tensor] = None): | |
#Check prompt_embeds is None -> not using prompt as input | |
if prompt_embeds is not None or negative_prompt_embeds is not None : | |
return None,None | |
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 isinstance(class_name, TextualInversionLoaderMixin): | |
prompt = class_name.maybe_convert_prompt(prompt, class_name.tokenizer) | |
text_inputs = class_name.tokenizer( | |
prompt, | |
padding="max_length", | |
max_length=class_name.tokenizer.model_max_length, | |
truncation=True, | |
return_tensors="pt", | |
) | |
text_input_ids = text_inputs.input_ids.detach().cpu().numpy() | |
uncond_tokens: List[str] | |
if negative_prompt is None: | |
uncond_tokens = [""] * batch_size | |
elif 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 isinstance(negative_prompt, str): | |
uncond_tokens = [negative_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 | |
# textual inversion: procecss multi-vector tokens if necessary | |
if isinstance(class_name, TextualInversionLoaderMixin): | |
uncond_tokens = class_name.maybe_convert_prompt(uncond_tokens, class_name.tokenizer) | |
uncond_input = class_name.tokenizer( | |
uncond_tokens, | |
padding="max_length", | |
max_length=max_length, | |
truncation=True, | |
return_tensors="pt", | |
) | |
uncond_input_ids = uncond_input.input_ids.detach().cpu().numpy() | |
if batch_size == 1: | |
return text_input_ids.reshape((1,-1)),uncond_input_ids.reshape((1,-1)) | |
return text_input_ids,uncond_input_ids | |
# from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg | |
def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): | |
""" | |
Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and | |
Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4 | |
""" | |
std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True) | |
std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True) | |
# rescale the results from guidance (fixes overexposure) | |
noise_pred_rescaled = noise_cfg * (std_text / std_cfg) | |
# mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images | |
noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg | |
return noise_cfg | |
def retrieve_timesteps( | |
scheduler, | |
num_inference_steps: Optional[int] = None, | |
device: Optional[Union[str, torch.device]] = None, | |
timesteps: Optional[List[int]] = None, | |
sigmas: Optional[List[float]] = None, | |
**kwargs, | |
): | |
""" | |
Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles | |
custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. | |
Args: | |
scheduler (`SchedulerMixin`): | |
The scheduler to get timesteps from. | |
num_inference_steps (`int`): | |
The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` | |
must be `None`. | |
device (`str` or `torch.device`, *optional*): | |
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. | |
timesteps (`List[int]`, *optional*): | |
Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, | |
`num_inference_steps` and `sigmas` must be `None`. | |
sigmas (`List[float]`, *optional*): | |
Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, | |
`num_inference_steps` and `timesteps` must be `None`. | |
Returns: | |
`Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the | |
second element is the number of inference steps. | |
""" | |
if timesteps is not None and sigmas is not None: | |
raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") | |
if timesteps is not None: | |
accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) | |
if not accepts_timesteps: | |
raise ValueError( | |
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" | |
f" timestep schedules. Please check whether you are using the correct scheduler." | |
) | |
scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) | |
timesteps = scheduler.timesteps | |
num_inference_steps = len(timesteps) | |
elif sigmas is not None: | |
accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) | |
if not accept_sigmas: | |
raise ValueError( | |
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" | |
f" sigmas schedules. Please check whether you are using the correct scheduler." | |
) | |
scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) | |
timesteps = scheduler.timesteps | |
num_inference_steps = len(timesteps) | |
else: | |
scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) | |
timesteps = scheduler.timesteps | |
return timesteps, num_inference_steps | |
class StableDiffusionPipeline_finetune(IPAdapterMixin,StableDiffusionPipeline): | |
def type_output(self,output_type,device,d_type,return_dict,latents,generator): | |
if not output_type == "latent": | |
image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False,generator=generator)[0] | |
image, has_nsfw_concept = self.run_safety_checker(image, device, d_type) | |
else: | |
image = latents | |
has_nsfw_concept = None | |
if has_nsfw_concept is None: | |
do_denormalize = [True] * image.shape[0] | |
else: | |
do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept] | |
image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize) | |
# Offload all models | |
self.maybe_free_model_hooks() | |
if not return_dict: | |
return (image, has_nsfw_concept) | |
return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) | |
def __call__( | |
self, | |
prompt: Union[str, List[str]] = None, | |
height: Optional[int] = None, | |
width: Optional[int] = None, | |
num_inference_steps: int = 50, | |
timesteps: List[int] = None, | |
sigmas: List[float] = None, | |
guidance_scale: float = 7.5, | |
negative_prompt: Optional[Union[str, List[str]]] = None, | |
num_images_per_prompt: Optional[int] = 1, | |
eta: float = 0.0, | |
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, | |
latents: Optional[torch.Tensor] = None, | |
prompt_embeds: Optional[torch.Tensor] = None, | |
negative_prompt_embeds: Optional[torch.Tensor] = None, | |
ip_adapter_image: Optional[PipelineImageInput] = None, | |
ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, | |
output_type: Optional[str] = "pil", | |
return_dict: bool = True, | |
#callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, | |
#callback_steps: int = 1, | |
cross_attention_kwargs: Optional[Dict[str, Any]] = None, | |
guidance_rescale: float = 0.0, | |
clip_skip: Optional[int] = 0, | |
callback_on_step_end: Optional[ | |
Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] | |
] = None, | |
callback_on_step_end_tensor_inputs: List[str] = ["latents"], | |
region_map_state=None, | |
weight_func = lambda w, sigma, qk: w * sigma * qk.std(), | |
latent_processing = 0, | |
image_t2i_adapter : Optional[PipelineImageInput] = None, | |
adapter_conditioning_scale: Union[float, List[float]] = 1.0, | |
adapter_conditioning_factor: float = 1.0, | |
long_encode: int = 0, | |
**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 using `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 using `callback_on_step_end`", | |
) | |
if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): | |
callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs | |
# 0. Default height and width to unet | |
height = height or self.unet.config.sample_size * self.vae_scale_factor | |
width = width or self.unet.config.sample_size * self.vae_scale_factor | |
# to deal with lora scaling and other possible forward hooks | |
# 1. Check inputs. Raise error if not correct | |
self.check_inputs( | |
prompt, | |
height, | |
width, | |
callback_steps, | |
negative_prompt, | |
prompt_embeds, | |
negative_prompt_embeds, | |
ip_adapter_image, | |
ip_adapter_image_embeds, | |
callback_on_step_end_tensor_inputs, | |
) | |
self._guidance_scale = guidance_scale | |
self._guidance_rescale = guidance_rescale | |
self._clip_skip = clip_skip | |
self._cross_attention_kwargs = cross_attention_kwargs | |
self._interrupt = False | |
adapter_state = None | |
if image_t2i_adapter is not None: | |
height, width = default_height_width(self,height, width, image_t2i_adapter) | |
adapter_state = preprocessing_t2i_adapter(self,image_t2i_adapter,width,height,adapter_conditioning_scale,num_images_per_prompt) | |
# 2. Define call parameters | |
if prompt is not None and isinstance(prompt, str): | |
batch_size = 1 | |
elif prompt is not None and isinstance(prompt, list): | |
batch_size = len(prompt) | |
else: | |
batch_size = prompt_embeds.shape[0] | |
device = self._execution_device | |
# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) | |
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` | |
# corresponds to doing no classifier free guidance. | |
do_classifier_free_guidance = guidance_scale > 1.0 | |
# 3. Encode input prompt | |
lora_scale = ( | |
self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None | |
) | |
#print(type(negative_prompt)) | |
#print(type(prompt)) | |
'''if negative_prompt is None: | |
negative_prompt = '' | |
if prompt is None: | |
prompt =''' | |
#text_ids, text_embeddings = self.prompt_parser([negative_prompt, prompt]) | |
#text_embeddings = text_embeddings.to(self.unet.dtype) | |
#print(text_embeddings) | |
#Copy prompt_embed of input for support get token_id | |
prompt_embeds_copy = None | |
negative_prompt_embeds_copy = None | |
if prompt_embeds is not None: | |
prompt_embeds_copy = prompt_embeds.clone().detach() | |
if negative_prompt_embeds is not None: | |
negative_prompt_embeds_copy = negative_prompt_embeds.clone().detach() | |
prompt_embeds, negative_prompt_embeds,text_input_ids = encode_prompt_function( | |
self, | |
prompt, | |
device, | |
num_images_per_prompt, | |
self.do_classifier_free_guidance, | |
negative_prompt, | |
prompt_embeds=prompt_embeds, | |
negative_prompt_embeds=negative_prompt_embeds, | |
lora_scale=lora_scale, | |
clip_skip=self.clip_skip, | |
long_encode = long_encode, | |
) | |
#Get token_id | |
#text_input_ids,uncond_input_ids = get_id_text(self,prompt,max_length = prompt_embeds.shape[1],negative_prompt = negative_prompt,prompt_embeds = prompt_embeds_copy,negative_prompt_embeds = negative_prompt_embeds_copy) | |
# For classifier free guidance, we need to do two forward passes. | |
# Here we concatenate the unconditional and text embeddings into a single batch | |
# to avoid doing two forward passes | |
'''if text_input_ids is not None: | |
text_input_ids = np.concatenate([uncond_input_ids, text_input_ids])''' | |
if self.do_classifier_free_guidance: | |
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) | |
if ip_adapter_image is not None or ip_adapter_image_embeds is not None: | |
image_embeds = self.prepare_ip_adapter_image_embeds( | |
ip_adapter_image, | |
ip_adapter_image_embeds, | |
device, | |
batch_size * num_images_per_prompt, | |
self.do_classifier_free_guidance, | |
) | |
# 4. Prepare timesteps | |
#print(prompt_embeds) | |
timesteps, num_inference_steps = retrieve_timesteps( | |
self.scheduler, num_inference_steps, device, timesteps, sigmas | |
) | |
#4.1 Prepare region | |
region_state = encode_region_map( | |
self, | |
region_map_state, | |
width = width, | |
height = height, | |
num_images_per_prompt = num_images_per_prompt, | |
text_ids=text_input_ids, | |
) | |
if self.cross_attention_kwargs is None: | |
self._cross_attention_kwargs ={} | |
# 5. Prepare latent variables | |
num_channels_latents = self.unet.config.in_channels | |
latents = self.prepare_latents( | |
batch_size * num_images_per_prompt, | |
num_channels_latents, | |
height, | |
width, | |
prompt_embeds.dtype, | |
device, | |
generator, | |
latents, | |
) | |
lst_latent = [] | |
if latent_processing == 1: | |
lst_latent = [self.type_output("pil",device,prompt_embeds.dtype,return_dict,latents,generator).images[0]] | |
# 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline | |
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) | |
# 6.1 Add image embeds for IP-Adapter | |
added_cond_kwargs = ( | |
{"image_embeds": image_embeds} | |
if (ip_adapter_image is not None or ip_adapter_image_embeds is not None) | |
else None | |
) | |
# 6.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) | |
#print(self.scheduler.sigmas) | |
#print(len(self.scheduler.sigmas)) | |
#values, indices = torch.sort(self.scheduler.sigmas, descending=True) | |
#print(self.scheduler.sigmas) | |
# 7. Denoising loop | |
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order | |
self._num_timesteps = len(timesteps) | |
with self.progress_bar(total=num_inference_steps) as progress_bar: | |
#step_x = 0 | |
for i, t in enumerate(timesteps): | |
if self.interrupt: | |
continue | |
# 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) | |
#print(self.scheduler.sigmas[step_x]) | |
region_prompt = { | |
"region_state": region_state, | |
"sigma": self.scheduler.sigmas[i], | |
"weight_func": weight_func, | |
} | |
self._cross_attention_kwargs["region_prompt"] = region_prompt | |
#print(t) | |
#step_x=step_x+1 | |
#tensor_data = {k: torch.Tensor(v) for k, v in encoder_state.items()} | |
# predict the noise residual | |
down_intrablock_additional_residuals = None | |
if adapter_state is not None: | |
if i < int(num_inference_steps * adapter_conditioning_factor): | |
down_intrablock_additional_residuals = [state.clone() for state in adapter_state] | |
else: | |
down_intrablock_additional_residuals = None | |
noise_pred = self.unet( | |
latent_model_input, | |
t, | |
encoder_hidden_states=prompt_embeds, | |
timestep_cond=timestep_cond, | |
cross_attention_kwargs=self.cross_attention_kwargs, | |
down_intrablock_additional_residuals = down_intrablock_additional_residuals, | |
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 + self.guidance_scale * (noise_pred_text - noise_pred_uncond) | |
if self.do_classifier_free_guidance and self.guidance_rescale > 0.0: | |
# Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf | |
noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale) | |
# compute the previous noisy sample x_t -> x_t-1 | |
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] | |
if latent_processing == 1: | |
lst_latent.append(self.type_output("pil",device,prompt_embeds.dtype,return_dict,latents,generator).images[0]) | |
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) | |
# 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) | |
torch.cuda.empty_cache() | |
if latent_processing == 1: | |
if output_type == 'latent': | |
lst_latent.append(self.type_output(output_type,device,prompt_embeds.dtype,return_dict,latents,generator).images[0]) | |
return lst_latent | |
if output_type == 'latent': | |
return [self.type_output("pil",device,prompt_embeds.dtype,return_dict,latents,generator).images[0],self.type_output(output_type,device,prompt_embeds.dtype,return_dict,latents,generator).images[0]] | |
return [self.type_output(output_type,device,prompt_embeds.dtype,return_dict,latents,generator).images[0]] | |
class StableDiffusionControlNetPipeline_finetune(IPAdapterMixin,StableDiffusionControlNetPipeline): | |
def type_output(self,output_type,device,d_type,return_dict,latents,generator): | |
if not output_type == "latent": | |
image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False,generator=generator)[0] | |
image, has_nsfw_concept = self.run_safety_checker(image, device, d_type) | |
else: | |
image = latents | |
has_nsfw_concept = None | |
if has_nsfw_concept is None: | |
do_denormalize = [True] * image.shape[0] | |
else: | |
do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept] | |
image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize) | |
# Offload all models | |
self.maybe_free_model_hooks() | |
if not return_dict: | |
return (image, has_nsfw_concept) | |
return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) | |
def __call__( | |
self, | |
prompt: Union[str, List[str]] = None, | |
image: PipelineImageInput = None, | |
height: Optional[int] = None, | |
width: Optional[int] = None, | |
num_inference_steps: int = 50, | |
timesteps: List[int] = None, | |
sigmas: List[float] = None, | |
guidance_scale: float = 7.5, | |
negative_prompt: Optional[Union[str, List[str]]] = None, | |
num_images_per_prompt: Optional[int] = 1, | |
eta: float = 0.0, | |
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, | |
latents: Optional[torch.Tensor] = None, | |
prompt_embeds: Optional[torch.Tensor] = None, | |
negative_prompt_embeds: Optional[torch.Tensor] = None, | |
ip_adapter_image: Optional[PipelineImageInput] = None, | |
ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, | |
output_type: Optional[str] = "pil", | |
return_dict: bool = True, | |
guidance_rescale: float = 0.0, | |
#callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, | |
#callback_steps: int = 1, | |
cross_attention_kwargs: Optional[Dict[str, Any]] = None, | |
controlnet_conditioning_scale: Union[float, List[float]] = 1.0, | |
guess_mode: bool = False, | |
control_guidance_start: Union[float, List[float]] = 0.0, | |
control_guidance_end: Union[float, List[float]] = 1.0, | |
clip_skip: Optional[int] = 0, | |
callback_on_step_end: Optional[ | |
Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] | |
] = None, | |
callback_on_step_end_tensor_inputs: List[str] = ["latents"], | |
region_map_state=None, | |
weight_func = lambda w, sigma, qk: w * sigma * qk.std(), | |
latent_processing = 0, | |
image_t2i_adapter : Optional[PipelineImageInput] = None, | |
adapter_conditioning_scale: Union[float, List[float]] = 1.0, | |
adapter_conditioning_factor: float = 1.0, | |
long_encode: int = 0, | |
**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 using `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 using `callback_on_step_end`", | |
) | |
if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): | |
callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs | |
controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet | |
# align format for control guidance | |
if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): | |
control_guidance_start = len(control_guidance_end) * [control_guidance_start] | |
elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): | |
control_guidance_end = len(control_guidance_start) * [control_guidance_end] | |
elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list): | |
mult = len(controlnet.nets) if isinstance(controlnet, MultiControlNetModel) else 1 | |
control_guidance_start, control_guidance_end = ( | |
mult * [control_guidance_start], | |
mult * [control_guidance_end], | |
) | |
if height is None: | |
_,height = get_image_size(image) | |
height = int((height // 8)*8) | |
if width is None: | |
width,_ = get_image_size(image) | |
width = int((width // 8)*8) | |
# 1. Check inputs. Raise error if not correct | |
self.check_inputs( | |
prompt, | |
image, | |
callback_steps, | |
negative_prompt, | |
prompt_embeds, | |
negative_prompt_embeds, | |
ip_adapter_image, | |
ip_adapter_image_embeds, | |
controlnet_conditioning_scale, | |
control_guidance_start, | |
control_guidance_end, | |
callback_on_step_end_tensor_inputs, | |
) | |
self._guidance_scale = guidance_scale | |
self._clip_skip = clip_skip | |
self._cross_attention_kwargs = cross_attention_kwargs | |
adapter_state = None | |
if image_t2i_adapter is not None: | |
height, width = default_height_width(self,height, width, image_t2i_adapter) | |
adapter_state = preprocessing_t2i_adapter(self,image_t2i_adapter,width,height,adapter_conditioning_scale,num_images_per_prompt) | |
# 2. Define call parameters | |
if prompt is not None and isinstance(prompt, str): | |
batch_size = 1 | |
elif prompt is not None and isinstance(prompt, list): | |
batch_size = len(prompt) | |
else: | |
batch_size = prompt_embeds.shape[0] | |
device = self._execution_device | |
# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) | |
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` | |
# corresponds to doing no classifier free guidance. | |
#do_classifier_free_guidance = guidance_scale > 1.0 | |
if isinstance(controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float): | |
controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(controlnet.nets) | |
global_pool_conditions = ( | |
controlnet.config.global_pool_conditions | |
if isinstance(controlnet, ControlNetModel) | |
else controlnet.nets[0].config.global_pool_conditions | |
) | |
guess_mode = guess_mode or global_pool_conditions | |
# 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 | |
) | |
#text_ids, text_embeddings = self.prompt_parser([negative_prompt, prompt]) | |
#text_embeddings = text_embeddings.to(self.unet.dtype) | |
#Copy input prompt_embeds and negative_prompt_embeds | |
prompt_embeds_copy = None | |
negative_prompt_embeds_copy = None | |
if prompt_embeds is not None: | |
prompt_embeds_copy = prompt_embeds.clone().detach() | |
if negative_prompt_embeds is not None: | |
negative_prompt_embeds_copy = negative_prompt_embeds.clone().detach() | |
prompt_embeds, negative_prompt_embeds,text_input_ids = encode_prompt_function( | |
self, | |
prompt, | |
device, | |
num_images_per_prompt, | |
self.do_classifier_free_guidance, | |
negative_prompt, | |
prompt_embeds=prompt_embeds, | |
negative_prompt_embeds=negative_prompt_embeds, | |
lora_scale=text_encoder_lora_scale, | |
clip_skip=self.clip_skip, | |
long_encode = long_encode, | |
) | |
#Get token_id | |
#text_input_ids,uncond_input_ids = get_id_text(self,prompt,max_length = prompt_embeds.shape[1],negative_prompt = negative_prompt,prompt_embeds = prompt_embeds_copy,negative_prompt_embeds = negative_prompt_embeds_copy) | |
# For classifier free guidance, we need to do two forward passes. | |
# Here we concatenate the unconditional and text embeddings into a single batch | |
# to avoid doing two forward passes | |
'''if text_input_ids is not None: | |
text_input_ids = np.concatenate([uncond_input_ids, text_input_ids])''' | |
# For classifier free guidance, we need to do two forward passes. | |
# Here we concatenate the unconditional and text embeddings into a single batch | |
# to avoid doing two forward passes | |
if self.do_classifier_free_guidance: | |
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) | |
if ip_adapter_image is not None or ip_adapter_image_embeds is not None: | |
image_embeds = self.prepare_ip_adapter_image_embeds( | |
ip_adapter_image, | |
ip_adapter_image_embeds, | |
device, | |
batch_size * num_images_per_prompt, | |
self.do_classifier_free_guidance, | |
) | |
#if height is None and width is None: | |
#height, width = image.shape[-2:] | |
# 4. Prepare image | |
if isinstance(controlnet, ControlNetModel): | |
image = self.prepare_image( | |
image=image, | |
width=width, | |
height=height, | |
batch_size=batch_size * num_images_per_prompt, | |
num_images_per_prompt=num_images_per_prompt, | |
device=device, | |
dtype=controlnet.dtype, | |
do_classifier_free_guidance=self.do_classifier_free_guidance, | |
guess_mode=guess_mode, | |
) | |
elif isinstance(controlnet, MultiControlNetModel): | |
images = [] | |
# Nested lists as ControlNet condition | |
if isinstance(image[0], list): | |
# Transpose the nested image list | |
image = [list(t) for t in zip(*image)] | |
for image_ in image: | |
image_ = self.prepare_image( | |
image=image_, | |
width=width, | |
height=height, | |
batch_size=batch_size * num_images_per_prompt, | |
num_images_per_prompt=num_images_per_prompt, | |
device=device, | |
dtype=controlnet.dtype, | |
do_classifier_free_guidance=self.do_classifier_free_guidance, | |
guess_mode=guess_mode, | |
) | |
images.append(image_) | |
image = images | |
height, width = image[0].shape[-2:] | |
else: | |
assert False | |
# 5. Prepare timesteps | |
timesteps, num_inference_steps = retrieve_timesteps( | |
self.scheduler, num_inference_steps, device, timesteps, sigmas | |
) | |
self._num_timesteps = len(timesteps) | |
# 6. Prepare latent variables | |
region_state = encode_region_map( | |
self, | |
region_map_state, | |
width = width, | |
height = height, | |
num_images_per_prompt = num_images_per_prompt, | |
text_ids=text_input_ids, | |
) | |
if self.cross_attention_kwargs is None: | |
self._cross_attention_kwargs ={} | |
num_channels_latents = self.unet.config.in_channels | |
latents = self.prepare_latents( | |
batch_size * num_images_per_prompt, | |
num_channels_latents, | |
height, | |
width, | |
prompt_embeds.dtype, | |
device, | |
generator, | |
latents, | |
) | |
# 6.5 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) | |
lst_latent = [] | |
if latent_processing == 1: | |
lst_latent = [self.type_output("pil",device,prompt_embeds.dtype,return_dict,latents,generator).images[0]] | |
# 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline | |
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) | |
# 7.1 Add image embeds for IP-Adapter | |
added_cond_kwargs = ( | |
{"image_embeds": image_embeds} | |
if ip_adapter_image is not None or ip_adapter_image_embeds is not None | |
else None | |
) | |
# 7.2 Create tensor stating which controlnets to keep | |
controlnet_keep = [] | |
for i in range(len(timesteps)): | |
keeps = [ | |
1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e) | |
for s, e in zip(control_guidance_start, control_guidance_end) | |
] | |
controlnet_keep.append(keeps[0] if isinstance(controlnet, ControlNetModel) else keeps) | |
# 8. Denoising loop | |
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order | |
is_unet_compiled = is_compiled_module(self.unet) | |
is_controlnet_compiled = is_compiled_module(self.controlnet) | |
is_torch_higher_equal_2_1 = is_torch_version(">=", "2.1") | |
with self.progress_bar(total=num_inference_steps) as progress_bar: | |
#step_x = 0 | |
for i, t in enumerate(timesteps): | |
# Relevant thread: | |
# https://dev-discuss.pytorch.org/t/cudagraphs-in-pytorch-2-0/1428 | |
if (is_unet_compiled and is_controlnet_compiled) and is_torch_higher_equal_2_1: | |
torch._inductor.cudagraph_mark_step_begin() | |
# 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 | |
if guess_mode and self.do_classifier_free_guidance: | |
# Infer ControlNet only for the conditional batch. | |
control_model_input = latents | |
control_model_input = self.scheduler.scale_model_input(control_model_input, t) | |
controlnet_prompt_embeds = prompt_embeds.chunk(2)[1] | |
else: | |
control_model_input = latent_model_input | |
controlnet_prompt_embeds = prompt_embeds | |
if isinstance(controlnet_keep[i], list): | |
cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])] | |
else: | |
controlnet_cond_scale = controlnet_conditioning_scale | |
if isinstance(controlnet_cond_scale, list): | |
controlnet_cond_scale = controlnet_cond_scale[0] | |
cond_scale = controlnet_cond_scale * controlnet_keep[i] | |
down_block_res_samples, mid_block_res_sample = self.controlnet( | |
control_model_input, | |
t, | |
encoder_hidden_states=controlnet_prompt_embeds, | |
controlnet_cond=image, | |
conditioning_scale=cond_scale, | |
guess_mode=guess_mode, | |
return_dict=False, | |
) | |
if guess_mode and self.do_classifier_free_guidance: | |
# Infered ControlNet only for the conditional batch. | |
# To apply the output of ControlNet to both the unconditional and conditional batches, | |
# add 0 to the unconditional batch to keep it unchanged. | |
down_block_res_samples = [torch.cat([torch.zeros_like(d), d]) for d in down_block_res_samples] | |
mid_block_res_sample = torch.cat([torch.zeros_like(mid_block_res_sample), mid_block_res_sample]) | |
region_prompt = { | |
"region_state": region_state, | |
"sigma": self.scheduler.sigmas[i], | |
"weight_func": weight_func, | |
} | |
self._cross_attention_kwargs["region_prompt"] = region_prompt | |
#print(t) | |
#step_x=step_x+1 | |
down_intrablock_additional_residuals = None | |
if adapter_state is not None: | |
if i < int(num_inference_steps * adapter_conditioning_factor): | |
down_intrablock_additional_residuals = [state.clone() for state in adapter_state] | |
else: | |
down_intrablock_additional_residuals = None | |
# predict the noise residual | |
noise_pred = self.unet( | |
latent_model_input, | |
t, | |
encoder_hidden_states=prompt_embeds, | |
timestep_cond=timestep_cond, | |
cross_attention_kwargs=self.cross_attention_kwargs, | |
down_block_additional_residuals=down_block_res_samples, | |
mid_block_additional_residual=mid_block_res_sample, | |
down_intrablock_additional_residuals = down_intrablock_additional_residuals, | |
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 + self.guidance_scale * (noise_pred_text - noise_pred_uncond) | |
if self.do_classifier_free_guidance and guidance_rescale > 0.0: | |
# Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf | |
noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) | |
# compute the previous noisy sample x_t -> x_t-1 | |
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] | |
if latent_processing == 1: | |
lst_latent.append(self.type_output("pil",device,prompt_embeds.dtype,return_dict,latents,generator).images[0]) | |
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) | |
# 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) | |
# If we do sequential model offloading, let's offload unet and controlnet | |
# manually for max memory savings | |
if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: | |
self.unet.to("cpu") | |
self.controlnet.to("cpu") | |
torch.cuda.empty_cache() | |
if latent_processing == 1: | |
if output_type == 'latent': | |
lst_latent.append(self.type_output(output_type,device,prompt_embeds.dtype,return_dict,latents,generator).images[0]) | |
return lst_latent | |
if output_type == 'latent': | |
return [self.type_output("pil",device,prompt_embeds.dtype,return_dict,latents,generator).images[0],self.type_output(output_type,device,prompt_embeds.dtype,return_dict,latents,generator).images[0]] | |
return [self.type_output(output_type,device,prompt_embeds.dtype,return_dict,latents,generator).images[0]] | |
class StableDiffusionControlNetImg2ImgPipeline_finetune(IPAdapterMixin,StableDiffusionControlNetImg2ImgPipeline): | |
def type_output(self,output_type,device,d_type,return_dict,latents,generator): | |
if not output_type == "latent": | |
image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False,generator=generator)[0] | |
image, has_nsfw_concept = self.run_safety_checker(image, device, d_type) | |
else: | |
image = latents | |
has_nsfw_concept = None | |
if has_nsfw_concept is None: | |
do_denormalize = [True] * image.shape[0] | |
else: | |
do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept] | |
image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize) | |
# Offload all models | |
self.maybe_free_model_hooks() | |
if not return_dict: | |
return (image, has_nsfw_concept) | |
return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) | |
def __call__( | |
self, | |
prompt: Union[str, List[str]] = None, | |
image: PipelineImageInput = None, | |
control_image: PipelineImageInput = None, | |
height: Optional[int] = None, | |
width: Optional[int] = None, | |
strength: float = 0.8, | |
num_inference_steps: int = 50, | |
guidance_scale: float = 7.5, | |
negative_prompt: Optional[Union[str, List[str]]] = None, | |
num_images_per_prompt: Optional[int] = 1, | |
eta: float = 0.0, | |
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, | |
latents: Optional[torch.Tensor] = None, | |
prompt_embeds: Optional[torch.Tensor] = None, | |
negative_prompt_embeds: Optional[torch.Tensor] = None, | |
ip_adapter_image: Optional[PipelineImageInput] = None, | |
ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, | |
output_type: Optional[str] = "pil", | |
return_dict: bool = True, | |
guidance_rescale: float = 0.0, | |
#callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, | |
#callback_steps: int = 1, | |
cross_attention_kwargs: Optional[Dict[str, Any]] = None, | |
controlnet_conditioning_scale: Union[float, List[float]] = 0.8, | |
guess_mode: bool = False, | |
control_guidance_start: Union[float, List[float]] = 0.0, | |
control_guidance_end: Union[float, List[float]] = 1.0, | |
clip_skip: Optional[int] = 0, | |
callback_on_step_end: Optional[ | |
Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] | |
] = None, | |
callback_on_step_end_tensor_inputs: List[str] = ["latents"], | |
region_map_state=None, | |
weight_func = lambda w, sigma, qk: w * sigma * qk.std(), | |
latent_processing = 0, | |
image_t2i_adapter : Optional[PipelineImageInput] = None, | |
adapter_conditioning_scale: Union[float, List[float]] = 1.0, | |
adapter_conditioning_factor: float = 1.0, | |
long_encode: int = 0, | |
**kwargs, | |
): | |
init_step = num_inference_steps | |
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 using `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 using `callback_on_step_end`", | |
) | |
if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): | |
callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs | |
controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet | |
if height is None: | |
_,height = get_image_size(image) | |
height = int((height // 8)*8) | |
if width is None: | |
width,_ = get_image_size(image) | |
width = int((width // 8)*8) | |
# align format for control guidance | |
if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): | |
control_guidance_start = len(control_guidance_end) * [control_guidance_start] | |
elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): | |
control_guidance_end = len(control_guidance_start) * [control_guidance_end] | |
elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list): | |
mult = len(controlnet.nets) if isinstance(controlnet, MultiControlNetModel) else 1 | |
control_guidance_start, control_guidance_end = ( | |
mult * [control_guidance_start], | |
mult * [control_guidance_end], | |
) | |
# 1. Check inputs. Raise error if not correct | |
self.check_inputs( | |
prompt, | |
control_image, | |
callback_steps, | |
negative_prompt, | |
prompt_embeds, | |
negative_prompt_embeds, | |
ip_adapter_image, | |
ip_adapter_image_embeds, | |
controlnet_conditioning_scale, | |
control_guidance_start, | |
control_guidance_end, | |
callback_on_step_end_tensor_inputs, | |
) | |
self._guidance_scale = guidance_scale | |
self._clip_skip = clip_skip | |
self._cross_attention_kwargs = cross_attention_kwargs | |
adapter_state = None | |
if image_t2i_adapter is not None: | |
height, width = default_height_width(self,height, width, image_t2i_adapter) | |
adapter_state = preprocessing_t2i_adapter(self,image_t2i_adapter,width,height,adapter_conditioning_scale,num_images_per_prompt) | |
#self.prompt_parser = FrozenCLIPEmbedderWithCustomWords(self.tokenizer, self.text_encoder,clip_skip+1) | |
# 2. Define call parameters | |
if prompt is not None and isinstance(prompt, str): | |
batch_size = 1 | |
elif prompt is not None and isinstance(prompt, list): | |
batch_size = len(prompt) | |
else: | |
batch_size = prompt_embeds.shape[0] | |
device = self._execution_device | |
# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) | |
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` | |
# corresponds to doing no classifier free guidance. | |
#do_classifier_free_guidance = guidance_scale > 1.0 | |
if isinstance(controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float): | |
controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(controlnet.nets) | |
global_pool_conditions = ( | |
controlnet.config.global_pool_conditions | |
if isinstance(controlnet, ControlNetModel) | |
else controlnet.nets[0].config.global_pool_conditions | |
) | |
guess_mode = guess_mode or global_pool_conditions | |
# 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 | |
) | |
#text_ids, text_embeddings = self.prompt_parser([negative_prompt, prompt]) | |
#text_embeddings = text_embeddings.to(self.unet.dtype) | |
#Copy input prompt_embeds and negative_prompt_embeds | |
prompt_embeds_copy = None | |
negative_prompt_embeds_copy = None | |
if prompt_embeds is not None: | |
prompt_embeds_copy = prompt_embeds.clone().detach() | |
if negative_prompt_embeds is not None: | |
negative_prompt_embeds_copy = negative_prompt_embeds.clone().detach() | |
prompt_embeds, negative_prompt_embeds,text_input_ids = encode_prompt_function( | |
self, | |
prompt, | |
device, | |
num_images_per_prompt, | |
self.do_classifier_free_guidance, | |
negative_prompt, | |
prompt_embeds=prompt_embeds, | |
negative_prompt_embeds=negative_prompt_embeds, | |
lora_scale=text_encoder_lora_scale, | |
clip_skip=self.clip_skip, | |
long_encode = long_encode, | |
) | |
#Get token_id | |
#text_input_ids,uncond_input_ids = get_id_text(self,prompt,max_length = prompt_embeds.shape[1],negative_prompt = negative_prompt,prompt_embeds = prompt_embeds_copy,negative_prompt_embeds = negative_prompt_embeds_copy) | |
# For classifier free guidance, we need to do two forward passes. | |
# Here we concatenate the unconditional and text embeddings into a single batch | |
# to avoid doing two forward passes | |
'''if text_input_ids is not None: | |
text_input_ids = np.concatenate([uncond_input_ids, text_input_ids])''' | |
# For classifier free guidance, we need to do two forward passes. | |
# Here we concatenate the unconditional and text embeddings into a single batch | |
# to avoid doing two forward passes | |
if self.do_classifier_free_guidance: | |
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) | |
if ip_adapter_image is not None or ip_adapter_image_embeds is not None: | |
image_embeds = self.prepare_ip_adapter_image_embeds( | |
ip_adapter_image, | |
ip_adapter_image_embeds, | |
device, | |
batch_size * num_images_per_prompt, | |
self.do_classifier_free_guidance, | |
) | |
# 4. Prepare image | |
image = self.image_processor.preprocess(image, height=height, width=width).to(dtype=torch.float32) | |
# 5. Prepare controlnet_conditioning_image | |
if isinstance(controlnet, ControlNetModel): | |
control_image = self.prepare_control_image( | |
image=control_image, | |
width=width, | |
height=height, | |
batch_size=batch_size * num_images_per_prompt, | |
num_images_per_prompt=num_images_per_prompt, | |
device=device, | |
dtype=controlnet.dtype, | |
do_classifier_free_guidance=self.do_classifier_free_guidance, | |
guess_mode=guess_mode, | |
) | |
elif isinstance(controlnet, MultiControlNetModel): | |
control_images = [] | |
# Nested lists as ControlNet condition | |
if isinstance(image[0], list): | |
# Transpose the nested image list | |
image = [list(t) for t in zip(*image)] | |
for control_image_ in control_image: | |
control_image_ = self.prepare_control_image( | |
image=control_image_, | |
width=width, | |
height=height, | |
batch_size=batch_size * num_images_per_prompt, | |
num_images_per_prompt=num_images_per_prompt, | |
device=device, | |
dtype=controlnet.dtype, | |
do_classifier_free_guidance=self.do_classifier_free_guidance, | |
guess_mode=guess_mode, | |
) | |
control_images.append(control_image_) | |
control_image = control_images | |
else: | |
assert False | |
# 5. Prepare timesteps | |
region_state = encode_region_map( | |
self, | |
region_map_state, | |
width = width, | |
height = height, | |
num_images_per_prompt = num_images_per_prompt, | |
text_ids=text_input_ids, | |
) | |
if self.cross_attention_kwargs is None: | |
self._cross_attention_kwargs ={} | |
self.scheduler.set_timesteps(num_inference_steps, device=device) | |
timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) | |
latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) | |
self._num_timesteps = len(timesteps) | |
# 6. Prepare latent variables | |
if latents is None: | |
latents = self.prepare_latents( | |
image, | |
latent_timestep, | |
batch_size, | |
num_images_per_prompt, | |
prompt_embeds.dtype, | |
device, | |
generator, | |
) | |
lst_latent = [] | |
if latent_processing == 1: | |
lst_latent = [self.type_output("pil",device,prompt_embeds.dtype,return_dict,latents,generator).images[0]] | |
# 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline | |
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) | |
# 7.1 Add image embeds for IP-Adapter | |
added_cond_kwargs = ( | |
{"image_embeds": image_embeds} | |
if ip_adapter_image is not None or ip_adapter_image_embeds is not None | |
else None | |
) | |
# 7.2 Create tensor stating which controlnets to keep | |
controlnet_keep = [] | |
for i in range(len(timesteps)): | |
keeps = [ | |
1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e) | |
for s, e in zip(control_guidance_start, control_guidance_end) | |
] | |
controlnet_keep.append(keeps[0] if isinstance(controlnet, ControlNetModel) else keeps) | |
sigmas = self.scheduler.sigmas[init_step-len(timesteps):] | |
# 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: | |
#step_x = 0 | |
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 | |
if guess_mode and self.do_classifier_free_guidance: | |
# Infer ControlNet only for the conditional batch. | |
control_model_input = latents | |
control_model_input = self.scheduler.scale_model_input(control_model_input, t) | |
controlnet_prompt_embeds = prompt_embeds.chunk(2)[1] | |
else: | |
control_model_input = latent_model_input | |
controlnet_prompt_embeds = prompt_embeds | |
if isinstance(controlnet_keep[i], list): | |
cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])] | |
else: | |
controlnet_cond_scale = controlnet_conditioning_scale | |
if isinstance(controlnet_cond_scale, list): | |
controlnet_cond_scale = controlnet_cond_scale[0] | |
cond_scale = controlnet_cond_scale * controlnet_keep[i] | |
down_block_res_samples, mid_block_res_sample = self.controlnet( | |
control_model_input, | |
t, | |
encoder_hidden_states=controlnet_prompt_embeds, | |
controlnet_cond=control_image, | |
conditioning_scale=cond_scale, | |
guess_mode=guess_mode, | |
return_dict=False, | |
) | |
if guess_mode and self.do_classifier_free_guidance: | |
# Infered ControlNet only for the conditional batch. | |
# To apply the output of ControlNet to both the unconditional and conditional batches, | |
# add 0 to the unconditional batch to keep it unchanged. | |
down_block_res_samples = [torch.cat([torch.zeros_like(d), d]) for d in down_block_res_samples] | |
mid_block_res_sample = torch.cat([torch.zeros_like(mid_block_res_sample), mid_block_res_sample]) | |
region_prompt = { | |
"region_state": region_state, | |
"sigma": self.scheduler.sigmas[i], | |
"weight_func": weight_func, | |
} | |
self._cross_attention_kwargs["region_prompt"] = region_prompt | |
#print(t) | |
#step_x=step_x+1 | |
down_intrablock_additional_residuals = None | |
if adapter_state is not None: | |
if i < int(num_inference_steps * adapter_conditioning_factor): | |
down_intrablock_additional_residuals = [state.clone() for state in adapter_state] | |
else: | |
down_intrablock_additional_residuals = None | |
# predict the noise residual | |
# predict the noise residual | |
noise_pred = self.unet( | |
latent_model_input, | |
t, | |
encoder_hidden_states=prompt_embeds, | |
cross_attention_kwargs=self.cross_attention_kwargs, | |
down_block_additional_residuals=down_block_res_samples, | |
mid_block_additional_residual=mid_block_res_sample, | |
down_intrablock_additional_residuals = down_intrablock_additional_residuals, | |
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) | |
if self.do_classifier_free_guidance and guidance_rescale > 0.0: | |
# Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf | |
noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) | |
# compute the previous noisy sample x_t -> x_t-1 | |
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] | |
if latent_processing == 1: | |
lst_latent.append(self.type_output("pil",device,prompt_embeds.dtype,return_dict,latents,generator).images[0]) | |
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) | |
# 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) | |
# If we do sequential model offloading, let's offload unet and controlnet | |
# manually for max memory savings | |
if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: | |
self.unet.to("cpu") | |
self.controlnet.to("cpu") | |
torch.cuda.empty_cache() | |
if latent_processing == 1: | |
if output_type == 'latent': | |
lst_latent.append(self.type_output(output_type,device,prompt_embeds.dtype,return_dict,latents,generator).images[0]) | |
return lst_latent | |
if output_type == 'latent': | |
return [self.type_output("pil",device,prompt_embeds.dtype,return_dict,latents,generator).images[0],self.type_output(output_type,device,prompt_embeds.dtype,return_dict,latents,generator).images[0]] | |
return [self.type_output(output_type,device,prompt_embeds.dtype,return_dict,latents,generator).images[0]] | |
class StableDiffusionImg2ImgPipeline_finetune(IPAdapterMixin,StableDiffusionImg2ImgPipeline): | |
def type_output(self,output_type,device,d_type,return_dict,latents,generator): | |
if not output_type == "latent": | |
image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False,generator=generator)[0] | |
image, has_nsfw_concept = self.run_safety_checker(image, device, d_type) | |
else: | |
image = latents | |
has_nsfw_concept = None | |
if has_nsfw_concept is None: | |
do_denormalize = [True] * image.shape[0] | |
else: | |
do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept] | |
image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize) | |
# Offload all models | |
self.maybe_free_model_hooks() | |
if not return_dict: | |
return (image, has_nsfw_concept) | |
return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) | |
def __call__( | |
self, | |
prompt: Union[str, List[str]] = None, | |
image: PipelineImageInput = None, | |
height: Optional[int] = None, | |
width: Optional[int] = None, | |
strength: float = 0.8, | |
num_inference_steps: Optional[int] = 50, | |
timesteps: List[int] = None, | |
sigmas: List[float] = None, | |
guidance_scale: Optional[float] = 7.5, | |
negative_prompt: Optional[Union[str, List[str]]] = None, | |
num_images_per_prompt: Optional[int] = 1, | |
eta: Optional[float] = 0.0, | |
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, | |
prompt_embeds: Optional[torch.Tensor] = None, | |
negative_prompt_embeds: Optional[torch.Tensor] = None, | |
ip_adapter_image: Optional[PipelineImageInput] = None, | |
ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, | |
output_type: Optional[str] = "pil", | |
return_dict: bool = True, | |
guidance_rescale: float = 0.0, | |
#callback: Optional[Callable[[int, int, torch.Tensor], None]] = None, | |
#callback_steps: int = 1, | |
cross_attention_kwargs: Optional[Dict[str, Any]] = None, | |
clip_skip: int = 0, | |
callback_on_step_end: Optional[ | |
Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] | |
] = None, | |
callback_on_step_end_tensor_inputs: List[str] = ["latents"], | |
region_map_state=None, | |
weight_func = lambda w, sigma, qk: w * sigma * qk.std(), | |
latent_processing = 0, | |
image_t2i_adapter : Optional[PipelineImageInput] = None, | |
adapter_conditioning_scale: Union[float, List[float]] = 1.0, | |
adapter_conditioning_factor: float = 1.0, | |
long_encode: int = 0, | |
**kwargs, | |
): | |
init_step = num_inference_steps | |
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`", | |
) | |
if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): | |
callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs | |
# 1. Check inputs. Raise error if not correct | |
self.check_inputs( | |
prompt, | |
strength, | |
callback_steps, | |
negative_prompt, | |
prompt_embeds, | |
negative_prompt_embeds, | |
ip_adapter_image, | |
ip_adapter_image_embeds, | |
callback_on_step_end_tensor_inputs, | |
) | |
#self.prompt_parser = FrozenCLIPEmbedderWithCustomWords(self.tokenizer, self.text_encoder,clip_skip+1) | |
self._guidance_scale = guidance_scale | |
self._clip_skip = clip_skip | |
self._cross_attention_kwargs = cross_attention_kwargs | |
self._interrupt = False | |
if height is None: | |
_,height = get_image_size(image) | |
height = int((height // 8)*8) | |
if width is None: | |
width,_ = get_image_size(image) | |
width = int((width // 8)*8) | |
adapter_state = None | |
if image_t2i_adapter is not None: | |
height, width = default_height_width(self,height, width, image_t2i_adapter) | |
adapter_state = preprocessing_t2i_adapter(self,image_t2i_adapter,width,height,adapter_conditioning_scale,num_images_per_prompt) | |
# 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 | |
) | |
#Copy input prompt_embeds and negative_prompt_embeds | |
prompt_embeds_copy = None | |
negative_prompt_embeds_copy = None | |
if prompt_embeds is not None: | |
prompt_embeds_copy = prompt_embeds.clone().detach() | |
if negative_prompt_embeds is not None: | |
negative_prompt_embeds_copy = negative_prompt_embeds.clone().detach() | |
prompt_embeds, negative_prompt_embeds,text_input_ids = encode_prompt_function( | |
self, | |
prompt, | |
device, | |
num_images_per_prompt, | |
self.do_classifier_free_guidance, | |
negative_prompt, | |
prompt_embeds=prompt_embeds, | |
negative_prompt_embeds=negative_prompt_embeds, | |
lora_scale=text_encoder_lora_scale, | |
clip_skip=self.clip_skip, | |
long_encode = long_encode, | |
) | |
#Get token_id | |
#text_input_ids,uncond_input_ids = get_id_text(self,prompt,max_length = prompt_embeds.shape[1],negative_prompt = negative_prompt,prompt_embeds = prompt_embeds_copy,negative_prompt_embeds = negative_prompt_embeds_copy) | |
# For classifier free guidance, we need to do two forward passes. | |
# Here we concatenate the unconditional and text embeddings into a single batch | |
# to avoid doing two forward passes | |
'''if text_input_ids is not None: | |
text_input_ids = np.concatenate([uncond_input_ids, text_input_ids])''' | |
#text_ids, text_embeddings = self.prompt_parser([negative_prompt, prompt]) | |
#text_embeddings = text_embeddings.to(self.unet.dtype) | |
# For classifier free guidance, we need to do two forward passes. | |
# Here we concatenate the unconditional and text embeddings into a single batch | |
# to avoid doing two forward passes | |
if self.do_classifier_free_guidance: | |
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) | |
if ip_adapter_image is not None or ip_adapter_image_embeds is not None: | |
image_embeds = self.prepare_ip_adapter_image_embeds( | |
ip_adapter_image, | |
ip_adapter_image_embeds, | |
device, | |
batch_size * num_images_per_prompt, | |
self.do_classifier_free_guidance, | |
) | |
# 4. Preprocess image | |
image = self.image_processor.preprocess(image) | |
# 5. set timesteps | |
region_state = encode_region_map( | |
self, | |
region_map_state, | |
width = width, | |
height = height, | |
num_images_per_prompt = num_images_per_prompt, | |
text_ids=text_input_ids, | |
) | |
if self.cross_attention_kwargs is None: | |
self._cross_attention_kwargs ={} | |
timesteps, num_inference_steps = retrieve_timesteps( | |
self.scheduler, num_inference_steps, device, timesteps, sigmas | |
) | |
timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) | |
latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) | |
# 6. Prepare latent variables | |
latents = self.prepare_latents( | |
image, | |
latent_timestep, | |
batch_size, | |
num_images_per_prompt, | |
prompt_embeds.dtype, | |
device, | |
generator, | |
) | |
lst_latent =[] | |
if latent_processing == 1: | |
lst_latent = [self.type_output("pil",device,prompt_embeds.dtype,return_dict,latents,generator).images[0]] | |
# 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline | |
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) | |
# 7.1 Add image embeds for IP-Adapter | |
added_cond_kwargs = ( | |
{"image_embeds": image_embeds} | |
if ip_adapter_image is not None or ip_adapter_image_embeds is not None | |
else None | |
) | |
# 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) | |
sigmas = self.scheduler.sigmas[init_step-len(timesteps):] | |
# 8. Denoising loop | |
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order | |
self._num_timesteps = len(timesteps) | |
with self.progress_bar(total=num_inference_steps) as progress_bar: | |
#step_x = 0 | |
for i, t in enumerate(timesteps): | |
if self.interrupt: | |
continue | |
# 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) | |
region_prompt = { | |
"region_state": region_state, | |
"sigma": self.scheduler.sigmas[i], | |
"weight_func": weight_func, | |
} | |
self._cross_attention_kwargs["region_prompt"] = region_prompt | |
#print(t) | |
#step_x=step_x+1 | |
down_intrablock_additional_residuals = None | |
if adapter_state is not None: | |
if i < int(num_inference_steps * adapter_conditioning_factor): | |
down_intrablock_additional_residuals = [state.clone() for state in adapter_state] | |
else: | |
down_intrablock_additional_residuals = None | |
# predict the noise residual | |
noise_pred = self.unet( | |
latent_model_input, | |
t, | |
encoder_hidden_states=prompt_embeds, | |
timestep_cond=timestep_cond, | |
cross_attention_kwargs=self.cross_attention_kwargs, | |
down_intrablock_additional_residuals = down_intrablock_additional_residuals, | |
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 + self.guidance_scale * (noise_pred_text - noise_pred_uncond) | |
if self.do_classifier_free_guidance and guidance_rescale > 0.0: | |
# Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf | |
noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) | |
# compute the previous noisy sample x_t -> x_t-1 | |
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] | |
if latent_processing == 1: | |
lst_latent.append(self.type_output("pil",device,prompt_embeds.dtype,return_dict,latents,generator).images[0]) | |
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) | |
# 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) | |
if latent_processing == 1: | |
if output_type == 'latent': | |
lst_latent.append(self.type_output(output_type,device,prompt_embeds.dtype,return_dict,latents,generator).images[0]) | |
return lst_latent | |
if output_type == 'latent': | |
return [self.type_output("pil",device,prompt_embeds.dtype,return_dict,latents,generator).images[0],self.type_output(output_type,device,prompt_embeds.dtype,return_dict,latents,generator).images[0]] | |
return [self.type_output(output_type,device,prompt_embeds.dtype,return_dict,latents,generator).images[0]] | |
class StableDiffusionInpaintPipeline_finetune(IPAdapterMixin,StableDiffusionInpaintPipeline): | |
def type_output(self,output_type,device,d_type,return_dict,latents,generator,init_image,padding_mask_crop,mask_image,original_image,crops_coords): | |
if not output_type == "latent": | |
condition_kwargs = {} | |
if isinstance(self.vae, AsymmetricAutoencoderKL): | |
init_image = init_image.to(device=device, dtype=masked_image_latents.dtype) | |
init_image_condition = init_image.clone() | |
init_image = self._encode_vae_image(init_image, generator=generator) | |
mask_condition = mask_condition.to(device=device, dtype=masked_image_latents.dtype) | |
condition_kwargs = {"image": init_image_condition, "mask": mask_condition} | |
image = self.vae.decode( | |
latents / self.vae.config.scaling_factor, return_dict=False, generator=generator, **condition_kwargs | |
)[0] | |
image, has_nsfw_concept = self.run_safety_checker(image, device, d_type) | |
else: | |
image = latents | |
has_nsfw_concept = None | |
if has_nsfw_concept is None: | |
do_denormalize = [True] * image.shape[0] | |
else: | |
do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept] | |
image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize) | |
if padding_mask_crop is not None: | |
image = [self.image_processor.apply_overlay(mask_image, original_image, i, crops_coords) for i in image] | |
# Offload all models | |
self.maybe_free_model_hooks() | |
if not return_dict: | |
return (image, has_nsfw_concept) | |
return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) | |
def __call__( | |
self, | |
prompt: Union[str, List[str]] = None, | |
image: PipelineImageInput = None, | |
mask_image: PipelineImageInput = None, | |
masked_image_latents: torch.Tensor = None, | |
height: Optional[int] = None, | |
width: Optional[int] = None, | |
padding_mask_crop: Optional[int] = None, | |
strength: float = 1.0, | |
num_inference_steps: int = 50, | |
timesteps: List[int] = None, | |
sigmas: List[float] = None, | |
guidance_scale: float = 7.5, | |
negative_prompt: Optional[Union[str, List[str]]] = None, | |
num_images_per_prompt: Optional[int] = 1, | |
eta: float = 0.0, | |
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, | |
latents: Optional[torch.Tensor] = None, | |
prompt_embeds: Optional[torch.Tensor] = None, | |
negative_prompt_embeds: Optional[torch.Tensor] = None, | |
ip_adapter_image: Optional[PipelineImageInput] = None, | |
ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, | |
output_type: Optional[str] = "pil", | |
return_dict: bool = True, | |
cross_attention_kwargs: Optional[Dict[str, Any]] = None, | |
clip_skip: int = None, | |
callback_on_step_end: Optional[ | |
Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] | |
] = None, | |
callback_on_step_end_tensor_inputs: List[str] = ["latents"], | |
region_map_state=None, | |
weight_func = lambda w, sigma, qk: w * sigma * qk.std(), | |
latent_processing = 0, | |
image_t2i_adapter : Optional[PipelineImageInput] = None, | |
adapter_conditioning_scale: Union[float, List[float]] = 1.0, | |
adapter_conditioning_factor: float = 1.0, | |
long_encode: int = 0, | |
guidance_rescale: float = 0.0, | |
**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`", | |
) | |
if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): | |
callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs | |
# 0. Default height and width to unet | |
'''height = height or self.unet.config.sample_size * self.vae_scale_factor | |
width = width or self.unet.config.sample_size * self.vae_scale_factor''' | |
if height is None: | |
_,height = get_image_size(image) | |
height = int((height // 8)*8) | |
if width is None: | |
width,_ = get_image_size(image) | |
width = int((width // 8)*8) | |
adapter_state = None | |
if image_t2i_adapter is not None: | |
height, width = default_height_width(self,height, width, image_t2i_adapter) | |
adapter_state = preprocessing_t2i_adapter(self,image_t2i_adapter,width,height,adapter_conditioning_scale,num_images_per_prompt) | |
# 1. Check inputs | |
self.check_inputs( | |
prompt, | |
image, | |
mask_image, | |
height, | |
width, | |
strength, | |
callback_steps, | |
output_type, | |
negative_prompt, | |
prompt_embeds, | |
negative_prompt_embeds, | |
ip_adapter_image, | |
ip_adapter_image_embeds, | |
callback_on_step_end_tensor_inputs, | |
padding_mask_crop, | |
) | |
self._guidance_scale = guidance_scale | |
self._clip_skip = clip_skip | |
self._cross_attention_kwargs = cross_attention_kwargs | |
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] | |
device = self._execution_device | |
# 3. Encode input prompt | |
text_encoder_lora_scale = ( | |
cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None | |
) | |
#Copy input prompt_embeds and negative_prompt_embeds | |
prompt_embeds_copy = None | |
negative_prompt_embeds_copy = None | |
if prompt_embeds is not None: | |
prompt_embeds_copy = prompt_embeds.clone().detach() | |
if negative_prompt_embeds is not None: | |
negative_prompt_embeds_copy = negative_prompt_embeds.clone().detach() | |
prompt_embeds, negative_prompt_embeds,text_input_ids = encode_prompt_function( | |
self, | |
prompt, | |
device, | |
num_images_per_prompt, | |
self.do_classifier_free_guidance, | |
negative_prompt, | |
prompt_embeds=prompt_embeds, | |
negative_prompt_embeds=negative_prompt_embeds, | |
lora_scale=text_encoder_lora_scale, | |
clip_skip=self.clip_skip, | |
long_encode = long_encode, | |
) | |
#Get token_id | |
#text_input_ids,uncond_input_ids = get_id_text(self,prompt,max_length = prompt_embeds.shape[1],negative_prompt = negative_prompt,prompt_embeds = prompt_embeds_copy,negative_prompt_embeds = negative_prompt_embeds_copy) | |
# For classifier free guidance, we need to do two forward passes. | |
# Here we concatenate the unconditional and text embeddings into a single batch | |
# to avoid doing two forward passes | |
'''if text_input_ids is not None: | |
text_input_ids = np.concatenate([uncond_input_ids, text_input_ids])''' | |
#text_ids, text_embeddings = self.prompt_parser([negative_prompt, prompt]) | |
#text_embeddings = text_embeddings.to(self.unet.dtype) | |
# For classifier free guidance, we need to do two forward passes. | |
# Here we concatenate the unconditional and text embeddings into a single batch | |
# to avoid doing two forward passes | |
if self.do_classifier_free_guidance: | |
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) | |
if ip_adapter_image is not None or ip_adapter_image_embeds is not None: | |
image_embeds = self.prepare_ip_adapter_image_embeds( | |
ip_adapter_image, | |
ip_adapter_image_embeds, | |
device, | |
batch_size * num_images_per_prompt, | |
self.do_classifier_free_guidance, | |
) | |
# 4. set timesteps | |
timesteps, num_inference_steps = retrieve_timesteps( | |
self.scheduler, num_inference_steps, device, timesteps, sigmas | |
) | |
timesteps, num_inference_steps = self.get_timesteps( | |
num_inference_steps=num_inference_steps, strength=strength, device=device | |
) | |
# check that number of inference steps is not < 1 - as this doesn't make sense | |
if num_inference_steps < 1: | |
raise ValueError( | |
f"After adjusting the num_inference_steps by strength parameter: {strength}, the number of pipeline" | |
f"steps is {num_inference_steps} which is < 1 and not appropriate for this pipeline." | |
) | |
# at which timestep to set the initial noise (n.b. 50% if strength is 0.5) | |
latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) | |
# create a boolean to check if the strength is set to 1. if so then initialise the latents with pure noise | |
is_strength_max = strength == 1.0 | |
#4.1 Preprocess region mao | |
region_state = encode_region_map( | |
self, | |
region_map_state, | |
width = width, | |
height = height, | |
num_images_per_prompt = num_images_per_prompt, | |
text_ids=text_input_ids, | |
) | |
if self.cross_attention_kwargs is None: | |
self._cross_attention_kwargs ={} | |
# 5. Preprocess mask and image | |
if padding_mask_crop is not None: | |
crops_coords = self.mask_processor.get_crop_region(mask_image, width, height, pad=padding_mask_crop) | |
resize_mode = "fill" | |
else: | |
crops_coords = None | |
resize_mode = "default" | |
original_image = image | |
init_image = self.image_processor.preprocess( | |
image, height=height, width=width, crops_coords=crops_coords, resize_mode=resize_mode | |
) | |
init_image = init_image.to(dtype=torch.float32) | |
# 6. Prepare latent variables | |
num_channels_latents = self.vae.config.latent_channels | |
num_channels_unet = self.unet.config.in_channels | |
return_image_latents = num_channels_unet == 4 | |
latents_outputs = self.prepare_latents( | |
batch_size * num_images_per_prompt, | |
num_channels_latents, | |
height, | |
width, | |
prompt_embeds.dtype, | |
device, | |
generator, | |
latents, | |
image=init_image, | |
timestep=latent_timestep, | |
is_strength_max=is_strength_max, | |
return_noise=True, | |
return_image_latents=return_image_latents, | |
) | |
if return_image_latents: | |
latents, noise, image_latents = latents_outputs | |
else: | |
latents, noise = latents_outputs | |
# 7. Prepare mask latent variables | |
mask_condition = self.mask_processor.preprocess( | |
mask_image, height=height, width=width, resize_mode=resize_mode, crops_coords=crops_coords | |
) | |
if masked_image_latents is None: | |
masked_image = init_image * (mask_condition < 0.5) | |
else: | |
masked_image = masked_image_latents | |
mask, masked_image_latents = self.prepare_mask_latents( | |
mask_condition, | |
masked_image, | |
batch_size * num_images_per_prompt, | |
height, | |
width, | |
prompt_embeds.dtype, | |
device, | |
generator, | |
self.do_classifier_free_guidance, | |
) | |
# 8. Check that sizes of mask, masked image and latents match | |
if num_channels_unet == 9: | |
# default case for runwayml/stable-diffusion-inpainting | |
num_channels_mask = mask.shape[1] | |
num_channels_masked_image = masked_image_latents.shape[1] | |
if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels: | |
raise ValueError( | |
f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects" | |
f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +" | |
f" `num_channels_mask`: {num_channels_mask} + `num_channels_masked_image`: {num_channels_masked_image}" | |
f" = {num_channels_latents+num_channels_masked_image+num_channels_mask}. Please verify the config of" | |
" `pipeline.unet` or your `mask_image` or `image` input." | |
) | |
elif num_channels_unet != 4: | |
raise ValueError( | |
f"The unet {self.unet.__class__} should have either 4 or 9 input channels, not {self.unet.config.in_channels}." | |
) | |
# 9. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline | |
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) | |
# 9.1 Add image embeds for IP-Adapter | |
added_cond_kwargs = ( | |
{"image_embeds": image_embeds} | |
if ip_adapter_image is not None or ip_adapter_image_embeds is not None | |
else None | |
) | |
# 9.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) | |
lst_latent =[] | |
if latent_processing == 1: | |
lst_latent = [self.type_output("pil",device,prompt_embeds.dtype,return_dict,latents,generator,init_image,padding_mask_crop,mask_image,original_image,crops_coords).images[0]] | |
# 10. Denoising loop | |
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order | |
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 | |
# 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 | |
# concat latents, mask, masked_image_latents in the channel dimension | |
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) | |
if num_channels_unet == 9: | |
latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1) | |
region_prompt = { | |
"region_state": region_state, | |
"sigma": self.scheduler.sigmas[i], | |
"weight_func": weight_func, | |
} | |
self._cross_attention_kwargs["region_prompt"] = region_prompt | |
down_intrablock_additional_residuals = None | |
if adapter_state is not None: | |
if i < int(num_inference_steps * adapter_conditioning_factor): | |
down_intrablock_additional_residuals = [state.clone() for state in adapter_state] | |
else: | |
down_intrablock_additional_residuals = None | |
# predict the noise residual | |
noise_pred = self.unet( | |
latent_model_input, | |
t, | |
encoder_hidden_states=prompt_embeds, | |
timestep_cond=timestep_cond, | |
cross_attention_kwargs=self.cross_attention_kwargs, | |
down_intrablock_additional_residuals = down_intrablock_additional_residuals, | |
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 + self.guidance_scale * (noise_pred_text - noise_pred_uncond) | |
if self.do_classifier_free_guidance and guidance_rescale > 0.0: | |
# Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf | |
noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) | |
# compute the previous noisy sample x_t -> x_t-1 | |
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] | |
if num_channels_unet == 4: | |
init_latents_proper = image_latents | |
if self.do_classifier_free_guidance: | |
init_mask, _ = mask.chunk(2) | |
else: | |
init_mask = mask | |
if i < len(timesteps) - 1: | |
noise_timestep = timesteps[i + 1] | |
init_latents_proper = self.scheduler.add_noise( | |
init_latents_proper, noise, torch.tensor([noise_timestep]) | |
) | |
latents = (1 - init_mask) * init_latents_proper + init_mask * latents | |
if latent_processing == 1: | |
lst_latent.append(self.type_output("pil",device,prompt_embeds.dtype,return_dict,latents,generator,init_image,padding_mask_crop,mask_image,original_image,crops_coords).images[0]) | |
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) | |
mask = callback_outputs.pop("mask", mask) | |
masked_image_latents = callback_outputs.pop("masked_image_latents", masked_image_latents) | |
# 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) | |
if latent_processing == 1: | |
if output_type == 'latent': | |
lst_latent.append(self.type_output(output_type,device,prompt_embeds.dtype,return_dict,latents,generator,init_image,padding_mask_crop,mask_image,original_image,crops_coords).images[0]) | |
return lst_latent | |
if output_type == 'latent': | |
return [self.type_output("pil",device,prompt_embeds.dtype,return_dict,latents,generator,init_image,padding_mask_crop,mask_image,original_image,crops_coords).images[0],self.type_output(output_type,device,prompt_embeds.dtype,return_dict,latents,generator,init_image,padding_mask_crop,mask_image,original_image,crops_coords).images[0]] | |
return [self.type_output(output_type,device,prompt_embeds.dtype,return_dict,latents,generator,init_image,padding_mask_crop,mask_image,original_image,crops_coords).images[0]] | |
class StableDiffusionControlNetInpaintPipeline_finetune(IPAdapterMixin,StableDiffusionControlNetInpaintPipeline): | |
def type_output(self,output_type,device,d_type,return_dict,latents,generator,padding_mask_crop,mask_image,original_image,crops_coords): | |
if not output_type == "latent": | |
image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[ | |
0 | |
] | |
image, has_nsfw_concept = self.run_safety_checker(image, device,d_type) | |
else: | |
image = latents | |
has_nsfw_concept = None | |
if has_nsfw_concept is None: | |
do_denormalize = [True] * image.shape[0] | |
else: | |
do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept] | |
image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize) | |
if padding_mask_crop is not None: | |
image = [self.image_processor.apply_overlay(mask_image, original_image, i, crops_coords) for i in image] | |
# Offload all models | |
self.maybe_free_model_hooks() | |
if not return_dict: | |
return (image, has_nsfw_concept) | |
return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) | |
def __call__( | |
self, | |
prompt: Union[str, List[str]] = None, | |
image: PipelineImageInput = None, | |
mask_image: PipelineImageInput = None, | |
control_image: PipelineImageInput = None, | |
height: Optional[int] = None, | |
width: Optional[int] = None, | |
padding_mask_crop: Optional[int] = None, | |
strength: float = 1.0, | |
num_inference_steps: int = 50, | |
guidance_scale: float = 7.5, | |
negative_prompt: Optional[Union[str, List[str]]] = None, | |
num_images_per_prompt: Optional[int] = 1, | |
eta: float = 0.0, | |
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, | |
latents: Optional[torch.Tensor] = None, | |
prompt_embeds: Optional[torch.Tensor] = None, | |
negative_prompt_embeds: Optional[torch.Tensor] = None, | |
ip_adapter_image: Optional[PipelineImageInput] = None, | |
ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, | |
output_type: Optional[str] = "pil", | |
return_dict: bool = True, | |
cross_attention_kwargs: Optional[Dict[str, Any]] = None, | |
controlnet_conditioning_scale: Union[float, List[float]] = 0.5, | |
guess_mode: bool = False, | |
control_guidance_start: Union[float, List[float]] = 0.0, | |
control_guidance_end: Union[float, List[float]] = 1.0, | |
clip_skip: Optional[int] = None, | |
callback_on_step_end: Optional[ | |
Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] | |
] = None, | |
callback_on_step_end_tensor_inputs: List[str] = ["latents"], | |
region_map_state=None, | |
weight_func = lambda w, sigma, qk: w * sigma * qk.std(), | |
latent_processing = 0, | |
image_t2i_adapter : Optional[PipelineImageInput] = None, | |
adapter_conditioning_scale: Union[float, List[float]] = 1.0, | |
adapter_conditioning_factor: float = 1.0, | |
long_encode: int = 0, | |
guidance_rescale: float = 0.0, | |
**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 using `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 using `callback_on_step_end`", | |
) | |
if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): | |
callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs | |
if height is None: | |
_,height = get_image_size(image) | |
height = int((height // 8)*8) | |
if width is None: | |
width,_ = get_image_size(image) | |
width = int((width // 8)*8) | |
adapter_state = None | |
if image_t2i_adapter is not None: | |
height, width = default_height_width(self,height, width, image_t2i_adapter) | |
adapter_state = preprocessing_t2i_adapter(self,image_t2i_adapter,width,height,adapter_conditioning_scale,num_images_per_prompt) | |
controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet | |
# align format for control guidance | |
if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): | |
control_guidance_start = len(control_guidance_end) * [control_guidance_start] | |
elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): | |
control_guidance_end = len(control_guidance_start) * [control_guidance_end] | |
elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list): | |
mult = len(controlnet.nets) if isinstance(controlnet, MultiControlNetModel) else 1 | |
control_guidance_start, control_guidance_end = ( | |
mult * [control_guidance_start], | |
mult * [control_guidance_end], | |
) | |
# 1. Check inputs. Raise error if not correct | |
self.check_inputs( | |
prompt, | |
control_image, | |
mask_image, | |
height, | |
width, | |
callback_steps, | |
output_type, | |
negative_prompt, | |
prompt_embeds, | |
negative_prompt_embeds, | |
ip_adapter_image, | |
ip_adapter_image_embeds, | |
controlnet_conditioning_scale, | |
control_guidance_start, | |
control_guidance_end, | |
callback_on_step_end_tensor_inputs, | |
padding_mask_crop, | |
) | |
self._guidance_scale = guidance_scale | |
self._clip_skip = clip_skip | |
self._cross_attention_kwargs = cross_attention_kwargs | |
# 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 padding_mask_crop is not None: | |
height, width = self.image_processor.get_default_height_width(image, height, width) | |
crops_coords = self.mask_processor.get_crop_region(mask_image, width, height, pad=padding_mask_crop) | |
resize_mode = "fill" | |
else: | |
crops_coords = None | |
resize_mode = "default" | |
device = self._execution_device | |
if isinstance(controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float): | |
controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(controlnet.nets) | |
global_pool_conditions = ( | |
controlnet.config.global_pool_conditions | |
if isinstance(controlnet, ControlNetModel) | |
else controlnet.nets[0].config.global_pool_conditions | |
) | |
guess_mode = guess_mode or global_pool_conditions | |
# 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 | |
) | |
#Copy input prompt_embeds and negative_prompt_embeds | |
'''prompt_embeds_copy = None | |
negative_prompt_embeds_copy = None | |
if prompt_embeds is not None: | |
prompt_embeds_copy = prompt_embeds.clone().detach() | |
if negative_prompt_embeds is not None: | |
negative_prompt_embeds_copy = negative_prompt_embeds.clone().detach()''' | |
prompt_embeds, negative_prompt_embeds,text_input_ids = encode_prompt_function( | |
self, | |
prompt, | |
device, | |
num_images_per_prompt, | |
self.do_classifier_free_guidance, | |
negative_prompt, | |
prompt_embeds=prompt_embeds, | |
negative_prompt_embeds=negative_prompt_embeds, | |
lora_scale=text_encoder_lora_scale, | |
clip_skip=self.clip_skip, | |
long_encode = long_encode, | |
) | |
#Get token_id | |
#text_input_ids,uncond_input_ids = get_id_text(self,prompt,max_length = prompt_embeds.shape[1],negative_prompt = negative_prompt,prompt_embeds = prompt_embeds_copy,negative_prompt_embeds = negative_prompt_embeds_copy) | |
# For classifier free guidance, we need to do two forward passes. | |
# Here we concatenate the unconditional and text embeddings into a single batch | |
# to avoid doing two forward passes | |
'''if text_input_ids is not None: | |
text_input_ids = np.concatenate([uncond_input_ids, text_input_ids])''' | |
#text_ids, text_embeddings = self.prompt_parser([negative_prompt, prompt]) | |
#text_embeddings = text_embeddings.to(self.unet.dtype) | |
# For classifier free guidance, we need to do two forward passes. | |
# Here we concatenate the unconditional and text embeddings into a single batch | |
# to avoid doing two forward passes | |
if self.do_classifier_free_guidance: | |
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) | |
if ip_adapter_image is not None or ip_adapter_image_embeds is not None: | |
image_embeds = self.prepare_ip_adapter_image_embeds( | |
ip_adapter_image, | |
ip_adapter_image_embeds, | |
device, | |
batch_size * num_images_per_prompt, | |
self.do_classifier_free_guidance, | |
) | |
# 4. Prepare image | |
if isinstance(controlnet, ControlNetModel): | |
control_image = self.prepare_control_image( | |
image=control_image, | |
width=width, | |
height=height, | |
batch_size=batch_size * num_images_per_prompt, | |
num_images_per_prompt=num_images_per_prompt, | |
device=device, | |
dtype=controlnet.dtype, | |
crops_coords=crops_coords, | |
resize_mode=resize_mode, | |
do_classifier_free_guidance=self.do_classifier_free_guidance, | |
guess_mode=guess_mode, | |
) | |
elif isinstance(controlnet, MultiControlNetModel): | |
control_images = [] | |
for control_image_ in control_image: | |
control_image_ = self.prepare_control_image( | |
image=control_image_, | |
width=width, | |
height=height, | |
batch_size=batch_size * num_images_per_prompt, | |
num_images_per_prompt=num_images_per_prompt, | |
device=device, | |
dtype=controlnet.dtype, | |
crops_coords=crops_coords, | |
resize_mode=resize_mode, | |
do_classifier_free_guidance=self.do_classifier_free_guidance, | |
guess_mode=guess_mode, | |
) | |
control_images.append(control_image_) | |
control_image = control_images | |
else: | |
assert False | |
# 4.1 Preprocess mask and image - resizes image and mask w.r.t height and width | |
original_image = image | |
init_image = self.image_processor.preprocess( | |
image, height=height, width=width, crops_coords=crops_coords, resize_mode=resize_mode | |
) | |
init_image = init_image.to(dtype=torch.float32) | |
mask = self.mask_processor.preprocess( | |
mask_image, height=height, width=width, resize_mode=resize_mode, crops_coords=crops_coords | |
) | |
masked_image = init_image * (mask < 0.5) | |
_, _, height, width = init_image.shape | |
#4.2 Preprocess region mao | |
region_state = encode_region_map( | |
self, | |
region_map_state, | |
width = width, | |
height = height, | |
num_images_per_prompt = num_images_per_prompt, | |
text_ids=text_input_ids, | |
) | |
if self.cross_attention_kwargs is None: | |
self._cross_attention_kwargs ={} | |
# 5. Prepare timesteps | |
self.scheduler.set_timesteps(num_inference_steps, device=device) | |
timesteps, num_inference_steps = self.get_timesteps( | |
num_inference_steps=num_inference_steps, strength=strength, device=device | |
) | |
# at which timestep to set the initial noise (n.b. 50% if strength is 0.5) | |
latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) | |
# create a boolean to check if the strength is set to 1. if so then initialise the latents with pure noise | |
is_strength_max = strength == 1.0 | |
self._num_timesteps = len(timesteps) | |
# 6. Prepare latent variables | |
num_channels_latents = self.vae.config.latent_channels | |
num_channels_unet = self.unet.config.in_channels | |
return_image_latents = num_channels_unet == 4 | |
latents_outputs = self.prepare_latents( | |
batch_size * num_images_per_prompt, | |
num_channels_latents, | |
height, | |
width, | |
prompt_embeds.dtype, | |
device, | |
generator, | |
latents, | |
image=init_image, | |
timestep=latent_timestep, | |
is_strength_max=is_strength_max, | |
return_noise=True, | |
return_image_latents=return_image_latents, | |
) | |
if return_image_latents: | |
latents, noise, image_latents = latents_outputs | |
else: | |
latents, noise = latents_outputs | |
# 7. Prepare mask latent variables | |
mask, masked_image_latents = self.prepare_mask_latents( | |
mask, | |
masked_image, | |
batch_size * num_images_per_prompt, | |
height, | |
width, | |
prompt_embeds.dtype, | |
device, | |
generator, | |
self.do_classifier_free_guidance, | |
) | |
# 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline | |
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) | |
# 7.1 Add image embeds for IP-Adapter | |
added_cond_kwargs = ( | |
{"image_embeds": image_embeds} | |
if ip_adapter_image is not None or ip_adapter_image_embeds is not None | |
else None | |
) | |
# 7.2 Create tensor stating which controlnets to keep | |
controlnet_keep = [] | |
for i in range(len(timesteps)): | |
keeps = [ | |
1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e) | |
for s, e in zip(control_guidance_start, control_guidance_end) | |
] | |
controlnet_keep.append(keeps[0] if isinstance(controlnet, ControlNetModel) else keeps) | |
lst_latent =[] | |
if latent_processing == 1: | |
lst_latent = [self.type_output("pil",device,prompt_embeds.dtype,return_dict,latents,generator,padding_mask_crop,mask_image,original_image,crops_coords).images[0]] | |
# 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 | |
if guess_mode and self.do_classifier_free_guidance: | |
# Infer ControlNet only for the conditional batch. | |
control_model_input = latents | |
control_model_input = self.scheduler.scale_model_input(control_model_input, t) | |
controlnet_prompt_embeds = prompt_embeds.chunk(2)[1] | |
else: | |
control_model_input = latent_model_input | |
controlnet_prompt_embeds = prompt_embeds | |
if isinstance(controlnet_keep[i], list): | |
cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])] | |
else: | |
controlnet_cond_scale = controlnet_conditioning_scale | |
if isinstance(controlnet_cond_scale, list): | |
controlnet_cond_scale = controlnet_cond_scale[0] | |
cond_scale = controlnet_cond_scale * controlnet_keep[i] | |
down_block_res_samples, mid_block_res_sample = self.controlnet( | |
control_model_input, | |
t, | |
encoder_hidden_states=controlnet_prompt_embeds, | |
controlnet_cond=control_image, | |
conditioning_scale=cond_scale, | |
guess_mode=guess_mode, | |
return_dict=False, | |
) | |
if guess_mode and self.do_classifier_free_guidance: | |
# Infered ControlNet only for the conditional batch. | |
# To apply the output of ControlNet to both the unconditional and conditional batches, | |
# add 0 to the unconditional batch to keep it unchanged. | |
down_block_res_samples = [torch.cat([torch.zeros_like(d), d]) for d in down_block_res_samples] | |
mid_block_res_sample = torch.cat([torch.zeros_like(mid_block_res_sample), mid_block_res_sample]) | |
# predict the noise residual | |
if num_channels_unet == 9: | |
latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1) | |
region_prompt = { | |
"region_state": region_state, | |
"sigma": self.scheduler.sigmas[i], | |
"weight_func": weight_func, | |
} | |
self._cross_attention_kwargs["region_prompt"] = region_prompt | |
down_intrablock_additional_residuals = None | |
if adapter_state is not None: | |
if i < int(num_inference_steps * adapter_conditioning_factor): | |
down_intrablock_additional_residuals = [state.clone() for state in adapter_state] | |
else: | |
down_intrablock_additional_residuals = None | |
noise_pred = self.unet( | |
latent_model_input, | |
t, | |
encoder_hidden_states=prompt_embeds, | |
cross_attention_kwargs=self.cross_attention_kwargs, | |
down_block_additional_residuals=down_block_res_samples, | |
mid_block_additional_residual=mid_block_res_sample, | |
down_intrablock_additional_residuals = down_intrablock_additional_residuals, | |
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) | |
if self.do_classifier_free_guidance and guidance_rescale > 0.0: | |
# Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf | |
noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) | |
# compute the previous noisy sample x_t -> x_t-1 | |
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] | |
if num_channels_unet == 4: | |
init_latents_proper = image_latents | |
if self.do_classifier_free_guidance: | |
init_mask, _ = mask.chunk(2) | |
else: | |
init_mask = mask | |
if i < len(timesteps) - 1: | |
noise_timestep = timesteps[i + 1] | |
init_latents_proper = self.scheduler.add_noise( | |
init_latents_proper, noise, torch.tensor([noise_timestep]) | |
) | |
latents = (1 - init_mask) * init_latents_proper + init_mask * latents | |
if latent_processing == 1: | |
lst_latent.append(self.type_output("pil",device,prompt_embeds.dtype,return_dict,latents,generator,padding_mask_crop,mask_image,original_image,crops_coords).images[0]) | |
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) | |
# 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) | |
# If we do sequential model offloading, let's offload unet and controlnet | |
# manually for max memory savings | |
if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: | |
self.unet.to("cpu") | |
self.controlnet.to("cpu") | |
torch.cuda.empty_cache() | |
if latent_processing == 1: | |
if output_type == 'latent': | |
lst_latent.append(self.type_output(output_type,device,prompt_embeds.dtype,return_dict,latents,generator,padding_mask_crop,mask_image,original_image,crops_coords).images[0]) | |
return lst_latent | |
if output_type == 'latent': | |
return [self.type_output("pil",device,prompt_embeds.dtype,return_dict,latents,generator,padding_mask_crop,mask_image,original_image,crops_coords).images[0],self.type_output(output_type,device,prompt_embeds.dtype,return_dict,latents,generator,init_image,padding_mask_crop,mask_image,original_image,crops_coords).images[0]] | |
return [self.type_output(output_type,device,prompt_embeds.dtype,return_dict,latents,generator,padding_mask_crop,mask_image,original_image,crops_coords).images[0]] |