diffusers/community-pipelines-mirror
Community Pipeline Examples For more information about community pipelines, please have a look at this issue. Community pipeline examples consist pipelines that have been added by the community. Please have a look at the following tables to get an overview of all community examples. Click on the Code Example to get a copy-and-paste ready code example that you can try out. If a community pipeline doesn't work as expected, please open an issue and ping the author on it. Please… See the full description on the dataset page: https://huggingface.co/datasets/diffusers/community-pipelines-mirror.
922k
1# Inspired by: https://github.com/Mikubill/sd-webui-controlnet/discussions/1236 and https://github.com/Mikubill/sd-webui-controlnet/discussions/12802import inspect3from typing import Any, Callable, Dict, List, Optional, Tuple, Union4 5import numpy as np6import PIL.Image7import torch8from packaging import version9from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer10 11from diffusers import AutoencoderKL, DiffusionPipeline, UNet2DConditionModel12from diffusers.configuration_utils import FrozenDict, deprecate13from diffusers.image_processor import VaeImageProcessor14from diffusers.loaders import FromSingleFileMixin, IPAdapterMixin, LoraLoaderMixin, TextualInversionLoaderMixin15from diffusers.models.attention import BasicTransformerBlock16from diffusers.models.lora import adjust_lora_scale_text_encoder17from diffusers.models.unets.unet_2d_blocks import CrossAttnDownBlock2D, CrossAttnUpBlock2D, DownBlock2D, UpBlock2D18from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput19from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import rescale_noise_cfg20from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker21from diffusers.schedulers import KarrasDiffusionSchedulers22from diffusers.utils import (23 PIL_INTERPOLATION,24 USE_PEFT_BACKEND,25 logging,26 scale_lora_layers,27 unscale_lora_layers,28)29from diffusers.utils.torch_utils import randn_tensor30 31 32logger = logging.get_logger(__name__) # pylint: disable=invalid-name33 34EXAMPLE_DOC_STRING = """35 Examples:36 ```py37 >>> import torch38 >>> from diffusers import UniPCMultistepScheduler39 >>> from diffusers.utils import load_image40 41 >>> input_image = load_image("https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png")42 43 >>> pipe = StableDiffusionReferencePipeline.from_pretrained(44 "runwayml/stable-diffusion-v1-5",45 safety_checker=None,46 torch_dtype=torch.float1647 ).to('cuda:0')48 49 >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)50 51 >>> result_img = pipe(ref_image=input_image,52 prompt="1girl",53 num_inference_steps=20,54 reference_attn=True,55 reference_adain=True).images[0]56 57 >>> result_img.show()58 ```59"""60 61 62def torch_dfs(model: torch.nn.Module):63 r"""64 Performs a depth-first search on the given PyTorch model and returns a list of all its child modules.65 66 Args:67 model (torch.nn.Module): The PyTorch model to perform the depth-first search on.68 69 Returns:70 list: A list of all child modules of the given model.71 """72 result = [model]73 for child in model.children():74 result += torch_dfs(child)75 return result76 77 78class StableDiffusionReferencePipeline(79 DiffusionPipeline, TextualInversionLoaderMixin, LoraLoaderMixin, IPAdapterMixin, FromSingleFileMixin80):81 r"""82 Pipeline for Stable Diffusion Reference.83 84 This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods85 implemented for all pipelines (downloading, saving, running on a particular device, etc.).86 87 The pipeline also inherits the following loading methods:88 - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings89 - [`~loaders.LoraLoaderMixin.load_lora_weights`] for loading LoRA weights90 - [`~loaders.LoraLoaderMixin.save_lora_weights`] for saving LoRA weights91 - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files92 - [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters93 94 Args:95 vae ([`AutoencoderKL`]):96 Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.97 text_encoder ([`CLIPTextModel`]):98 Frozen text-encoder. Stable Diffusion uses the text portion of99 [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically100 the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.101 tokenizer (`CLIPTokenizer`):102 Tokenizer of class103 [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).104 unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.105 scheduler ([`SchedulerMixin`]):106 A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of107 [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].108 safety_checker ([`StableDiffusionSafetyChecker`]):109 Classification module that estimates whether generated images could be considered offensive or harmful.110 Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details.111 feature_extractor ([`CLIPImageProcessor`]):112 Model that extracts features from generated images to be used as inputs for the `safety_checker`.113 """114 115 _optional_components = ["safety_checker", "feature_extractor"]116 117 def __init__(118 self,119 vae: AutoencoderKL,120 text_encoder: CLIPTextModel,121 tokenizer: CLIPTokenizer,122 unet: UNet2DConditionModel,123 scheduler: KarrasDiffusionSchedulers,124 safety_checker: StableDiffusionSafetyChecker,125 feature_extractor: CLIPImageProcessor,126 requires_safety_checker: bool = True,127 ):128 super().__init__()129 130 if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:131 deprecation_message = (132 f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"133 f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "134 "to update the config accordingly as leaving `steps_offset` might led to incorrect results"135 " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"136 " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"137 " file"138 )139 deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)140 new_config = dict(scheduler.config)141 new_config["steps_offset"] = 1142 scheduler._internal_dict = FrozenDict(new_config)143 144 if hasattr(scheduler.config, "skip_prk_steps") and scheduler.config.skip_prk_steps is False:145 deprecation_message = (146 f"The configuration file of this scheduler: {scheduler} has not set the configuration"147 " `skip_prk_steps`. `skip_prk_steps` should be set to True in the configuration file. Please make"148 " sure to update the config accordingly as not setting `skip_prk_steps` in the config might lead to"149 " incorrect results in future versions. If you have downloaded this checkpoint from the Hugging Face"150 " Hub, it would be very nice if you could open a Pull request for the"151 " `scheduler/scheduler_config.json` file"152 )153 deprecate(154 "skip_prk_steps not set",155 "1.0.0",156 deprecation_message,157 standard_warn=False,158 )159 new_config = dict(scheduler.config)160 new_config["skip_prk_steps"] = True161 scheduler._internal_dict = FrozenDict(new_config)162 163 if safety_checker is None and requires_safety_checker:164 logger.warning(165 f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"166 " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"167 " results in services or applications open to the public. Both the diffusers team and Hugging Face"168 " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"169 " it only for use-cases that involve analyzing network behavior or auditing its results. For more"170 " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."171 )172 173 if safety_checker is not None and feature_extractor is None:174 raise ValueError(175 "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"176 " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."177 )178 179 is_unet_version_less_0_9_0 = hasattr(unet.config, "_diffusers_version") and version.parse(180 version.parse(unet.config._diffusers_version).base_version181 ) < version.parse("0.9.0.dev0")182 is_unet_sample_size_less_64 = hasattr(unet.config, "sample_size") and unet.config.sample_size < 64183 if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64:184 deprecation_message = (185 "The configuration file of the unet has set the default `sample_size` to smaller than"186 " 64 which seems highly unlikely .If you're checkpoint is a fine-tuned version of any of the"187 " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-"188 " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5"189 " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the"190 " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`"191 " in the config might lead to incorrect results in future versions. If you have downloaded this"192 " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for"193 " the `unet/config.json` file"194 )195 deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False)196 new_config = dict(unet.config)197 new_config["sample_size"] = 64198 unet._internal_dict = FrozenDict(new_config)199 # Check shapes, assume num_channels_latents == 4, num_channels_mask == 1, num_channels_masked == 4200 if unet.config.in_channels != 4:201 logger.warning(202 f"You have loaded a UNet with {unet.config.in_channels} input channels, whereas by default,"203 f" {self.__class__} assumes that `pipeline.unet` has 4 input channels: 4 for `num_channels_latents`,"204 ". If you did not intend to modify"205 " this behavior, please check whether you have loaded the right checkpoint."206 )207 208 self.register_modules(209 vae=vae,210 text_encoder=text_encoder,211 tokenizer=tokenizer,212 unet=unet,213 scheduler=scheduler,214 safety_checker=safety_checker,215 feature_extractor=feature_extractor,216 )217 self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)218 self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)219 self.register_to_config(requires_safety_checker=requires_safety_checker)220 221 def _default_height_width(222 self,223 height: Optional[int],224 width: Optional[int],225 image: Union[PIL.Image.Image, torch.Tensor, List[PIL.Image.Image]],226 ) -> Tuple[int, int]:227 r"""228 Calculate the default height and width for the given image.229 230 Args:231 height (int or None): The desired height of the image. If None, the height will be determined based on the input image.232 width (int or None): The desired width of the image. If None, the width will be determined based on the input image.233 image (PIL.Image.Image or torch.Tensor or list[PIL.Image.Image]): The input image or a list of images.234 235 Returns:236 Tuple[int, int]: A tuple containing the calculated height and width.237 238 """239 # NOTE: It is possible that a list of images have different240 # dimensions for each image, so just checking the first image241 # is not _exactly_ correct, but it is simple.242 while isinstance(image, list):243 image = image[0]244 245 if height is None:246 if isinstance(image, PIL.Image.Image):247 height = image.height248 elif isinstance(image, torch.Tensor):249 height = image.shape[2]250 251 height = (height // 8) * 8 # round down to nearest multiple of 8252 253 if width is None:254 if isinstance(image, PIL.Image.Image):255 width = image.width256 elif isinstance(image, torch.Tensor):257 width = image.shape[3]258 259 width = (width // 8) * 8 # round down to nearest multiple of 8260 261 return height, width262 263 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.check_inputs264 def check_inputs(265 self,266 prompt: Optional[Union[str, List[str]]],267 height: int,268 width: int,269 callback_steps: Optional[int],270 negative_prompt: Optional[str] = None,271 prompt_embeds: Optional[torch.Tensor] = None,272 negative_prompt_embeds: Optional[torch.Tensor] = None,273 ip_adapter_image: Optional[torch.Tensor] = None,274 ip_adapter_image_embeds: Optional[torch.Tensor] = None,275 callback_on_step_end_tensor_inputs: Optional[List[str]] = None,276 ) -> None:277 """278 Check the validity of the input arguments for the diffusion model.279 280 Args:281 prompt (Optional[Union[str, List[str]]]): The prompt text or list of prompt texts.282 height (int): The height of the input image.283 width (int): The width of the input image.284 callback_steps (Optional[int]): The number of steps to perform the callback on.285 negative_prompt (Optional[str]): The negative prompt text.286 prompt_embeds (Optional[torch.Tensor]): The prompt embeddings.287 negative_prompt_embeds (Optional[torch.Tensor]): The negative prompt embeddings.288 ip_adapter_image (Optional[torch.Tensor]): The input adapter image.289 ip_adapter_image_embeds (Optional[torch.Tensor]): The input adapter image embeddings.290 callback_on_step_end_tensor_inputs (Optional[List[str]]): The list of tensor inputs to perform the callback on.291 292 Raises:293 ValueError: If `height` or `width` is not divisible by 8.294 ValueError: If `callback_steps` is not a positive integer.295 ValueError: If `callback_on_step_end_tensor_inputs` contains invalid tensor inputs.296 ValueError: If both `prompt` and `prompt_embeds` are provided.297 ValueError: If neither `prompt` nor `prompt_embeds` are provided.298 ValueError: If `prompt` is not of type `str` or `list`.299 ValueError: If both `negative_prompt` and `negative_prompt_embeds` are provided.300 ValueError: If both `prompt_embeds` and `negative_prompt_embeds` are provided and have different shapes.301 ValueError: If both `ip_adapter_image` and `ip_adapter_image_embeds` are provided.302 303 Returns:304 None305 """306 if height % 8 != 0 or width % 8 != 0:307 raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")308 309 if callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0):310 raise ValueError(311 f"`callback_steps` has to be a positive integer but is {callback_steps} of type"312 f" {type(callback_steps)}."313 )314 if callback_on_step_end_tensor_inputs is not None and not all(315 k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs316 ):317 raise ValueError(318 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]}"319 )320 321 if prompt is not None and prompt_embeds is not None:322 raise ValueError(323 f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"324 " only forward one of the two."325 )326 elif prompt is None and prompt_embeds is None:327 raise ValueError(328 "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."329 )330 elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):331 raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")332 333 if negative_prompt is not None and negative_prompt_embeds is not None:334 raise ValueError(335 f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"336 f" {negative_prompt_embeds}. Please make sure to only forward one of the two."337 )338 339 if prompt_embeds is not None and negative_prompt_embeds is not None:340 if prompt_embeds.shape != negative_prompt_embeds.shape:341 raise ValueError(342 "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"343 f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"344 f" {negative_prompt_embeds.shape}."345 )346 347 if ip_adapter_image is not None and ip_adapter_image_embeds is not None:348 raise ValueError(349 "Provide either `ip_adapter_image` or `ip_adapter_image_embeds`. Cannot leave both `ip_adapter_image` and `ip_adapter_image_embeds` defined."350 )351 352 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._encode_prompt353 def _encode_prompt(354 self,355 prompt: Union[str, List[str]],356 device: torch.device,357 num_images_per_prompt: int,358 do_classifier_free_guidance: bool,359 negative_prompt: Optional[Union[str, List[str]]] = None,360 prompt_embeds: Optional[torch.Tensor] = None,361 negative_prompt_embeds: Optional[torch.Tensor] = None,362 lora_scale: Optional[float] = None,363 **kwargs,364 ) -> torch.Tensor:365 r"""366 Encodes the prompt into embeddings.367 368 Args:369 prompt (Union[str, List[str]]): The prompt text or a list of prompt texts.370 device (torch.device): The device to use for encoding.371 num_images_per_prompt (int): The number of images per prompt.372 do_classifier_free_guidance (bool): Whether to use classifier-free guidance.373 negative_prompt (Optional[Union[str, List[str]]], optional): The negative prompt text or a list of negative prompt texts. Defaults to None.374 prompt_embeds (Optional[torch.Tensor], optional): The prompt embeddings. Defaults to None.375 negative_prompt_embeds (Optional[torch.Tensor], optional): The negative prompt embeddings. Defaults to None.376 lora_scale (Optional[float], optional): The LoRA scale. Defaults to None.377 **kwargs: Additional keyword arguments.378 379 Returns:380 torch.Tensor: The encoded prompt embeddings.381 """382 deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."383 deprecate("_encode_prompt()", "1.0.0", deprecation_message, standard_warn=False)384 385 prompt_embeds_tuple = self.encode_prompt(386 prompt=prompt,387 device=device,388 num_images_per_prompt=num_images_per_prompt,389 do_classifier_free_guidance=do_classifier_free_guidance,390 negative_prompt=negative_prompt,391 prompt_embeds=prompt_embeds,392 negative_prompt_embeds=negative_prompt_embeds,393 lora_scale=lora_scale,394 **kwargs,395 )396 397 # concatenate for backwards comp398 prompt_embeds = torch.cat([prompt_embeds_tuple[1], prompt_embeds_tuple[0]])399 400 return prompt_embeds401 402 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_prompt403 def encode_prompt(404 self,405 prompt: Optional[str],406 device: torch.device,407 num_images_per_prompt: int,408 do_classifier_free_guidance: bool,409 negative_prompt: Optional[str] = None,410 prompt_embeds: Optional[torch.Tensor] = None,411 negative_prompt_embeds: Optional[torch.Tensor] = None,412 lora_scale: Optional[float] = None,413 clip_skip: Optional[int] = None,414 ) -> torch.Tensor:415 r"""416 Encodes the prompt into text encoder hidden states.417 418 Args:419 prompt (`str` or `List[str]`, *optional*):420 prompt to be encoded421 device: (`torch.device`):422 torch device423 num_images_per_prompt (`int`):424 number of images that should be generated per prompt425 do_classifier_free_guidance (`bool`):426 whether to use classifier free guidance or not427 negative_prompt (`str` or `List[str]`, *optional*):428 The prompt or prompts not to guide the image generation. If not defined, one has to pass429 `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is430 less than `1`).431 prompt_embeds (`torch.Tensor`, *optional*):432 Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not433 provided, text embeddings will be generated from `prompt` input argument.434 negative_prompt_embeds (`torch.Tensor`, *optional*):435 Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt436 weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input437 argument.438 lora_scale (`float`, *optional*):439 A LoRA scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.440 clip_skip (`int`, *optional*):441 Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that442 the output of the pre-final layer will be used for computing the prompt embeddings.443 """444 # set lora scale so that monkey patched LoRA445 # function of text encoder can correctly access it446 if lora_scale is not None and isinstance(self, LoraLoaderMixin):447 self._lora_scale = lora_scale448 449 # dynamically adjust the LoRA scale450 if not USE_PEFT_BACKEND:451 adjust_lora_scale_text_encoder(self.text_encoder, lora_scale)452 else:453 scale_lora_layers(self.text_encoder, lora_scale)454 455 if prompt is not None and isinstance(prompt, str):456 batch_size = 1457 elif prompt is not None and isinstance(prompt, list):458 batch_size = len(prompt)459 else:460 batch_size = prompt_embeds.shape[0]461 462 if prompt_embeds is None:463 # textual inversion: process multi-vector tokens if necessary464 if isinstance(self, TextualInversionLoaderMixin):465 prompt = self.maybe_convert_prompt(prompt, self.tokenizer)466 467 text_inputs = self.tokenizer(468 prompt,469 padding="max_length",470 max_length=self.tokenizer.model_max_length,471 truncation=True,472 return_tensors="pt",473 )474 text_input_ids = text_inputs.input_ids475 untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids476 477 if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(478 text_input_ids, untruncated_ids479 ):480 removed_text = self.tokenizer.batch_decode(481 untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]482 )483 logger.warning(484 "The following part of your input was truncated because CLIP can only handle sequences up to"485 f" {self.tokenizer.model_max_length} tokens: {removed_text}"486 )487 488 if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:489 attention_mask = text_inputs.attention_mask.to(device)490 else:491 attention_mask = None492 493 if clip_skip is None:494 prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=attention_mask)495 prompt_embeds = prompt_embeds[0]496 else:497 prompt_embeds = self.text_encoder(498 text_input_ids.to(device), attention_mask=attention_mask, output_hidden_states=True499 )500 # Access the `hidden_states` first, that contains a tuple of501 # all the hidden states from the encoder layers. Then index into502 # the tuple to access the hidden states from the desired layer.503 prompt_embeds = prompt_embeds[-1][-(clip_skip + 1)]504 # We also need to apply the final LayerNorm here to not mess with the505 # representations. The `last_hidden_states` that we typically use for506 # obtaining the final prompt representations passes through the LayerNorm507 # layer.508 prompt_embeds = self.text_encoder.text_model.final_layer_norm(prompt_embeds)509 510 if self.text_encoder is not None:511 prompt_embeds_dtype = self.text_encoder.dtype512 elif self.unet is not None:513 prompt_embeds_dtype = self.unet.dtype514 else:515 prompt_embeds_dtype = prompt_embeds.dtype516 517 prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)518 519 bs_embed, seq_len, _ = prompt_embeds.shape520 # duplicate text embeddings for each generation per prompt, using mps friendly method521 prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)522 prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)523 524 # get unconditional embeddings for classifier free guidance525 if do_classifier_free_guidance and negative_prompt_embeds is None:526 uncond_tokens: List[str]527 if negative_prompt is None:528 uncond_tokens = [""] * batch_size529 elif prompt is not None and type(prompt) is not type(negative_prompt):530 raise TypeError(531 f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="532 f" {type(prompt)}."533 )534 elif isinstance(negative_prompt, str):535 uncond_tokens = [negative_prompt]536 elif batch_size != len(negative_prompt):537 raise ValueError(538 f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"539 f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"540 " the batch size of `prompt`."541 )542 else:543 uncond_tokens = negative_prompt544 545 # textual inversion: process multi-vector tokens if necessary546 if isinstance(self, TextualInversionLoaderMixin):547 uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)548 549 max_length = prompt_embeds.shape[1]550 uncond_input = self.tokenizer(551 uncond_tokens,552 padding="max_length",553 max_length=max_length,554 truncation=True,555 return_tensors="pt",556 )557 558 if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:559 attention_mask = uncond_input.attention_mask.to(device)560 else:561 attention_mask = None562 563 negative_prompt_embeds = self.text_encoder(564 uncond_input.input_ids.to(device),565 attention_mask=attention_mask,566 )567 negative_prompt_embeds = negative_prompt_embeds[0]568 569 if do_classifier_free_guidance:570 # duplicate unconditional embeddings for each generation per prompt, using mps friendly method571 seq_len = negative_prompt_embeds.shape[1]572 573 negative_prompt_embeds = negative_prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)574 575 negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)576 negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)577 578 if isinstance(self, LoraLoaderMixin) and USE_PEFT_BACKEND:579 # Retrieve the original scale by scaling back the LoRA layers580 unscale_lora_layers(self.text_encoder, lora_scale)581 582 return prompt_embeds, negative_prompt_embeds583 584 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents585 def prepare_latents(586 self,587 batch_size: int,588 num_channels_latents: int,589 height: int,590 width: int,591 dtype: torch.dtype,592 device: torch.device,593 generator: Union[torch.Generator, List[torch.Generator]],594 latents: Optional[torch.Tensor] = None,595 ) -> torch.Tensor:596 r"""597 Prepare the latent vectors for diffusion.598 599 Args:600 batch_size (int): The number of samples in the batch.601 num_channels_latents (int): The number of channels in the latent vectors.602 height (int): The height of the latent vectors.603 width (int): The width of the latent vectors.604 dtype (torch.dtype): The data type of the latent vectors.605 device (torch.device): The device to place the latent vectors on.606 generator (Union[torch.Generator, List[torch.Generator]]): The generator(s) to use for random number generation.607 latents (Optional[torch.Tensor]): The pre-existing latent vectors. If None, new latent vectors will be generated.608 609 Returns:610 torch.Tensor: The prepared latent vectors.611 """612 shape = (613 batch_size,614 num_channels_latents,615 int(height) // self.vae_scale_factor,616 int(width) // self.vae_scale_factor,617 )618 if isinstance(generator, list) and len(generator) != batch_size:619 raise ValueError(620 f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"621 f" size of {batch_size}. Make sure the batch size matches the length of the generators."622 )623 624 if latents is None:625 latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)626 else:627 latents = latents.to(device)628 629 # scale the initial noise by the standard deviation required by the scheduler630 latents = latents * self.scheduler.init_noise_sigma631 return latents632 633 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs634 def prepare_extra_step_kwargs(635 self, generator: Union[torch.Generator, List[torch.Generator]], eta: float636 ) -> Dict[str, Any]:637 r"""638 Prepare extra keyword arguments for the scheduler step.639 640 Args:641 generator (Union[torch.Generator, List[torch.Generator]]): The generator used for sampling.642 eta (float): The value of eta (η) used with the DDIMScheduler. Should be between 0 and 1.643 644 Returns:645 Dict[str, Any]: A dictionary containing the extra keyword arguments for the scheduler step.646 """647 # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature648 # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.649 # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502650 # and should be between [0, 1]651 652 accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())653 extra_step_kwargs = {}654 if accepts_eta:655 extra_step_kwargs["eta"] = eta656 657 # check if the scheduler accepts generator658 accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())659 if accepts_generator:660 extra_step_kwargs["generator"] = generator661 return extra_step_kwargs662 663 def prepare_image(664 self,665 image: Union[torch.Tensor, PIL.Image.Image, List[Union[torch.Tensor, PIL.Image.Image]]],666 width: int,667 height: int,668 batch_size: int,669 num_images_per_prompt: int,670 device: torch.device,671 dtype: torch.dtype,672 do_classifier_free_guidance: bool = False,673 guess_mode: bool = False,674 ) -> torch.Tensor:675 r"""676 Prepares the input image for processing.677 678 Args:679 image (torch.Tensor or PIL.Image.Image or list): The input image(s).680 width (int): The desired width of the image.681 height (int): The desired height of the image.682 batch_size (int): The batch size for processing.683 num_images_per_prompt (int): The number of images per prompt.684 device (torch.device): The device to use for processing.685 dtype (torch.dtype): The data type of the image.686 do_classifier_free_guidance (bool, optional): Whether to perform classifier-free guidance. Defaults to False.687 guess_mode (bool, optional): Whether to use guess mode. Defaults to False.688 689 Returns:690 torch.Tensor: The prepared image for processing.691 """692 if not isinstance(image, torch.Tensor):693 if isinstance(image, PIL.Image.Image):694 image = [image]695 696 if isinstance(image[0], PIL.Image.Image):697 images = []698 699 for image_ in image:700 image_ = image_.convert("RGB")701 image_ = image_.resize((width, height), resample=PIL_INTERPOLATION["lanczos"])702 image_ = np.array(image_)703 image_ = image_[None, :]704 images.append(image_)705 706 image = images707 708 image = np.concatenate(image, axis=0)709 image = np.array(image).astype(np.float32) / 255.0710 image = (image - 0.5) / 0.5711 image = image.transpose(0, 3, 1, 2)712 image = torch.from_numpy(image)713 elif isinstance(image[0], torch.Tensor):714 image = torch.cat(image, dim=0)715 716 image_batch_size = image.shape[0]717 718 if image_batch_size == 1:719 repeat_by = batch_size720 else:721 # image batch size is the same as prompt batch size722 repeat_by = num_images_per_prompt723 724 image = image.repeat_interleave(repeat_by, dim=0)725 726 image = image.to(device=device, dtype=dtype)727 728 if do_classifier_free_guidance and not guess_mode:729 image = torch.cat([image] * 2)730 731 return image732 733 def prepare_ref_latents(734 self,735 refimage: torch.Tensor,736 batch_size: int,737 dtype: torch.dtype,738 device: torch.device,739 generator: Union[int, List[int]],740 do_classifier_free_guidance: bool,741 ) -> torch.Tensor:742 r"""743 Prepares reference latents for generating images.744 745 Args:746 refimage (torch.Tensor): The reference image.747 batch_size (int): The desired batch size.748 dtype (torch.dtype): The data type of the tensors.749 device (torch.device): The device to perform computations on.750 generator (int or list): The generator index or a list of generator indices.751 do_classifier_free_guidance (bool): Whether to use classifier-free guidance.752 753 Returns:754 torch.Tensor: The prepared reference latents.755 """756 refimage = refimage.to(device=device, dtype=dtype)757 758 # encode the mask image into latents space so we can concatenate it to the latents759 if isinstance(generator, list):760 ref_image_latents = [761 self.vae.encode(refimage[i : i + 1]).latent_dist.sample(generator=generator[i])762 for i in range(batch_size)763 ]764 ref_image_latents = torch.cat(ref_image_latents, dim=0)765 else:766 ref_image_latents = self.vae.encode(refimage).latent_dist.sample(generator=generator)767 ref_image_latents = self.vae.config.scaling_factor * ref_image_latents768 769 # duplicate mask and ref_image_latents for each generation per prompt, using mps friendly method770 if ref_image_latents.shape[0] < batch_size:771 if not batch_size % ref_image_latents.shape[0] == 0:772 raise ValueError(773 "The passed images and the required batch size don't match. Images are supposed to be duplicated"774 f" to a total batch size of {batch_size}, but {ref_image_latents.shape[0]} images were passed."775 " Make sure the number of images that you pass is divisible by the total requested batch size."776 )777 ref_image_latents = ref_image_latents.repeat(batch_size // ref_image_latents.shape[0], 1, 1, 1)778 779 # aligning device to prevent device errors when concating it with the latent model input780 ref_image_latents = ref_image_latents.to(device=device, dtype=dtype)781 return ref_image_latents782 783 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker784 def run_safety_checker(785 self, image: Union[torch.Tensor, PIL.Image.Image], device: torch.device, dtype: torch.dtype786 ) -> Tuple[Union[torch.Tensor, PIL.Image.Image], Optional[bool]]:787 r"""788 Runs the safety checker on the given image.789 790 Args:791 image (Union[torch.Tensor, PIL.Image.Image]): The input image to be checked.792 device (torch.device): The device to run the safety checker on.793 dtype (torch.dtype): The data type of the input image.794 795 Returns:796 (image, has_nsfw_concept) Tuple[Union[torch.Tensor, PIL.Image.Image], Optional[bool]]: A tuple containing the processed image and797 a boolean indicating whether the image has a NSFW (Not Safe for Work) concept.798 """799 if self.safety_checker is None:800 has_nsfw_concept = None801 else:802 if torch.is_tensor(image):803 feature_extractor_input = self.image_processor.postprocess(image, output_type="pil")804 else:805 feature_extractor_input = self.image_processor.numpy_to_pil(image)806 safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device)807 image, has_nsfw_concept = self.safety_checker(808 images=image, clip_input=safety_checker_input.pixel_values.to(dtype)809 )810 return image, has_nsfw_concept811 812 @torch.no_grad()813 def __call__(814 self,815 prompt: Union[str, List[str]] = None,816 ref_image: Union[torch.Tensor, PIL.Image.Image] = None,817 height: Optional[int] = None,818 width: Optional[int] = None,819 num_inference_steps: int = 50,820 guidance_scale: float = 7.5,821 negative_prompt: Optional[Union[str, List[str]]] = None,822 num_images_per_prompt: Optional[int] = 1,823 eta: float = 0.0,824 generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,825 latents: Optional[torch.Tensor] = None,826 prompt_embeds: Optional[torch.Tensor] = None,827 negative_prompt_embeds: Optional[torch.Tensor] = None,828 output_type: Optional[str] = "pil",829 return_dict: bool = True,830 callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,831 callback_steps: int = 1,832 cross_attention_kwargs: Optional[Dict[str, Any]] = None,833 guidance_rescale: float = 0.0,834 attention_auto_machine_weight: float = 1.0,835 gn_auto_machine_weight: float = 1.0,836 style_fidelity: float = 0.5,837 reference_attn: bool = True,838 reference_adain: bool = True,839 ):840 r"""841 Function invoked when calling the pipeline for generation.842 843 Args:844 prompt (`str` or `List[str]`, *optional*):845 The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.846 instead.847 ref_image (`torch.Tensor`, `PIL.Image.Image`):848 The Reference Control input condition. Reference Control uses this input condition to generate guidance to Unet. If849 the type is specified as `torch.Tensor`, it is passed to Reference Control as is. `PIL.Image.Image` can850 also be accepted as an image.851 height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):852 The height in pixels of the generated image.853 width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):854 The width in pixels of the generated image.855 num_inference_steps (`int`, *optional*, defaults to 50):856 The number of denoising steps. More denoising steps usually lead to a higher quality image at the857 expense of slower inference.858 guidance_scale (`float`, *optional*, defaults to 7.5):859 Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).860 `guidance_scale` is defined as `w` of equation 2. of [Imagen861 Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >862 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,863 usually at the expense of lower image quality.864 negative_prompt (`str` or `List[str]`, *optional*):865 The prompt or prompts not to guide the image generation. If not defined, one has to pass866 `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is867 less than `1`).868 num_images_per_prompt (`int`, *optional*, defaults to 1):869 The number of images to generate per prompt.870 eta (`float`, *optional*, defaults to 0.0):871 Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to872 [`schedulers.DDIMScheduler`], will be ignored for others.873 generator (`torch.Generator` or `List[torch.Generator]`, *optional*):874 One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)875 to make generation deterministic.876 latents (`torch.Tensor`, *optional*):877 Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image878 generation. Can be used to tweak the same generation with different prompts. If not provided, a latents879 tensor will ge generated by sampling using the supplied random `generator`.880 prompt_embeds (`torch.Tensor`, *optional*):881 Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not882 provided, text embeddings will be generated from `prompt` input argument.883 negative_prompt_embeds (`torch.Tensor`, *optional*):884 Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt885 weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input886 argument.887 output_type (`str`, *optional*, defaults to `"pil"`):888 The output format of the generate image. Choose between889 [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.890 return_dict (`bool`, *optional*, defaults to `True`):891 Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a892 plain tuple.893 callback (`Callable`, *optional*):894 A function that will be called every `callback_steps` steps during inference. The function will be895 called with the following arguments: `callback(step: int, timestep: int, latents: torch.Tensor)`.896 callback_steps (`int`, *optional*, defaults to 1):897 The frequency at which the `callback` function will be called. If not specified, the callback will be898 called at every step.899 cross_attention_kwargs (`dict`, *optional*):900 A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under901 `self.processor` in902 [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).903 guidance_rescale (`float`, *optional*, defaults to 0.0):904 Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are905 Flawed](https://arxiv.org/pdf/2305.08891.pdf) `guidance_scale` is defined as `φ` in equation 16. of906 [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf).907 Guidance rescale factor should fix overexposure when using zero terminal SNR.908 attention_auto_machine_weight (`float`):909 Weight of using reference query for self attention's context.910 If attention_auto_machine_weight=1.0, use reference query for all self attention's context.911 gn_auto_machine_weight (`float`):912 Weight of using reference adain. If gn_auto_machine_weight=2.0, use all reference adain plugins.913 style_fidelity (`float`):914 style fidelity of ref_uncond_xt. If style_fidelity=1.0, control more important,915 elif style_fidelity=0.0, prompt more important, else balanced.916 reference_attn (`bool`):917 Whether to use reference query for self attention's context.918 reference_adain (`bool`):919 Whether to use reference adain.920 921 Examples:922 923 Returns:924 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:925 [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.926 When returning a tuple, the first element is a list with the generated images, and the second element is a927 list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"928 (nsfw) content, according to the `safety_checker`.929 """930 assert reference_attn or reference_adain, "`reference_attn` or `reference_adain` must be True."931 932 # 0. Default height and width to unet933 height, width = self._default_height_width(height, width, ref_image)934 935 # 1. Check inputs. Raise error if not correct936 self.check_inputs(937 prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds938 )939 940 # 2. Define call parameters941 if prompt is not None and isinstance(prompt, str):942 batch_size = 1943 elif prompt is not None and isinstance(prompt, list):944 batch_size = len(prompt)945 else:946 batch_size = prompt_embeds.shape[0]947 948 device = self._execution_device949 # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)950 # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`951 # corresponds to doing no classifier free guidance.952 do_classifier_free_guidance = guidance_scale > 1.0953 954 # 3. Encode input prompt955 text_encoder_lora_scale = (956 cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None957 )958 prompt_embeds = self._encode_prompt(959 prompt,960 device,961 num_images_per_prompt,962 do_classifier_free_guidance,963 negative_prompt,964 prompt_embeds=prompt_embeds,965 negative_prompt_embeds=negative_prompt_embeds,966 lora_scale=text_encoder_lora_scale,967 )968 969 # 4. Preprocess reference image970 ref_image = self.prepare_image(971 image=ref_image,972 width=width,973 height=height,974 batch_size=batch_size * num_images_per_prompt,975 num_images_per_prompt=num_images_per_prompt,976 device=device,977 dtype=prompt_embeds.dtype,978 )979 980 # 5. Prepare timesteps981 self.scheduler.set_timesteps(num_inference_steps, device=device)982 timesteps = self.scheduler.timesteps983 984 # 6. Prepare latent variables985 num_channels_latents = self.unet.config.in_channels986 latents = self.prepare_latents(987 batch_size * num_images_per_prompt,988 num_channels_latents,989 height,990 width,991 prompt_embeds.dtype,992 device,993 generator,994 latents,995 )996 997 # 7. Prepare reference latent variables998 ref_image_latents = self.prepare_ref_latents(999 ref_image,1000 batch_size * num_images_per_prompt,1001 prompt_embeds.dtype,1002 device,1003 generator,1004 do_classifier_free_guidance,1005 )1006 1007 # 8. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline1008 extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)1009 1010 # 9. Modify self attention and group norm1011 MODE = "write"1012 uc_mask = (1013 torch.Tensor([1] * batch_size * num_images_per_prompt + [0] * batch_size * num_images_per_prompt)1014 .type_as(ref_image_latents)1015 .bool()1016 )1017 1018 def hacked_basic_transformer_inner_forward(1019 self,1020 hidden_states: torch.Tensor,1021 attention_mask: Optional[torch.Tensor] = None,1022 encoder_hidden_states: Optional[torch.Tensor] = None,1023 encoder_attention_mask: Optional[torch.Tensor] = None,1024 timestep: Optional[torch.LongTensor] = None,1025 cross_attention_kwargs: Dict[str, Any] = None,1026 class_labels: Optional[torch.LongTensor] = None,1027 ):1028 if self.use_ada_layer_norm:1029 norm_hidden_states = self.norm1(hidden_states, timestep)1030 elif self.use_ada_layer_norm_zero:1031 norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(1032 hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype1033 )1034 else:1035 norm_hidden_states = self.norm1(hidden_states)1036 1037 # 1. Self-Attention1038 cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}1039 if self.only_cross_attention:1040 attn_output = self.attn1(1041 norm_hidden_states,1042 encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,1043 attention_mask=attention_mask,1044 **cross_attention_kwargs,1045 )1046 else:1047 if MODE == "write":1048 self.bank.append(norm_hidden_states.detach().clone())1049 attn_output = self.attn1(1050 norm_hidden_states,1051 encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,1052 attention_mask=attention_mask,1053 **cross_attention_kwargs,1054 )1055 if MODE == "read":1056 if attention_auto_machine_weight > self.attn_weight:1057 attn_output_uc = self.attn1(1058 norm_hidden_states,1059 encoder_hidden_states=torch.cat([norm_hidden_states] + self.bank, dim=1),1060 # attention_mask=attention_mask,1061 **cross_attention_kwargs,1062 )1063 attn_output_c = attn_output_uc.clone()1064 if do_classifier_free_guidance and style_fidelity > 0:1065 attn_output_c[uc_mask] = self.attn1(1066 norm_hidden_states[uc_mask],1067 encoder_hidden_states=norm_hidden_states[uc_mask],1068 **cross_attention_kwargs,1069 )1070 attn_output = style_fidelity * attn_output_c + (1.0 - style_fidelity) * attn_output_uc1071 self.bank.clear()1072 else:1073 attn_output = self.attn1(1074 norm_hidden_states,1075 encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,1076 attention_mask=attention_mask,1077 **cross_attention_kwargs,1078 )1079 if self.use_ada_layer_norm_zero:1080 attn_output = gate_msa.unsqueeze(1) * attn_output1081 hidden_states = attn_output + hidden_states1082 1083 if self.attn2 is not None:1084 norm_hidden_states = (1085 self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)1086 )1087 1088 # 2. Cross-Attention1089 attn_output = self.attn2(1090 norm_hidden_states,1091 encoder_hidden_states=encoder_hidden_states,1092 attention_mask=encoder_attention_mask,1093 **cross_attention_kwargs,1094 )1095 hidden_states = attn_output + hidden_states1096 1097 # 3. Feed-forward1098 norm_hidden_states = self.norm3(hidden_states)1099 1100 if self.use_ada_layer_norm_zero:1101 norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]1102 1103 ff_output = self.ff(norm_hidden_states)1104 1105 if self.use_ada_layer_norm_zero:1106 ff_output = gate_mlp.unsqueeze(1) * ff_output1107 1108 hidden_states = ff_output + hidden_states1109 1110 return hidden_states1111 1112 def hacked_mid_forward(self, *args, **kwargs):1113 eps = 1e-61114 x = self.original_forward(*args, **kwargs)1115 if MODE == "write":1116 if gn_auto_machine_weight >= self.gn_weight:1117 var, mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0)1118 self.mean_bank.append(mean)1119 self.var_bank.append(var)1120 if MODE == "read":1121 if len(self.mean_bank) > 0 and len(self.var_bank) > 0:1122 var, mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0)1123 std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.51124 mean_acc = sum(self.mean_bank) / float(len(self.mean_bank))1125 var_acc = sum(self.var_bank) / float(len(self.var_bank))1126 std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.51127 x_uc = (((x - mean) / std) * std_acc) + mean_acc1128 x_c = x_uc.clone()1129 if do_classifier_free_guidance and style_fidelity > 0:1130 x_c[uc_mask] = x[uc_mask]1131 x = style_fidelity * x_c + (1.0 - style_fidelity) * x_uc1132 self.mean_bank = []1133 self.var_bank = []1134 return x1135 1136 def hack_CrossAttnDownBlock2D_forward(1137 self,1138 hidden_states: torch.Tensor,1139 temb: Optional[torch.Tensor] = None,1140 encoder_hidden_states: Optional[torch.Tensor] = None,1141 attention_mask: Optional[torch.Tensor] = None,1142 cross_attention_kwargs: Optional[Dict[str, Any]] = None,1143 encoder_attention_mask: Optional[torch.Tensor] = None,1144 ):1145 eps = 1e-61146 1147 # TODO(Patrick, William) - attention mask is not used1148 output_states = ()1149 1150 for i, (resnet, attn) in enumerate(zip(self.resnets, self.attentions)):1151 hidden_states = resnet(hidden_states, temb)1152 hidden_states = attn(1153 hidden_states,1154 encoder_hidden_states=encoder_hidden_states,1155 cross_attention_kwargs=cross_attention_kwargs,1156 attention_mask=attention_mask,1157 encoder_attention_mask=encoder_attention_mask,1158 return_dict=False,1159 )[0]1160 if MODE == "write":1161 if gn_auto_machine_weight >= self.gn_weight:1162 var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)1163 self.mean_bank.append([mean])1164 self.var_bank.append([var])1165 if MODE == "read":1166 if len(self.mean_bank) > 0 and len(self.var_bank) > 0:1167 var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)1168 std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.51169 mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))1170 var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))1171 std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.51172 hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc1173 hidden_states_c = hidden_states_uc.clone()1174 if do_classifier_free_guidance and style_fidelity > 0:1175 hidden_states_c[uc_mask] = hidden_states[uc_mask]1176 hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc1177 1178 output_states = output_states + (hidden_states,)1179 1180 if MODE == "read":1181 self.mean_bank = []1182 self.var_bank = []1183 1184 if self.downsamplers is not None:1185 for downsampler in self.downsamplers:1186 hidden_states = downsampler(hidden_states)1187 1188 output_states = output_states + (hidden_states,)1189 1190 return hidden_states, output_states1191 1192 def hacked_DownBlock2D_forward(1193 self,1194 hidden_states: torch.Tensor,1195 temb: Optional[torch.Tensor] = None,1196 **kwargs: Any,1197 ) -> Tuple[torch.Tensor, ...]:1198 eps = 1e-61199 1200 output_states = ()