EcoTry/IDM-VTON
0
1# Copyright 2023 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15import inspect16from typing import Any, Callable, Dict, List, Optional, Tuple, Union17 18import numpy as np19import PIL.Image20import torch21from transformers import (22 CLIPImageProcessor,23 CLIPTextModel,24 CLIPTextModelWithProjection,25 CLIPTokenizer,26 CLIPVisionModelWithProjection,27)28 29from diffusers.image_processor import PipelineImageInput, VaeImageProcessor30from diffusers.loaders import (31 FromSingleFileMixin,32 IPAdapterMixin,33 StableDiffusionXLLoraLoaderMixin,34 TextualInversionLoaderMixin,35)36from diffusers.models import AutoencoderKL, ImageProjection, UNet2DConditionModel37from diffusers.models.attention_processor import (38 AttnProcessor2_0,39 FusedAttnProcessor2_0,40 LoRAAttnProcessor2_0,41 LoRAXFormersAttnProcessor,42 XFormersAttnProcessor,43)44from diffusers.models.lora import adjust_lora_scale_text_encoder45from diffusers.schedulers import KarrasDiffusionSchedulers46from diffusers.utils import (47 USE_PEFT_BACKEND,48 deprecate,49 is_invisible_watermark_available,50 is_torch_xla_available,51 logging,52 replace_example_docstring,53 scale_lora_layers,54 unscale_lora_layers,55)56from diffusers.utils.torch_utils import randn_tensor57from diffusers.pipelines.pipeline_utils import DiffusionPipeline58 59 60 61if is_torch_xla_available():62 import torch_xla.core.xla_model as xm63 64 XLA_AVAILABLE = True65else:66 XLA_AVAILABLE = False67 68 69logger = logging.get_logger(__name__) # pylint: disable=invalid-name70 71 72EXAMPLE_DOC_STRING = """73 Examples:74 ```py75 >>> import torch76 >>> from diffusers import StableDiffusionXLInpaintPipeline77 >>> from diffusers.utils import load_image78 79 >>> pipe = StableDiffusionXLInpaintPipeline.from_pretrained(80 ... "stabilityai/stable-diffusion-xl-base-1.0",81 ... torch_dtype=torch.float16,82 ... variant="fp16",83 ... use_safetensors=True,84 ... )85 >>> pipe.to("cuda")86 87 >>> img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"88 >>> mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"89 90 >>> init_image = load_image(img_url).convert("RGB")91 >>> mask_image = load_image(mask_url).convert("RGB")92 93 >>> prompt = "A majestic tiger sitting on a bench"94 >>> image = pipe(95 ... prompt=prompt, image=init_image, mask_image=mask_image, num_inference_steps=50, strength=0.8096 ... ).images[0]97 ```98"""99 100 101# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg102def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):103 """104 Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and105 Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4106 """107 std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)108 std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)109 # rescale the results from guidance (fixes overexposure)110 noise_pred_rescaled = noise_cfg * (std_text / std_cfg)111 # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images112 noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg113 return noise_cfg114 115 116def mask_pil_to_torch(mask, height, width):117 # preprocess mask118 if isinstance(mask, (PIL.Image.Image, np.ndarray)):119 mask = [mask]120 121 if isinstance(mask, list) and isinstance(mask[0], PIL.Image.Image):122 mask = [i.resize((width, height), resample=PIL.Image.LANCZOS) for i in mask]123 mask = np.concatenate([np.array(m.convert("L"))[None, None, :] for m in mask], axis=0)124 mask = mask.astype(np.float32) / 255.0125 elif isinstance(mask, list) and isinstance(mask[0], np.ndarray):126 mask = np.concatenate([m[None, None, :] for m in mask], axis=0)127 128 mask = torch.from_numpy(mask)129 return mask130 131 132def prepare_mask_and_masked_image(image, mask, height, width, return_image: bool = False):133 """134 Prepares a pair (image, mask) to be consumed by the Stable Diffusion pipeline. This means that those inputs will be135 converted to ``torch.Tensor`` with shapes ``batch x channels x height x width`` where ``channels`` is ``3`` for the136 ``image`` and ``1`` for the ``mask``.137 138 The ``image`` will be converted to ``torch.float32`` and normalized to be in ``[-1, 1]``. The ``mask`` will be139 binarized (``mask > 0.5``) and cast to ``torch.float32`` too.140 141 Args:142 image (Union[np.array, PIL.Image, torch.Tensor]): The image to inpaint.143 It can be a ``PIL.Image``, or a ``height x width x 3`` ``np.array`` or a ``channels x height x width``144 ``torch.Tensor`` or a ``batch x channels x height x width`` ``torch.Tensor``.145 mask (_type_): The mask to apply to the image, i.e. regions to inpaint.146 It can be a ``PIL.Image``, or a ``height x width`` ``np.array`` or a ``1 x height x width``147 ``torch.Tensor`` or a ``batch x 1 x height x width`` ``torch.Tensor``.148 149 150 Raises:151 ValueError: ``torch.Tensor`` images should be in the ``[-1, 1]`` range. ValueError: ``torch.Tensor`` mask152 should be in the ``[0, 1]`` range. ValueError: ``mask`` and ``image`` should have the same spatial dimensions.153 TypeError: ``mask`` is a ``torch.Tensor`` but ``image`` is not154 (ot the other way around).155 156 Returns:157 tuple[torch.Tensor]: The pair (mask, masked_image) as ``torch.Tensor`` with 4158 dimensions: ``batch x channels x height x width``.159 """160 161 # checkpoint. TOD(Yiyi) - need to clean this up later162 deprecation_message = "The prepare_mask_and_masked_image method is deprecated and will be removed in a future version. Please use VaeImageProcessor.preprocess instead"163 deprecate(164 "prepare_mask_and_masked_image",165 "0.30.0",166 deprecation_message,167 )168 if image is None:169 raise ValueError("`image` input cannot be undefined.")170 171 if mask is None:172 raise ValueError("`mask_image` input cannot be undefined.")173 174 if isinstance(image, torch.Tensor):175 if not isinstance(mask, torch.Tensor):176 mask = mask_pil_to_torch(mask, height, width)177 178 if image.ndim == 3:179 image = image.unsqueeze(0)180 181 # Batch and add channel dim for single mask182 if mask.ndim == 2:183 mask = mask.unsqueeze(0).unsqueeze(0)184 185 # Batch single mask or add channel dim186 if mask.ndim == 3:187 # Single batched mask, no channel dim or single mask not batched but channel dim188 if mask.shape[0] == 1:189 mask = mask.unsqueeze(0)190 191 # Batched masks no channel dim192 else:193 mask = mask.unsqueeze(1)194 195 assert image.ndim == 4 and mask.ndim == 4, "Image and Mask must have 4 dimensions"196 # assert image.shape[-2:] == mask.shape[-2:], "Image and Mask must have the same spatial dimensions"197 assert image.shape[0] == mask.shape[0], "Image and Mask must have the same batch size"198 199 # Check image is in [-1, 1]200 # if image.min() < -1 or image.max() > 1:201 # raise ValueError("Image should be in [-1, 1] range")202 203 # Check mask is in [0, 1]204 if mask.min() < 0 or mask.max() > 1:205 raise ValueError("Mask should be in [0, 1] range")206 207 # Binarize mask208 mask[mask < 0.5] = 0209 mask[mask >= 0.5] = 1210 211 # Image as float32212 image = image.to(dtype=torch.float32)213 elif isinstance(mask, torch.Tensor):214 raise TypeError(f"`mask` is a torch.Tensor but `image` (type: {type(image)} is not")215 else:216 # preprocess image217 if isinstance(image, (PIL.Image.Image, np.ndarray)):218 image = [image]219 if isinstance(image, list) and isinstance(image[0], PIL.Image.Image):220 # resize all images w.r.t passed height an width221 image = [i.resize((width, height), resample=PIL.Image.LANCZOS) for i in image]222 image = [np.array(i.convert("RGB"))[None, :] for i in image]223 image = np.concatenate(image, axis=0)224 elif isinstance(image, list) and isinstance(image[0], np.ndarray):225 image = np.concatenate([i[None, :] for i in image], axis=0)226 227 image = image.transpose(0, 3, 1, 2)228 image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0229 230 mask = mask_pil_to_torch(mask, height, width)231 mask[mask < 0.5] = 0232 mask[mask >= 0.5] = 1233 234 if image.shape[1] == 4:235 # images are in latent space and thus can't236 # be masked set masked_image to None237 # we assume that the checkpoint is not an inpainting238 # checkpoint. TOD(Yiyi) - need to clean this up later239 masked_image = None240 else:241 masked_image = image * (mask < 0.5)242 243 # n.b. ensure backwards compatibility as old function does not return image244 if return_image:245 return mask, masked_image, image246 247 return mask, masked_image248 249 250# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents251def retrieve_latents(252 encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"253):254 if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":255 return encoder_output.latent_dist.sample(generator)256 elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":257 return encoder_output.latent_dist.mode()258 elif hasattr(encoder_output, "latents"):259 return encoder_output.latents260 else:261 raise AttributeError("Could not access latents of provided encoder_output")262 263 264# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps265def retrieve_timesteps(266 scheduler,267 num_inference_steps: Optional[int] = None,268 device: Optional[Union[str, torch.device]] = None,269 timesteps: Optional[List[int]] = None,270 **kwargs,271):272 """273 Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles274 custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.275 276 Args:277 scheduler (`SchedulerMixin`):278 The scheduler to get timesteps from.279 num_inference_steps (`int`):280 The number of diffusion steps used when generating samples with a pre-trained model. If used,281 `timesteps` must be `None`.282 device (`str` or `torch.device`, *optional*):283 The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.284 timesteps (`List[int]`, *optional*):285 Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default286 timestep spacing strategy of the scheduler is used. If `timesteps` is passed, `num_inference_steps`287 must be `None`.288 289 Returns:290 `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the291 second element is the number of inference steps.292 """293 if timesteps is not None:294 accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())295 if not accepts_timesteps:296 raise ValueError(297 f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"298 f" timestep schedules. Please check whether you are using the correct scheduler."299 )300 scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)301 timesteps = scheduler.timesteps302 num_inference_steps = len(timesteps)303 else:304 scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)305 timesteps = scheduler.timesteps306 return timesteps, num_inference_steps307 308 309class StableDiffusionXLInpaintPipeline(310 DiffusionPipeline,311 TextualInversionLoaderMixin,312 StableDiffusionXLLoraLoaderMixin,313 FromSingleFileMixin,314 IPAdapterMixin,315):316 r"""317 Pipeline for text-to-image generation using Stable Diffusion XL.318 319 This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the320 library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)321 322 The pipeline also inherits the following loading methods:323 - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings324 - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files325 - [`~loaders.StableDiffusionXLLoraLoaderMixin.load_lora_weights`] for loading LoRA weights326 - [`~loaders.StableDiffusionXLLoraLoaderMixin.save_lora_weights`] for saving LoRA weights327 - [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters328 329 Args:330 vae ([`AutoencoderKL`]):331 Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.332 text_encoder ([`CLIPTextModel`]):333 Frozen text-encoder. Stable Diffusion XL uses the text portion of334 [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically335 the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.336 text_encoder_2 ([` CLIPTextModelWithProjection`]):337 Second frozen text-encoder. Stable Diffusion XL uses the text and pool portion of338 [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection),339 specifically the340 [laion/CLIP-ViT-bigG-14-laion2B-39B-b160k](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k)341 variant.342 tokenizer (`CLIPTokenizer`):343 Tokenizer of class344 [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).345 tokenizer_2 (`CLIPTokenizer`):346 Second Tokenizer of class347 [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).348 unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.349 scheduler ([`SchedulerMixin`]):350 A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of351 [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].352 requires_aesthetics_score (`bool`, *optional*, defaults to `"False"`):353 Whether the `unet` requires a aesthetic_score condition to be passed during inference. Also see the config354 of `stabilityai/stable-diffusion-xl-refiner-1-0`.355 force_zeros_for_empty_prompt (`bool`, *optional*, defaults to `"True"`):356 Whether the negative prompt embeddings shall be forced to always be set to 0. Also see the config of357 `stabilityai/stable-diffusion-xl-base-1-0`.358 add_watermarker (`bool`, *optional*):359 Whether to use the [invisible_watermark library](https://github.com/ShieldMnt/invisible-watermark/) to360 watermark output images. If not defined, it will default to True if the package is installed, otherwise no361 watermarker will be used.362 """363 364 model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->unet->vae"365 366 _optional_components = [367 "tokenizer",368 "tokenizer_2",369 "text_encoder",370 "text_encoder_2",371 "image_encoder",372 "feature_extractor",373 ]374 _callback_tensor_inputs = [375 "latents",376 "prompt_embeds",377 "negative_prompt_embeds",378 "add_text_embeds",379 "add_time_ids",380 "negative_pooled_prompt_embeds",381 "add_neg_time_ids",382 "mask",383 "masked_image_latents",384 ]385 386 def __init__(387 self,388 vae: AutoencoderKL,389 text_encoder: CLIPTextModel,390 text_encoder_2: CLIPTextModelWithProjection,391 tokenizer: CLIPTokenizer,392 tokenizer_2: CLIPTokenizer,393 unet: UNet2DConditionModel,394 scheduler: KarrasDiffusionSchedulers,395 image_encoder: CLIPVisionModelWithProjection = None,396 feature_extractor: CLIPImageProcessor = None,397 requires_aesthetics_score: bool = False,398 force_zeros_for_empty_prompt: bool = True,399 ):400 super().__init__()401 402 self.register_modules(403 vae=vae,404 text_encoder=text_encoder,405 text_encoder_2=text_encoder_2,406 tokenizer=tokenizer,407 tokenizer_2=tokenizer_2,408 unet=unet,409 image_encoder=image_encoder,410 feature_extractor=feature_extractor,411 scheduler=scheduler,412 )413 self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt)414 self.register_to_config(requires_aesthetics_score=requires_aesthetics_score)415 self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)416 self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)417 self.mask_processor = VaeImageProcessor(418 vae_scale_factor=self.vae_scale_factor, do_normalize=False, do_binarize=True, do_convert_grayscale=True419 )420 421 422 423 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_vae_slicing424 def enable_vae_slicing(self):425 r"""426 Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to427 compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.428 """429 self.vae.enable_slicing()430 431 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.disable_vae_slicing432 def disable_vae_slicing(self):433 r"""434 Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to435 computing decoding in one step.436 """437 self.vae.disable_slicing()438 439 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_vae_tiling440 def enable_vae_tiling(self):441 r"""442 Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to443 compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow444 processing larger images.445 """446 self.vae.enable_tiling()447 448 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.disable_vae_tiling449 def disable_vae_tiling(self):450 r"""451 Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to452 computing decoding in one step.453 """454 self.vae.disable_tiling()455 456 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image457 def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None):458 dtype = next(self.image_encoder.parameters()).dtype459 # print(image.shape)460 if not isinstance(image, torch.Tensor):461 image = self.feature_extractor(image, return_tensors="pt").pixel_values462 463 image = image.to(device=device, dtype=dtype)464 if output_hidden_states:465 image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]466 image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)467 uncond_image_enc_hidden_states = self.image_encoder(468 torch.zeros_like(image), output_hidden_states=True469 ).hidden_states[-2]470 uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave(471 num_images_per_prompt, dim=0472 )473 return image_enc_hidden_states, uncond_image_enc_hidden_states474 else:475 image_embeds = self.image_encoder(image).image_embeds476 image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)477 uncond_image_embeds = torch.zeros_like(image_embeds)478 479 return image_embeds, uncond_image_embeds480 481 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds482 def prepare_ip_adapter_image_embeds(self, ip_adapter_image, device, num_images_per_prompt):483 # if not isinstance(ip_adapter_image, list):484 # ip_adapter_image = [ip_adapter_image]485 486 # if len(ip_adapter_image) != len(self.unet.encoder_hid_proj.image_projection_layers):487 # raise ValueError(488 # f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters."489 # )490 output_hidden_state = not isinstance(self.unet.encoder_hid_proj, ImageProjection)491 # print(output_hidden_state)492 image_embeds, negative_image_embeds = self.encode_image(493 ip_adapter_image, device, 1, output_hidden_state494 )495 # print(single_image_embeds.shape)496 # single_image_embeds = torch.stack([single_image_embeds] * num_images_per_prompt, dim=0)497 # single_negative_image_embeds = torch.stack([single_negative_image_embeds] * num_images_per_prompt, dim=0)498 # print(single_image_embeds.shape)499 if self.do_classifier_free_guidance:500 image_embeds = torch.cat([negative_image_embeds, image_embeds])501 image_embeds = image_embeds.to(device)502 503 504 return image_embeds505 506 507 # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.encode_prompt508 def encode_prompt(509 self,510 prompt: str,511 prompt_2: Optional[str] = None,512 device: Optional[torch.device] = None,513 num_images_per_prompt: int = 1,514 do_classifier_free_guidance: bool = True,515 negative_prompt: Optional[str] = None,516 negative_prompt_2: Optional[str] = None,517 prompt_embeds: Optional[torch.FloatTensor] = None,518 negative_prompt_embeds: Optional[torch.FloatTensor] = None,519 pooled_prompt_embeds: Optional[torch.FloatTensor] = None,520 negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,521 lora_scale: Optional[float] = None,522 clip_skip: Optional[int] = None,523 ):524 r"""525 Encodes the prompt into text encoder hidden states.526 527 Args:528 prompt (`str` or `List[str]`, *optional*):529 prompt to be encoded530 prompt_2 (`str` or `List[str]`, *optional*):531 The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is532 used in both text-encoders533 device: (`torch.device`):534 torch device535 num_images_per_prompt (`int`):536 number of images that should be generated per prompt537 do_classifier_free_guidance (`bool`):538 whether to use classifier free guidance or not539 negative_prompt (`str` or `List[str]`, *optional*):540 The prompt or prompts not to guide the image generation. If not defined, one has to pass541 `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is542 less than `1`).543 negative_prompt_2 (`str` or `List[str]`, *optional*):544 The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and545 `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders546 prompt_embeds (`torch.FloatTensor`, *optional*):547 Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not548 provided, text embeddings will be generated from `prompt` input argument.549 negative_prompt_embeds (`torch.FloatTensor`, *optional*):550 Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt551 weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input552 argument.553 pooled_prompt_embeds (`torch.FloatTensor`, *optional*):554 Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.555 If not provided, pooled text embeddings will be generated from `prompt` input argument.556 negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):557 Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt558 weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`559 input argument.560 lora_scale (`float`, *optional*):561 A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.562 clip_skip (`int`, *optional*):563 Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that564 the output of the pre-final layer will be used for computing the prompt embeddings.565 """566 device = device or self._execution_device567 568 # set lora scale so that monkey patched LoRA569 # function of text encoder can correctly access it570 if lora_scale is not None and isinstance(self, StableDiffusionXLLoraLoaderMixin):571 self._lora_scale = lora_scale572 573 # dynamically adjust the LoRA scale574 if self.text_encoder is not None:575 if not USE_PEFT_BACKEND:576 adjust_lora_scale_text_encoder(self.text_encoder, lora_scale)577 else:578 scale_lora_layers(self.text_encoder, lora_scale)579 580 if self.text_encoder_2 is not None:581 if not USE_PEFT_BACKEND:582 adjust_lora_scale_text_encoder(self.text_encoder_2, lora_scale)583 else:584 scale_lora_layers(self.text_encoder_2, lora_scale)585 586 prompt = [prompt] if isinstance(prompt, str) else prompt587 588 if prompt is not None:589 batch_size = len(prompt)590 else:591 batch_size = prompt_embeds.shape[0]592 593 # Define tokenizers and text encoders594 tokenizers = [self.tokenizer, self.tokenizer_2] if self.tokenizer is not None else [self.tokenizer_2]595 text_encoders = (596 [self.text_encoder, self.text_encoder_2] if self.text_encoder is not None else [self.text_encoder_2]597 )598 599 if prompt_embeds is None:600 prompt_2 = prompt_2 or prompt601 prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2602 603 # textual inversion: procecss multi-vector tokens if necessary604 prompt_embeds_list = []605 prompts = [prompt, prompt_2]606 for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders):607 if isinstance(self, TextualInversionLoaderMixin):608 prompt = self.maybe_convert_prompt(prompt, tokenizer)609 610 text_inputs = tokenizer(611 prompt,612 padding="max_length",613 max_length=tokenizer.model_max_length,614 truncation=True,615 return_tensors="pt",616 )617 618 text_input_ids = text_inputs.input_ids619 untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids620 621 if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(622 text_input_ids, untruncated_ids623 ):624 removed_text = tokenizer.batch_decode(untruncated_ids[:, tokenizer.model_max_length - 1 : -1])625 logger.warning(626 "The following part of your input was truncated because CLIP can only handle sequences up to"627 f" {tokenizer.model_max_length} tokens: {removed_text}"628 )629 630 prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True)631 632 # We are only ALWAYS interested in the pooled output of the final text encoder633 pooled_prompt_embeds = prompt_embeds[0]634 if clip_skip is None:635 prompt_embeds = prompt_embeds.hidden_states[-2]636 else:637 # "2" because SDXL always indexes from the penultimate layer.638 prompt_embeds = prompt_embeds.hidden_states[-(clip_skip + 2)]639 640 prompt_embeds_list.append(prompt_embeds)641 642 prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)643 644 # get unconditional embeddings for classifier free guidance645 zero_out_negative_prompt = negative_prompt is None and self.config.force_zeros_for_empty_prompt646 if do_classifier_free_guidance and negative_prompt_embeds is None and zero_out_negative_prompt:647 negative_prompt_embeds = torch.zeros_like(prompt_embeds)648 negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds)649 elif do_classifier_free_guidance and negative_prompt_embeds is None:650 negative_prompt = negative_prompt or ""651 negative_prompt_2 = negative_prompt_2 or negative_prompt652 653 # normalize str to list654 negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt655 negative_prompt_2 = (656 batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2657 )658 659 uncond_tokens: List[str]660 if prompt is not None and type(prompt) is not type(negative_prompt):661 raise TypeError(662 f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="663 f" {type(prompt)}."664 )665 elif batch_size != len(negative_prompt):666 raise ValueError(667 f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"668 f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"669 " the batch size of `prompt`."670 )671 else:672 uncond_tokens = [negative_prompt, negative_prompt_2]673 674 negative_prompt_embeds_list = []675 for negative_prompt, tokenizer, text_encoder in zip(uncond_tokens, tokenizers, text_encoders):676 if isinstance(self, TextualInversionLoaderMixin):677 negative_prompt = self.maybe_convert_prompt(negative_prompt, tokenizer)678 679 max_length = prompt_embeds.shape[1]680 uncond_input = tokenizer(681 negative_prompt,682 padding="max_length",683 max_length=max_length,684 truncation=True,685 return_tensors="pt",686 )687 688 negative_prompt_embeds = text_encoder(689 uncond_input.input_ids.to(device),690 output_hidden_states=True,691 )692 # We are only ALWAYS interested in the pooled output of the final text encoder693 negative_pooled_prompt_embeds = negative_prompt_embeds[0]694 negative_prompt_embeds = negative_prompt_embeds.hidden_states[-2]695 696 negative_prompt_embeds_list.append(negative_prompt_embeds)697 698 negative_prompt_embeds = torch.concat(negative_prompt_embeds_list, dim=-1)699 700 if self.text_encoder_2 is not None:701 prompt_embeds = prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device)702 else:703 prompt_embeds = prompt_embeds.to(dtype=self.unet.dtype, device=device)704 705 bs_embed, seq_len, _ = prompt_embeds.shape706 # duplicate text embeddings for each generation per prompt, using mps friendly method707 prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)708 prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)709 710 if do_classifier_free_guidance:711 # duplicate unconditional embeddings for each generation per prompt, using mps friendly method712 seq_len = negative_prompt_embeds.shape[1]713 714 if self.text_encoder_2 is not None:715 negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device)716 else:717 negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.unet.dtype, device=device)718 719 negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)720 negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)721 722 pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(723 bs_embed * num_images_per_prompt, -1724 )725 if do_classifier_free_guidance:726 negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(727 bs_embed * num_images_per_prompt, -1728 )729 730 if self.text_encoder is not None:731 if isinstance(self, StableDiffusionXLLoraLoaderMixin) and USE_PEFT_BACKEND:732 # Retrieve the original scale by scaling back the LoRA layers733 unscale_lora_layers(self.text_encoder, lora_scale)734 735 if self.text_encoder_2 is not None:736 if isinstance(self, StableDiffusionXLLoraLoaderMixin) and USE_PEFT_BACKEND:737 # Retrieve the original scale by scaling back the LoRA layers738 unscale_lora_layers(self.text_encoder_2, lora_scale)739 740 return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds741 742 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs743 def prepare_extra_step_kwargs(self, generator, eta):744 # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature745 # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.746 # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502747 # and should be between [0, 1]748 749 accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())750 extra_step_kwargs = {}751 if accepts_eta:752 extra_step_kwargs["eta"] = eta753 754 # check if the scheduler accepts generator755 accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())756 if accepts_generator:757 extra_step_kwargs["generator"] = generator758 return extra_step_kwargs759 760 def check_inputs(761 self,762 prompt,763 prompt_2,764 image,765 mask_image,766 height,767 width,768 strength,769 callback_steps,770 output_type,771 negative_prompt=None,772 negative_prompt_2=None,773 prompt_embeds=None,774 negative_prompt_embeds=None,775 callback_on_step_end_tensor_inputs=None,776 padding_mask_crop=None,777 ):778 if strength < 0 or strength > 1:779 raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}")780 781 if height % 8 != 0 or width % 8 != 0:782 raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")783 784 if callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0):785 raise ValueError(786 f"`callback_steps` has to be a positive integer but is {callback_steps} of type"787 f" {type(callback_steps)}."788 )789 790 if callback_on_step_end_tensor_inputs is not None and not all(791 k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs792 ):793 raise ValueError(794 f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"795 )796 797 if prompt is not None and prompt_embeds is not None:798 raise ValueError(799 f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"800 " only forward one of the two."801 )802 elif prompt_2 is not None and prompt_embeds is not None:803 raise ValueError(804 f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"805 " only forward one of the two."806 )807 elif prompt is None and prompt_embeds is None:808 raise ValueError(809 "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."810 )811 elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):812 raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")813 elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)):814 raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}")815 816 if negative_prompt is not None and negative_prompt_embeds is not None:817 raise ValueError(818 f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"819 f" {negative_prompt_embeds}. Please make sure to only forward one of the two."820 )821 elif negative_prompt_2 is not None and negative_prompt_embeds is not None:822 raise ValueError(823 f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_prompt_embeds`:"824 f" {negative_prompt_embeds}. Please make sure to only forward one of the two."825 )826 827 if prompt_embeds is not None and negative_prompt_embeds is not None:828 if prompt_embeds.shape != negative_prompt_embeds.shape:829 raise ValueError(830 "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"831 f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"832 f" {negative_prompt_embeds.shape}."833 )834 if padding_mask_crop is not None:835 if not isinstance(image, PIL.Image.Image):836 raise ValueError(837 f"The image should be a PIL image when inpainting mask crop, but is of type" f" {type(image)}."838 )839 if not isinstance(mask_image, PIL.Image.Image):840 raise ValueError(841 f"The mask image should be a PIL image when inpainting mask crop, but is of type"842 f" {type(mask_image)}."843 )844 if output_type != "pil":845 raise ValueError(f"The output type should be PIL when inpainting mask crop, but is" f" {output_type}.")846 847 def prepare_latents(848 self,849 batch_size,850 num_channels_latents,851 height,852 width,853 dtype,854 device,855 generator,856 latents=None,857 image=None,858 timestep=None,859 is_strength_max=True,860 add_noise=True,861 return_noise=False,862 return_image_latents=False,863 ):864 shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)865 if isinstance(generator, list) and len(generator) != batch_size:866 raise ValueError(867 f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"868 f" size of {batch_size}. Make sure the batch size matches the length of the generators."869 )870 871 if (image is None or timestep is None) and not is_strength_max:872 raise ValueError(873 "Since strength < 1. initial latents are to be initialised as a combination of Image + Noise."874 "However, either the image or the noise timestep has not been provided."875 )876 877 if image.shape[1] == 4:878 image_latents = image.to(device=device, dtype=dtype)879 image_latents = image_latents.repeat(batch_size // image_latents.shape[0], 1, 1, 1)880 elif return_image_latents or (latents is None and not is_strength_max):881 image = image.to(device=device, dtype=dtype)882 image_latents = self._encode_vae_image(image=image, generator=generator)883 image_latents = image_latents.repeat(batch_size // image_latents.shape[0], 1, 1, 1)884 885 if latents is None and add_noise:886 noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)887 # if strength is 1. then initialise the latents to noise, else initial to image + noise888 latents = noise if is_strength_max else self.scheduler.add_noise(image_latents, noise, timestep)889 # if pure noise then scale the initial latents by the Scheduler's init sigma890 latents = latents * self.scheduler.init_noise_sigma if is_strength_max else latents891 elif add_noise:892 noise = latents.to(device)893 latents = noise * self.scheduler.init_noise_sigma894 else:895 noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)896 latents = image_latents.to(device)897 898 outputs = (latents,)899 900 if return_noise:901 outputs += (noise,)902 903 if return_image_latents:904 outputs += (image_latents,)905 906 return outputs907 908 def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator):909 dtype = image.dtype910 if self.vae.config.force_upcast:911 image = image.float()912 self.vae.to(dtype=torch.float32)913 914 if isinstance(generator, list):915 image_latents = [916 retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i])917 for i in range(image.shape[0])918 ]919 image_latents = torch.cat(image_latents, dim=0)920 else:921 image_latents = retrieve_latents(self.vae.encode(image), generator=generator)922 923 if self.vae.config.force_upcast:924 self.vae.to(dtype)925 926 image_latents = image_latents.to(dtype)927 image_latents = self.vae.config.scaling_factor * image_latents928 929 return image_latents930 931 def prepare_mask_latents(932 self, mask, masked_image, batch_size, height, width, dtype, device, generator, do_classifier_free_guidance933 ):934 # resize the mask to latents shape as we concatenate the mask to the latents935 # we do that before converting to dtype to avoid breaking in case we're using cpu_offload936 # and half precision937 mask = torch.nn.functional.interpolate(938 mask, size=(height // self.vae_scale_factor, width // self.vae_scale_factor)939 )940 mask = mask.to(device=device, dtype=dtype)941 942 # duplicate mask and masked_image_latents for each generation per prompt, using mps friendly method943 if mask.shape[0] < batch_size:944 if not batch_size % mask.shape[0] == 0:945 raise ValueError(946 "The passed mask and the required batch size don't match. Masks are supposed to be duplicated to"947 f" a total batch size of {batch_size}, but {mask.shape[0]} masks were passed. Make sure the number"948 " of masks that you pass is divisible by the total requested batch size."949 )950 mask = mask.repeat(batch_size // mask.shape[0], 1, 1, 1)951 952 mask = torch.cat([mask] * 2) if do_classifier_free_guidance else mask953 if masked_image is not None and masked_image.shape[1] == 4:954 masked_image_latents = masked_image955 else:956 masked_image_latents = None957 958 if masked_image is not None:959 if masked_image_latents is None:960 masked_image = masked_image.to(device=device, dtype=dtype)961 masked_image_latents = self._encode_vae_image(masked_image, generator=generator)962 963 if masked_image_latents.shape[0] < batch_size:964 if not batch_size % masked_image_latents.shape[0] == 0:965 raise ValueError(966 "The passed images and the required batch size don't match. Images are supposed to be duplicated"967 f" to a total batch size of {batch_size}, but {masked_image_latents.shape[0]} images were passed."968 " Make sure the number of images that you pass is divisible by the total requested batch size."969 )970 masked_image_latents = masked_image_latents.repeat(971 batch_size // masked_image_latents.shape[0], 1, 1, 1972 )973 974 masked_image_latents = (975 torch.cat([masked_image_latents] * 2) if do_classifier_free_guidance else masked_image_latents976 )977 978 # aligning device to prevent device errors when concating it with the latent model input979 masked_image_latents = masked_image_latents.to(device=device, dtype=dtype)980 981 return mask, masked_image_latents982 983 # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline.get_timesteps984 def get_timesteps(self, num_inference_steps, strength, device, denoising_start=None):985 # get the original timestep using init_timestep986 if denoising_start is None:987 init_timestep = min(int(num_inference_steps * strength), num_inference_steps)988 t_start = max(num_inference_steps - init_timestep, 0)989 else:990 t_start = 0991 992 timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :]993 994 # Strength is irrelevant if we directly request a timestep to start at;995 # that is, strength is determined by the denoising_start instead.996 if denoising_start is not None:997 discrete_timestep_cutoff = int(998 round(999 self.scheduler.config.num_train_timesteps1000 - (denoising_start * self.scheduler.config.num_train_timesteps)1001 )1002 )1003 1004 num_inference_steps = (timesteps < discrete_timestep_cutoff).sum().item()1005 if self.scheduler.order == 2 and num_inference_steps % 2 == 0:1006 # if the scheduler is a 2nd order scheduler we might have to do +11007 # because `num_inference_steps` might be even given that every timestep1008 # (except the highest one) is duplicated. If `num_inference_steps` is even it would1009 # mean that we cut the timesteps in the middle of the denoising step1010 # (between 1st and 2nd devirative) which leads to incorrect results. By adding 11011 # we ensure that the denoising process always ends after the 2nd derivate step of the scheduler1012 num_inference_steps = num_inference_steps + 11013 1014 # because t_n+1 >= t_n, we slice the timesteps starting from the end1015 timesteps = timesteps[-num_inference_steps:]1016 return timesteps, num_inference_steps1017 1018 return timesteps, num_inference_steps - t_start1019 1020 # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline._get_add_time_ids1021 def _get_add_time_ids(1022 self,1023 original_size,1024 crops_coords_top_left,1025 target_size,1026 aesthetic_score,1027 negative_aesthetic_score,1028 negative_original_size,1029 negative_crops_coords_top_left,1030 negative_target_size,1031 dtype,1032 text_encoder_projection_dim=None,1033 ):1034 if self.config.requires_aesthetics_score:1035 add_time_ids = list(original_size + crops_coords_top_left + (aesthetic_score,))1036 add_neg_time_ids = list(1037 negative_original_size + negative_crops_coords_top_left + (negative_aesthetic_score,)1038 )1039 else:1040 add_time_ids = list(original_size + crops_coords_top_left + target_size)1041 add_neg_time_ids = list(negative_original_size + crops_coords_top_left + negative_target_size)1042 1043 passed_add_embed_dim = (1044 self.unet.config.addition_time_embed_dim * len(add_time_ids) + text_encoder_projection_dim1045 )1046 expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features1047 1048 if (1049 expected_add_embed_dim > passed_add_embed_dim1050 and (expected_add_embed_dim - passed_add_embed_dim) == self.unet.config.addition_time_embed_dim1051 ):1052 raise ValueError(1053 f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. Please make sure to enable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=True)` to make sure `aesthetic_score` {aesthetic_score} and `negative_aesthetic_score` {negative_aesthetic_score} is correctly used by the model."1054 )1055 elif (1056 expected_add_embed_dim < passed_add_embed_dim1057 and (passed_add_embed_dim - expected_add_embed_dim) == self.unet.config.addition_time_embed_dim1058 ):1059 raise ValueError(1060 f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. Please make sure to disable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=False)` to make sure `target_size` {target_size} is correctly used by the model."1061 )1062 elif expected_add_embed_dim != passed_add_embed_dim:1063 raise ValueError(1064 f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`."1065 )1066 1067 add_time_ids = torch.tensor([add_time_ids], dtype=dtype)1068 add_neg_time_ids = torch.tensor([add_neg_time_ids], dtype=dtype)1069 1070 return add_time_ids, add_neg_time_ids1071 1072 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae1073 def upcast_vae(self):1074 dtype = self.vae.dtype1075 self.vae.to(dtype=torch.float32)1076 use_torch_2_0_or_xformers = isinstance(1077 self.vae.decoder.mid_block.attentions[0].processor,1078 (1079 AttnProcessor2_0,1080 XFormersAttnProcessor,1081 LoRAXFormersAttnProcessor,1082 LoRAAttnProcessor2_0,1083 ),1084 )1085 # if xformers or torch_2_0 is used attention block does not need1086 # to be in float32 which can save lots of memory1087 if use_torch_2_0_or_xformers:1088 self.vae.post_quant_conv.to(dtype)1089 self.vae.decoder.conv_in.to(dtype)1090 self.vae.decoder.mid_block.to(dtype)1091 1092 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.enable_freeu1093 def enable_freeu(self, s1: float, s2: float, b1: float, b2: float):1094 r"""Enables the FreeU mechanism as in https://arxiv.org/abs/2309.11497.1095 1096 The suffixes after the scaling factors represent the stages where they are being applied.1097 1098 Please refer to the [official repository](https://github.com/ChenyangSi/FreeU) for combinations of the values1099 that are known to work well for different pipelines such as Stable Diffusion v1, v2, and Stable Diffusion XL.1100 1101 Args:1102 s1 (`float`):1103 Scaling factor for stage 1 to attenuate the contributions of the skip features. This is done to1104 mitigate "oversmoothing effect" in the enhanced denoising process.1105 s2 (`float`):1106 Scaling factor for stage 2 to attenuate the contributions of the skip features. This is done to1107 mitigate "oversmoothing effect" in the enhanced denoising process.1108 b1 (`float`): Scaling factor for stage 1 to amplify the contributions of backbone features.1109 b2 (`float`): Scaling factor for stage 2 to amplify the contributions of backbone features.1110 """1111 if not hasattr(self, "unet"):1112 raise ValueError("The pipeline must have `unet` for using FreeU.")1113 self.unet.enable_freeu(s1=s1, s2=s2, b1=b1, b2=b2)1114 1115 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.disable_freeu1116 def disable_freeu(self):1117 """Disables the FreeU mechanism if enabled."""1118 self.unet.disable_freeu()1119 1120 # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.fuse_qkv_projections1121 def fuse_qkv_projections(self, unet: bool = True, vae: bool = True):1122 """1123 Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query,1124 key, value) are fused. For cross-attention modules, key and value projection matrices are fused.1125 1126 <Tip warning={true}>1127 1128 This API is 🧪 experimental.1129 1130 </Tip>1131 1132 Args:1133 unet (`bool`, defaults to `True`): To apply fusion on the UNet.1134 vae (`bool`, defaults to `True`): To apply fusion on the VAE.1135 """1136 self.fusing_unet = False1137 self.fusing_vae = False1138 1139 if unet:1140 self.fusing_unet = True1141 self.unet.fuse_qkv_projections()1142 self.unet.set_attn_processor(FusedAttnProcessor2_0())1143 1144 if vae:1145 if not isinstance(self.vae, AutoencoderKL):1146 raise ValueError("`fuse_qkv_projections()` is only supported for the VAE of type `AutoencoderKL`.")1147 1148 self.fusing_vae = True1149 self.vae.fuse_qkv_projections()1150 self.vae.set_attn_processor(FusedAttnProcessor2_0())1151 1152 # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.unfuse_qkv_projections1153 def unfuse_qkv_projections(self, unet: bool = True, vae: bool = True):1154 """Disable QKV projection fusion if enabled.1155 1156 <Tip warning={true}>1157 1158 This API is 🧪 experimental.1159 1160 </Tip>1161 1162 Args:1163 unet (`bool`, defaults to `True`): To apply fusion on the UNet.1164 vae (`bool`, defaults to `True`): To apply fusion on the VAE.1165 1166 """1167 if unet:1168 if not self.fusing_unet:1169 logger.warning("The UNet was not initially fused for QKV projections. Doing nothing.")1170 else:1171 self.unet.unfuse_qkv_projections()1172 self.fusing_unet = False1173 1174 if vae:1175 if not self.fusing_vae:1176 logger.warning("The VAE was not initially fused for QKV projections. Doing nothing.")1177 else:1178 self.vae.unfuse_qkv_projections()1179 self.fusing_vae = False1180 1181 # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding1182 def get_guidance_scale_embedding(self, w, embedding_dim=512, dtype=torch.float32):1183 """1184 See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L2981185 1186 Args:1187 timesteps (`torch.Tensor`):1188 generate embedding vectors at these timesteps1189 embedding_dim (`int`, *optional*, defaults to 512):1190 dimension of the embeddings to generate1191 dtype:1192 data type of the generated embeddings1193 1194 Returns:1195 `torch.FloatTensor`: Embedding vectors with shape `(len(timesteps), embedding_dim)`1196 """1197 assert len(w.shape) == 11198 w = w * 1000.01199 1200 half_dim = embedding_dim // 2