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# Based on stable_diffusion_reference.py2 3import inspect4from typing import Any, Callable, Dict, List, Optional, Tuple, Union5 6import numpy as np7import PIL.Image8import torch9 10from diffusers import StableDiffusionXLPipeline11from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback12from diffusers.image_processor import PipelineImageInput13from diffusers.models.attention import BasicTransformerBlock14from diffusers.models.unets.unet_2d_blocks import CrossAttnDownBlock2D, CrossAttnUpBlock2D, DownBlock2D, UpBlock2D15from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput16from diffusers.utils import PIL_INTERPOLATION, deprecate, is_torch_xla_available, logging, replace_example_docstring17from diffusers.utils.torch_utils import randn_tensor18 19 20if is_torch_xla_available():21 import torch_xla.core.xla_model as xm # type: ignore22 23 XLA_AVAILABLE = True24else:25 XLA_AVAILABLE = False26 27 28logger = logging.get_logger(__name__) # pylint: disable=invalid-name29 30EXAMPLE_DOC_STRING = """31 Examples:32 ```py33 >>> import torch34 >>> from diffusers.schedulers import UniPCMultistepScheduler35 >>> from diffusers.utils import load_image36 37 >>> input_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/sdxl_reference_input_cat.jpg")38 39 >>> pipe = StableDiffusionXLReferencePipeline.from_pretrained(40 "stabilityai/stable-diffusion-xl-base-1.0",41 torch_dtype=torch.float16,42 use_safetensors=True,43 variant="fp16").to('cuda:0')44 45 >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)46 >>> result_img = pipe(ref_image=input_image,47 prompt="a dog",48 num_inference_steps=20,49 reference_attn=True,50 reference_adain=True).images[0]51 52 >>> result_img.show()53 ```54"""55 56 57def torch_dfs(model: torch.nn.Module):58 result = [model]59 for child in model.children():60 result += torch_dfs(child)61 return result62 63 64# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg65def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):66 """67 Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and68 Sample Steps are Flawed](https://huggingface.co/papers/2305.08891). See Section 3.469 """70 std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)71 std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)72 # rescale the results from guidance (fixes overexposure)73 noise_pred_rescaled = noise_cfg * (std_text / std_cfg)74 # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images75 noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg76 return noise_cfg77 78 79# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps80def retrieve_timesteps(81 scheduler,82 num_inference_steps: Optional[int] = None,83 device: Optional[Union[str, torch.device]] = None,84 timesteps: Optional[List[int]] = None,85 sigmas: Optional[List[float]] = None,86 **kwargs,87):88 r"""89 Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles90 custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.91 92 Args:93 scheduler (`SchedulerMixin`):94 The scheduler to get timesteps from.95 num_inference_steps (`int`):96 The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`97 must be `None`.98 device (`str` or `torch.device`, *optional*):99 The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.100 timesteps (`List[int]`, *optional*):101 Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,102 `num_inference_steps` and `sigmas` must be `None`.103 sigmas (`List[float]`, *optional*):104 Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,105 `num_inference_steps` and `timesteps` must be `None`.106 107 Returns:108 `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the109 second element is the number of inference steps.110 """111 if timesteps is not None and sigmas is not None:112 raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")113 if timesteps is not None:114 accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())115 if not accepts_timesteps:116 raise ValueError(117 f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"118 f" timestep schedules. Please check whether you are using the correct scheduler."119 )120 scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)121 timesteps = scheduler.timesteps122 num_inference_steps = len(timesteps)123 elif sigmas is not None:124 accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())125 if not accept_sigmas:126 raise ValueError(127 f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"128 f" sigmas schedules. Please check whether you are using the correct scheduler."129 )130 scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)131 timesteps = scheduler.timesteps132 num_inference_steps = len(timesteps)133 else:134 scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)135 timesteps = scheduler.timesteps136 return timesteps, num_inference_steps137 138 139class StableDiffusionXLReferencePipeline(StableDiffusionXLPipeline):140 def prepare_ref_latents(self, refimage, batch_size, dtype, device, generator, do_classifier_free_guidance):141 refimage = refimage.to(device=device)142 needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast143 if needs_upcasting:144 self.upcast_vae()145 refimage = refimage.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)146 if refimage.dtype != self.vae.dtype:147 refimage = refimage.to(dtype=self.vae.dtype)148 # encode the mask image into latents space so we can concatenate it to the latents149 if isinstance(generator, list):150 ref_image_latents = [151 self.vae.encode(refimage[i : i + 1]).latent_dist.sample(generator=generator[i])152 for i in range(batch_size)153 ]154 ref_image_latents = torch.cat(ref_image_latents, dim=0)155 else:156 ref_image_latents = self.vae.encode(refimage).latent_dist.sample(generator=generator)157 ref_image_latents = self.vae.config.scaling_factor * ref_image_latents158 159 # duplicate mask and ref_image_latents for each generation per prompt, using mps friendly method160 if ref_image_latents.shape[0] < batch_size:161 if not batch_size % ref_image_latents.shape[0] == 0:162 raise ValueError(163 "The passed images and the required batch size don't match. Images are supposed to be duplicated"164 f" to a total batch size of {batch_size}, but {ref_image_latents.shape[0]} images were passed."165 " Make sure the number of images that you pass is divisible by the total requested batch size."166 )167 ref_image_latents = ref_image_latents.repeat(batch_size // ref_image_latents.shape[0], 1, 1, 1)168 169 ref_image_latents = torch.cat([ref_image_latents] * 2) if do_classifier_free_guidance else ref_image_latents170 171 # aligning device to prevent device errors when concating it with the latent model input172 ref_image_latents = ref_image_latents.to(device=device, dtype=dtype)173 174 # cast back to fp16 if needed175 if needs_upcasting:176 self.vae.to(dtype=torch.float16)177 178 return ref_image_latents179 180 def prepare_ref_image(181 self,182 image,183 width,184 height,185 batch_size,186 num_images_per_prompt,187 device,188 dtype,189 do_classifier_free_guidance=False,190 guess_mode=False,191 ):192 if not isinstance(image, torch.Tensor):193 if isinstance(image, PIL.Image.Image):194 image = [image]195 196 if isinstance(image[0], PIL.Image.Image):197 images = []198 199 for image_ in image:200 image_ = image_.convert("RGB")201 image_ = image_.resize((width, height), resample=PIL_INTERPOLATION["lanczos"])202 image_ = np.array(image_)203 image_ = image_[None, :]204 images.append(image_)205 206 image = images207 208 image = np.concatenate(image, axis=0)209 image = np.array(image).astype(np.float32) / 255.0210 image = (image - 0.5) / 0.5211 image = image.transpose(0, 3, 1, 2)212 image = torch.from_numpy(image)213 214 elif isinstance(image[0], torch.Tensor):215 image = torch.stack(image, dim=0)216 217 image_batch_size = image.shape[0]218 219 if image_batch_size == 1:220 repeat_by = batch_size221 else:222 repeat_by = num_images_per_prompt223 224 image = image.repeat_interleave(repeat_by, dim=0)225 226 image = image.to(device=device, dtype=dtype)227 228 if do_classifier_free_guidance and not guess_mode:229 image = torch.cat([image] * 2)230 231 return image232 233 def check_ref_inputs(234 self,235 ref_image,236 reference_guidance_start,237 reference_guidance_end,238 style_fidelity,239 reference_attn,240 reference_adain,241 ):242 ref_image_is_pil = isinstance(ref_image, PIL.Image.Image)243 ref_image_is_tensor = isinstance(ref_image, torch.Tensor)244 245 if not ref_image_is_pil and not ref_image_is_tensor:246 raise TypeError(247 f"ref image must be passed and be one of PIL image or torch tensor, but is {type(ref_image)}"248 )249 250 if not reference_attn and not reference_adain:251 raise ValueError("`reference_attn` or `reference_adain` must be True.")252 253 if style_fidelity < 0.0:254 raise ValueError(f"style fidelity: {style_fidelity} can't be smaller than 0.")255 if style_fidelity > 1.0:256 raise ValueError(f"style fidelity: {style_fidelity} can't be larger than 1.0.")257 258 if reference_guidance_start >= reference_guidance_end:259 raise ValueError(260 f"reference guidance start: {reference_guidance_start} cannot be larger or equal to reference guidance end: {reference_guidance_end}."261 )262 if reference_guidance_start < 0.0:263 raise ValueError(f"reference guidance start: {reference_guidance_start} can't be smaller than 0.")264 if reference_guidance_end > 1.0:265 raise ValueError(f"reference guidance end: {reference_guidance_end} can't be larger than 1.0.")266 267 @torch.no_grad()268 @replace_example_docstring(EXAMPLE_DOC_STRING)269 def __call__(270 self,271 prompt: Union[str, List[str]] = None,272 prompt_2: Optional[Union[str, List[str]]] = None,273 ref_image: Union[torch.Tensor, PIL.Image.Image] = None,274 height: Optional[int] = None,275 width: Optional[int] = None,276 num_inference_steps: int = 50,277 timesteps: List[int] = None,278 sigmas: List[float] = None,279 denoising_end: Optional[float] = None,280 guidance_scale: float = 5.0,281 negative_prompt: Optional[Union[str, List[str]]] = None,282 negative_prompt_2: Optional[Union[str, List[str]]] = None,283 num_images_per_prompt: Optional[int] = 1,284 eta: float = 0.0,285 generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,286 latents: Optional[torch.Tensor] = None,287 prompt_embeds: Optional[torch.Tensor] = None,288 negative_prompt_embeds: Optional[torch.Tensor] = None,289 pooled_prompt_embeds: Optional[torch.Tensor] = None,290 negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,291 ip_adapter_image: Optional[PipelineImageInput] = None,292 ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,293 output_type: Optional[str] = "pil",294 return_dict: bool = True,295 cross_attention_kwargs: Optional[Dict[str, Any]] = None,296 guidance_rescale: float = 0.0,297 original_size: Optional[Tuple[int, int]] = None,298 crops_coords_top_left: Tuple[int, int] = (0, 0),299 target_size: Optional[Tuple[int, int]] = None,300 negative_original_size: Optional[Tuple[int, int]] = None,301 negative_crops_coords_top_left: Tuple[int, int] = (0, 0),302 negative_target_size: Optional[Tuple[int, int]] = None,303 clip_skip: Optional[int] = None,304 callback_on_step_end: Optional[305 Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]306 ] = None,307 callback_on_step_end_tensor_inputs: List[str] = ["latents"],308 attention_auto_machine_weight: float = 1.0,309 gn_auto_machine_weight: float = 1.0,310 reference_guidance_start: float = 0.0,311 reference_guidance_end: float = 1.0,312 style_fidelity: float = 0.5,313 reference_attn: bool = True,314 reference_adain: bool = True,315 **kwargs,316 ):317 r"""318 Function invoked when calling the pipeline for generation.319 320 Args:321 prompt (`str` or `List[str]`, *optional*):322 The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.323 instead.324 prompt_2 (`str` or `List[str]`, *optional*):325 The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is326 used in both text-encoders327 ref_image (`torch.Tensor`, `PIL.Image.Image`):328 The Reference Control input condition. Reference Control uses this input condition to generate guidance to Unet. If329 the type is specified as `Torch.Tensor`, it is passed to Reference Control as is. `PIL.Image.Image` can330 also be accepted as an image.331 height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):332 The height in pixels of the generated image. This is set to 1024 by default for the best results.333 Anything below 512 pixels won't work well for334 [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)335 and checkpoints that are not specifically fine-tuned on low resolutions.336 width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):337 The width in pixels of the generated image. This is set to 1024 by default for the best results.338 Anything below 512 pixels won't work well for339 [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)340 and checkpoints that are not specifically fine-tuned on low resolutions.341 num_inference_steps (`int`, *optional*, defaults to 50):342 The number of denoising steps. More denoising steps usually lead to a higher quality image at the343 expense of slower inference.344 timesteps (`List[int]`, *optional*):345 Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument346 in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is347 passed will be used. Must be in descending order.348 sigmas (`List[float]`, *optional*):349 Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in350 their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed351 will be used.352 denoising_end (`float`, *optional*):353 When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be354 completed before it is intentionally prematurely terminated. As a result, the returned sample will355 still retain a substantial amount of noise as determined by the discrete timesteps selected by the356 scheduler. The denoising_end parameter should ideally be utilized when this pipeline forms a part of a357 "Mixture of Denoisers" multi-pipeline setup, as elaborated in [**Refining the Image358 Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output)359 guidance_scale (`float`, *optional*, defaults to 5.0):360 Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://huggingface.co/papers/2207.12598).361 `guidance_scale` is defined as `w` of equation 2. of [Imagen362 Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale >363 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,364 usually at the expense of lower image quality.365 negative_prompt (`str` or `List[str]`, *optional*):366 The prompt or prompts not to guide the image generation. If not defined, one has to pass367 `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is368 less than `1`).369 negative_prompt_2 (`str` or `List[str]`, *optional*):370 The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and371 `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders372 num_images_per_prompt (`int`, *optional*, defaults to 1):373 The number of images to generate per prompt.374 eta (`float`, *optional*, defaults to 0.0):375 Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies to376 [`schedulers.DDIMScheduler`], will be ignored for others.377 generator (`torch.Generator` or `List[torch.Generator]`, *optional*):378 One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)379 to make generation deterministic.380 latents (`torch.Tensor`, *optional*):381 Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image382 generation. Can be used to tweak the same generation with different prompts. If not provided, a latents383 tensor will ge generated by sampling using the supplied random `generator`.384 prompt_embeds (`torch.Tensor`, *optional*):385 Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not386 provided, text embeddings will be generated from `prompt` input argument.387 negative_prompt_embeds (`torch.Tensor`, *optional*):388 Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt389 weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input390 argument.391 pooled_prompt_embeds (`torch.Tensor`, *optional*):392 Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.393 If not provided, pooled text embeddings will be generated from `prompt` input argument.394 negative_pooled_prompt_embeds (`torch.Tensor`, *optional*):395 Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt396 weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`397 input argument.398 ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.399 ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):400 Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of401 IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should402 contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not403 provided, embeddings are computed from the `ip_adapter_image` input argument.404 output_type (`str`, *optional*, defaults to `"pil"`):405 The output format of the generate image. Choose between406 [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.407 return_dict (`bool`, *optional*, defaults to `True`):408 Whether or not to return a [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] instead409 of a plain tuple.410 cross_attention_kwargs (`dict`, *optional*):411 A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under412 `self.processor` in413 [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).414 guidance_rescale (`float`, *optional*, defaults to 0.0):415 Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are416 Flawed](https://huggingface.co/papers/2305.08891) `guidance_scale` is defined as `φ` in equation 16. of417 [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891).418 Guidance rescale factor should fix overexposure when using zero terminal SNR.419 original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):420 If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.421 `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as422 explained in section 2.2 of423 [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).424 crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):425 `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position426 `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting427 `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of428 [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).429 target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):430 For most cases, `target_size` should be set to the desired height and width of the generated image. If431 not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in432 section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).433 negative_original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):434 To negatively condition the generation process based on a specific image resolution. Part of SDXL's435 micro-conditioning as explained in section 2.2 of436 [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more437 information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.438 negative_crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):439 To negatively condition the generation process based on a specific crop coordinates. Part of SDXL's440 micro-conditioning as explained in section 2.2 of441 [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more442 information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.443 negative_target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):444 To negatively condition the generation process based on a target image resolution. It should be as same445 as the `target_size` for most cases. Part of SDXL's micro-conditioning as explained in section 2.2 of446 [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). For more447 information, refer to this issue thread: https://github.com/huggingface/diffusers/issues/4208.448 callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):449 A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of450 each denoising step during the inference. with the following arguments: `callback_on_step_end(self:451 DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a452 list of all tensors as specified by `callback_on_step_end_tensor_inputs`.453 callback_on_step_end_tensor_inputs (`List`, *optional*):454 The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list455 will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the456 `._callback_tensor_inputs` attribute of your pipeline class.457 attention_auto_machine_weight (`float`):458 Weight of using reference query for self attention's context.459 If attention_auto_machine_weight=1.0, use reference query for all self attention's context.460 gn_auto_machine_weight (`float`):461 Weight of using reference adain. If gn_auto_machine_weight=2.0, use all reference adain plugins.462 reference_guidance_start (`float`, *optional*, defaults to 0.0):463 The percentage of total steps at which the reference ControlNet starts applying.464 reference_guidance_end (`float`, *optional*, defaults to 1.0):465 The percentage of total steps at which the reference ControlNet stops applying.466 style_fidelity (`float`):467 style fidelity of ref_uncond_xt. If style_fidelity=1.0, control more important,468 elif style_fidelity=0.0, prompt more important, else balanced.469 reference_attn (`bool`):470 Whether to use reference query for self attention's context.471 reference_adain (`bool`):472 Whether to use reference adain.473 474 Examples:475 476 Returns:477 [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] or `tuple`:478 [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a479 `tuple`. When returning a tuple, the first element is a list with the generated images.480 """481 482 callback = kwargs.pop("callback", None)483 callback_steps = kwargs.pop("callback_steps", None)484 485 if callback is not None:486 deprecate(487 "callback",488 "1.0.0",489 "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",490 )491 if callback_steps is not None:492 deprecate(493 "callback_steps",494 "1.0.0",495 "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",496 )497 498 if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):499 callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs500 501 # 0. Default height and width to unet502 height = height or self.default_sample_size * self.vae_scale_factor503 width = width or self.default_sample_size * self.vae_scale_factor504 505 original_size = original_size or (height, width)506 target_size = target_size or (height, width)507 508 # 1. Check inputs. Raise error if not correct509 self.check_inputs(510 prompt,511 prompt_2,512 height,513 width,514 callback_steps,515 negative_prompt,516 negative_prompt_2,517 prompt_embeds,518 negative_prompt_embeds,519 pooled_prompt_embeds,520 negative_pooled_prompt_embeds,521 ip_adapter_image,522 ip_adapter_image_embeds,523 callback_on_step_end_tensor_inputs,524 )525 526 self.check_ref_inputs(527 ref_image,528 reference_guidance_start,529 reference_guidance_end,530 style_fidelity,531 reference_attn,532 reference_adain,533 )534 535 self._guidance_scale = guidance_scale536 self._guidance_rescale = guidance_rescale537 self._clip_skip = clip_skip538 self._cross_attention_kwargs = cross_attention_kwargs539 self._denoising_end = denoising_end540 self._interrupt = False541 542 # 2. Define call parameters543 if prompt is not None and isinstance(prompt, str):544 batch_size = 1545 elif prompt is not None and isinstance(prompt, list):546 batch_size = len(prompt)547 else:548 batch_size = prompt_embeds.shape[0]549 550 device = self._execution_device551 552 # 3. Encode input prompt553 lora_scale = (554 self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None555 )556 557 (558 prompt_embeds,559 negative_prompt_embeds,560 pooled_prompt_embeds,561 negative_pooled_prompt_embeds,562 ) = self.encode_prompt(563 prompt=prompt,564 prompt_2=prompt_2,565 device=device,566 num_images_per_prompt=num_images_per_prompt,567 do_classifier_free_guidance=self.do_classifier_free_guidance,568 negative_prompt=negative_prompt,569 negative_prompt_2=negative_prompt_2,570 prompt_embeds=prompt_embeds,571 negative_prompt_embeds=negative_prompt_embeds,572 pooled_prompt_embeds=pooled_prompt_embeds,573 negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,574 lora_scale=lora_scale,575 clip_skip=self.clip_skip,576 )577 578 # 4. Preprocess reference image579 ref_image = self.prepare_ref_image(580 image=ref_image,581 width=width,582 height=height,583 batch_size=batch_size * num_images_per_prompt,584 num_images_per_prompt=num_images_per_prompt,585 device=device,586 dtype=prompt_embeds.dtype,587 )588 589 # 5. Prepare timesteps590 timesteps, num_inference_steps = retrieve_timesteps(591 self.scheduler, num_inference_steps, device, timesteps, sigmas592 )593 594 # 6. Prepare latent variables595 num_channels_latents = self.unet.config.in_channels596 latents = self.prepare_latents(597 batch_size * num_images_per_prompt,598 num_channels_latents,599 height,600 width,601 prompt_embeds.dtype,602 device,603 generator,604 latents,605 )606 607 # 7. Prepare reference latent variables608 ref_image_latents = self.prepare_ref_latents(609 ref_image,610 batch_size * num_images_per_prompt,611 prompt_embeds.dtype,612 device,613 generator,614 self.do_classifier_free_guidance,615 )616 617 # 8. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline618 extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)619 620 # 8.1 Create tensor stating which reference controlnets to keep621 reference_keeps = []622 for i in range(len(timesteps)):623 reference_keep = 1.0 - float(624 i / len(timesteps) < reference_guidance_start or (i + 1) / len(timesteps) > reference_guidance_end625 )626 reference_keeps.append(reference_keep)627 628 # 8.2 Modify self attention and group norm629 MODE = "write"630 uc_mask = (631 torch.Tensor([1] * batch_size * num_images_per_prompt + [0] * batch_size * num_images_per_prompt)632 .type_as(ref_image_latents)633 .bool()634 )635 636 do_classifier_free_guidance = self.do_classifier_free_guidance637 638 def hacked_basic_transformer_inner_forward(639 self,640 hidden_states: torch.Tensor,641 attention_mask: Optional[torch.Tensor] = None,642 encoder_hidden_states: Optional[torch.Tensor] = None,643 encoder_attention_mask: Optional[torch.Tensor] = None,644 timestep: Optional[torch.LongTensor] = None,645 cross_attention_kwargs: Dict[str, Any] = None,646 class_labels: Optional[torch.LongTensor] = None,647 ):648 if self.use_ada_layer_norm:649 norm_hidden_states = self.norm1(hidden_states, timestep)650 elif self.use_ada_layer_norm_zero:651 norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(652 hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype653 )654 else:655 norm_hidden_states = self.norm1(hidden_states)656 657 # 1. Self-Attention658 cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}659 if self.only_cross_attention:660 attn_output = self.attn1(661 norm_hidden_states,662 encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,663 attention_mask=attention_mask,664 **cross_attention_kwargs,665 )666 else:667 if MODE == "write":668 self.bank.append(norm_hidden_states.detach().clone())669 attn_output = self.attn1(670 norm_hidden_states,671 encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,672 attention_mask=attention_mask,673 **cross_attention_kwargs,674 )675 if MODE == "read":676 if attention_auto_machine_weight > self.attn_weight:677 attn_output_uc = self.attn1(678 norm_hidden_states,679 encoder_hidden_states=torch.cat([norm_hidden_states] + self.bank, dim=1),680 # attention_mask=attention_mask,681 **cross_attention_kwargs,682 )683 attn_output_c = attn_output_uc.clone()684 if do_classifier_free_guidance and style_fidelity > 0:685 attn_output_c[uc_mask] = self.attn1(686 norm_hidden_states[uc_mask],687 encoder_hidden_states=norm_hidden_states[uc_mask],688 **cross_attention_kwargs,689 )690 attn_output = style_fidelity * attn_output_c + (1.0 - style_fidelity) * attn_output_uc691 self.bank.clear()692 else:693 attn_output = self.attn1(694 norm_hidden_states,695 encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,696 attention_mask=attention_mask,697 **cross_attention_kwargs,698 )699 if self.use_ada_layer_norm_zero:700 attn_output = gate_msa.unsqueeze(1) * attn_output701 hidden_states = attn_output + hidden_states702 703 if self.attn2 is not None:704 norm_hidden_states = (705 self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)706 )707 708 # 2. Cross-Attention709 attn_output = self.attn2(710 norm_hidden_states,711 encoder_hidden_states=encoder_hidden_states,712 attention_mask=encoder_attention_mask,713 **cross_attention_kwargs,714 )715 hidden_states = attn_output + hidden_states716 717 # 3. Feed-forward718 norm_hidden_states = self.norm3(hidden_states)719 720 if self.use_ada_layer_norm_zero:721 norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]722 723 ff_output = self.ff(norm_hidden_states)724 725 if self.use_ada_layer_norm_zero:726 ff_output = gate_mlp.unsqueeze(1) * ff_output727 728 hidden_states = ff_output + hidden_states729 730 return hidden_states731 732 def hacked_mid_forward(self, *args, **kwargs):733 eps = 1e-6734 x = self.original_forward(*args, **kwargs)735 if MODE == "write":736 if gn_auto_machine_weight >= self.gn_weight:737 var, mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0)738 self.mean_bank.append(mean)739 self.var_bank.append(var)740 if MODE == "read":741 if len(self.mean_bank) > 0 and len(self.var_bank) > 0:742 var, mean = torch.var_mean(x, dim=(2, 3), keepdim=True, correction=0)743 std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5744 mean_acc = sum(self.mean_bank) / float(len(self.mean_bank))745 var_acc = sum(self.var_bank) / float(len(self.var_bank))746 std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5747 x_uc = (((x - mean) / std) * std_acc) + mean_acc748 x_c = x_uc.clone()749 if do_classifier_free_guidance and style_fidelity > 0:750 x_c[uc_mask] = x[uc_mask]751 x = style_fidelity * x_c + (1.0 - style_fidelity) * x_uc752 self.mean_bank = []753 self.var_bank = []754 return x755 756 def hack_CrossAttnDownBlock2D_forward(757 self,758 hidden_states: torch.Tensor,759 temb: Optional[torch.Tensor] = None,760 encoder_hidden_states: Optional[torch.Tensor] = None,761 attention_mask: Optional[torch.Tensor] = None,762 cross_attention_kwargs: Optional[Dict[str, Any]] = None,763 encoder_attention_mask: Optional[torch.Tensor] = None,764 ):765 eps = 1e-6766 767 # TODO(Patrick, William) - attention mask is not used768 output_states = ()769 770 for i, (resnet, attn) in enumerate(zip(self.resnets, self.attentions)):771 hidden_states = resnet(hidden_states, temb)772 hidden_states = attn(773 hidden_states,774 encoder_hidden_states=encoder_hidden_states,775 cross_attention_kwargs=cross_attention_kwargs,776 attention_mask=attention_mask,777 encoder_attention_mask=encoder_attention_mask,778 return_dict=False,779 )[0]780 if MODE == "write":781 if gn_auto_machine_weight >= self.gn_weight:782 var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)783 self.mean_bank.append([mean])784 self.var_bank.append([var])785 if MODE == "read":786 if len(self.mean_bank) > 0 and len(self.var_bank) > 0:787 var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)788 std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5789 mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))790 var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))791 std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5792 hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc793 hidden_states_c = hidden_states_uc.clone()794 if do_classifier_free_guidance and style_fidelity > 0:795 hidden_states_c[uc_mask] = hidden_states[uc_mask]796 hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc797 798 output_states = output_states + (hidden_states,)799 800 if MODE == "read":801 self.mean_bank = []802 self.var_bank = []803 804 if self.downsamplers is not None:805 for downsampler in self.downsamplers:806 hidden_states = downsampler(hidden_states)807 808 output_states = output_states + (hidden_states,)809 810 return hidden_states, output_states811 812 def hacked_DownBlock2D_forward(self, hidden_states, temb=None, *args, **kwargs):813 eps = 1e-6814 815 output_states = ()816 817 for i, resnet in enumerate(self.resnets):818 hidden_states = resnet(hidden_states, temb)819 820 if MODE == "write":821 if gn_auto_machine_weight >= self.gn_weight:822 var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)823 self.mean_bank.append([mean])824 self.var_bank.append([var])825 if MODE == "read":826 if len(self.mean_bank) > 0 and len(self.var_bank) > 0:827 var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)828 std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5829 mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))830 var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))831 std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5832 hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc833 hidden_states_c = hidden_states_uc.clone()834 if do_classifier_free_guidance and style_fidelity > 0:835 hidden_states_c[uc_mask] = hidden_states[uc_mask]836 hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc837 838 output_states = output_states + (hidden_states,)839 840 if MODE == "read":841 self.mean_bank = []842 self.var_bank = []843 844 if self.downsamplers is not None:845 for downsampler in self.downsamplers:846 hidden_states = downsampler(hidden_states)847 848 output_states = output_states + (hidden_states,)849 850 return hidden_states, output_states851 852 def hacked_CrossAttnUpBlock2D_forward(853 self,854 hidden_states: torch.Tensor,855 res_hidden_states_tuple: Tuple[torch.Tensor, ...],856 temb: Optional[torch.Tensor] = None,857 encoder_hidden_states: Optional[torch.Tensor] = None,858 cross_attention_kwargs: Optional[Dict[str, Any]] = None,859 upsample_size: Optional[int] = None,860 attention_mask: Optional[torch.Tensor] = None,861 encoder_attention_mask: Optional[torch.Tensor] = None,862 ):863 eps = 1e-6864 # TODO(Patrick, William) - attention mask is not used865 for i, (resnet, attn) in enumerate(zip(self.resnets, self.attentions)):866 # pop res hidden states867 res_hidden_states = res_hidden_states_tuple[-1]868 res_hidden_states_tuple = res_hidden_states_tuple[:-1]869 hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)870 hidden_states = resnet(hidden_states, temb)871 hidden_states = attn(872 hidden_states,873 encoder_hidden_states=encoder_hidden_states,874 cross_attention_kwargs=cross_attention_kwargs,875 attention_mask=attention_mask,876 encoder_attention_mask=encoder_attention_mask,877 return_dict=False,878 )[0]879 880 if MODE == "write":881 if gn_auto_machine_weight >= self.gn_weight:882 var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)883 self.mean_bank.append([mean])884 self.var_bank.append([var])885 if MODE == "read":886 if len(self.mean_bank) > 0 and len(self.var_bank) > 0:887 var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)888 std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5889 mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))890 var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))891 std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5892 hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc893 hidden_states_c = hidden_states_uc.clone()894 if do_classifier_free_guidance and style_fidelity > 0:895 hidden_states_c[uc_mask] = hidden_states[uc_mask]896 hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc897 898 if MODE == "read":899 self.mean_bank = []900 self.var_bank = []901 902 if self.upsamplers is not None:903 for upsampler in self.upsamplers:904 hidden_states = upsampler(hidden_states, upsample_size)905 906 return hidden_states907 908 def hacked_UpBlock2D_forward(909 self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None, *args, **kwargs910 ):911 eps = 1e-6912 for i, resnet in enumerate(self.resnets):913 # pop res hidden states914 res_hidden_states = res_hidden_states_tuple[-1]915 res_hidden_states_tuple = res_hidden_states_tuple[:-1]916 hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)917 hidden_states = resnet(hidden_states, temb)918 919 if MODE == "write":920 if gn_auto_machine_weight >= self.gn_weight:921 var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)922 self.mean_bank.append([mean])923 self.var_bank.append([var])924 if MODE == "read":925 if len(self.mean_bank) > 0 and len(self.var_bank) > 0:926 var, mean = torch.var_mean(hidden_states, dim=(2, 3), keepdim=True, correction=0)927 std = torch.maximum(var, torch.zeros_like(var) + eps) ** 0.5928 mean_acc = sum(self.mean_bank[i]) / float(len(self.mean_bank[i]))929 var_acc = sum(self.var_bank[i]) / float(len(self.var_bank[i]))930 std_acc = torch.maximum(var_acc, torch.zeros_like(var_acc) + eps) ** 0.5931 hidden_states_uc = (((hidden_states - mean) / std) * std_acc) + mean_acc932 hidden_states_c = hidden_states_uc.clone()933 if do_classifier_free_guidance and style_fidelity > 0:934 hidden_states_c[uc_mask] = hidden_states[uc_mask]935 hidden_states = style_fidelity * hidden_states_c + (1.0 - style_fidelity) * hidden_states_uc936 937 if MODE == "read":938 self.mean_bank = []939 self.var_bank = []940 941 if self.upsamplers is not None:942 for upsampler in self.upsamplers:943 hidden_states = upsampler(hidden_states, upsample_size)944 945 return hidden_states946 947 if reference_attn:948 attn_modules = [module for module in torch_dfs(self.unet) if isinstance(module, BasicTransformerBlock)]949 attn_modules = sorted(attn_modules, key=lambda x: -x.norm1.normalized_shape[0])950 951 for i, module in enumerate(attn_modules):952 module._original_inner_forward = module.forward953 module.forward = hacked_basic_transformer_inner_forward.__get__(module, BasicTransformerBlock)954 module.bank = []955 module.attn_weight = float(i) / float(len(attn_modules))956 957 if reference_adain:958 gn_modules = [self.unet.mid_block]959 self.unet.mid_block.gn_weight = 0960 961 down_blocks = self.unet.down_blocks962 for w, module in enumerate(down_blocks):963 module.gn_weight = 1.0 - float(w) / float(len(down_blocks))964 gn_modules.append(module)965 966 up_blocks = self.unet.up_blocks967 for w, module in enumerate(up_blocks):968 module.gn_weight = float(w) / float(len(up_blocks))969 gn_modules.append(module)970 971 for i, module in enumerate(gn_modules):972 if getattr(module, "original_forward", None) is None:973 module.original_forward = module.forward974 if i == 0:975 # mid_block976 module.forward = hacked_mid_forward.__get__(module, torch.nn.Module)977 elif isinstance(module, CrossAttnDownBlock2D):978 module.forward = hack_CrossAttnDownBlock2D_forward.__get__(module, CrossAttnDownBlock2D)979 elif isinstance(module, DownBlock2D):980 module.forward = hacked_DownBlock2D_forward.__get__(module, DownBlock2D)981 elif isinstance(module, CrossAttnUpBlock2D):982 module.forward = hacked_CrossAttnUpBlock2D_forward.__get__(module, CrossAttnUpBlock2D)983 elif isinstance(module, UpBlock2D):984 module.forward = hacked_UpBlock2D_forward.__get__(module, UpBlock2D)985 module.mean_bank = []986 module.var_bank = []987 module.gn_weight *= 2988 989 # 9. Prepare added time ids & embeddings990 add_text_embeds = pooled_prompt_embeds991 if self.text_encoder_2 is None:992 text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1])993 else:994 text_encoder_projection_dim = self.text_encoder_2.config.projection_dim995 996 add_time_ids = self._get_add_time_ids(997 original_size,998 crops_coords_top_left,999 target_size,1000 dtype=prompt_embeds.dtype,1001 text_encoder_projection_dim=text_encoder_projection_dim,1002 )1003 if negative_original_size is not None and negative_target_size is not None:1004 negative_add_time_ids = self._get_add_time_ids(1005 negative_original_size,1006 negative_crops_coords_top_left,1007 negative_target_size,1008 dtype=prompt_embeds.dtype,1009 text_encoder_projection_dim=text_encoder_projection_dim,1010 )1011 else:1012 negative_add_time_ids = add_time_ids1013 1014 if self.do_classifier_free_guidance:1015 prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)1016 add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)1017 add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)1018 1019 prompt_embeds = prompt_embeds.to(device)1020 add_text_embeds = add_text_embeds.to(device)1021 add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1)1022 1023 if ip_adapter_image is not None or ip_adapter_image_embeds is not None:1024 image_embeds = self.prepare_ip_adapter_image_embeds(1025 ip_adapter_image,1026 ip_adapter_image_embeds,1027 device,1028 batch_size * num_images_per_prompt,1029 self.do_classifier_free_guidance,1030 )1031 1032 # 10. Denoising loop1033 num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)1034 1035 # 10.1 Apply denoising_end1036 if (1037 self.denoising_end is not None1038 and isinstance(self.denoising_end, float)1039 and self.denoising_end > 01040 and self.denoising_end < 11041 ):1042 discrete_timestep_cutoff = int(1043 round(1044 self.scheduler.config.num_train_timesteps1045 - (self.denoising_end * self.scheduler.config.num_train_timesteps)1046 )1047 )1048 num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps)))1049 timesteps = timesteps[:num_inference_steps]1050 1051 # 11. Optionally get Guidance Scale Embedding1052 timestep_cond = None1053 if self.unet.config.time_cond_proj_dim is not None:1054 guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(batch_size * num_images_per_prompt)1055 timestep_cond = self.get_guidance_scale_embedding(1056 guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim1057 ).to(device=device, dtype=latents.dtype)1058 1059 self._num_timesteps = len(timesteps)1060 with self.progress_bar(total=num_inference_steps) as progress_bar:1061 for i, t in enumerate(timesteps):1062 if self.interrupt:1063 continue1064 1065 # expand the latents if we are doing classifier free guidance1066 latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents1067 1068 latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)1069 1070 # predict the noise residual1071 added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}1072 if ip_adapter_image is not None or ip_adapter_image_embeds is not None:1073 added_cond_kwargs["image_embeds"] = image_embeds1074 1075 # ref only part1076 if reference_keeps[i] > 0:1077 noise = randn_tensor(1078 ref_image_latents.shape, generator=generator, device=device, dtype=ref_image_latents.dtype1079 )1080 ref_xt = self.scheduler.add_noise(1081 ref_image_latents,1082 noise,1083 t.reshape(1084 1,1085 ),1086 )1087 ref_xt = self.scheduler.scale_model_input(ref_xt, t)1088 1089 MODE = "write"1090 self.unet(1091 ref_xt,1092 t,1093 encoder_hidden_states=prompt_embeds,1094 cross_attention_kwargs=cross_attention_kwargs,1095 added_cond_kwargs=added_cond_kwargs,1096 return_dict=False,1097 )1098 1099 # predict the noise residual1100 MODE = "read"1101 noise_pred = self.unet(1102 latent_model_input,1103 t,1104 encoder_hidden_states=prompt_embeds,1105 timestep_cond=timestep_cond,1106 cross_attention_kwargs=self.cross_attention_kwargs,1107 added_cond_kwargs=added_cond_kwargs,1108 return_dict=False,1109 )[0]1110 1111 # perform guidance1112 if self.do_classifier_free_guidance:1113 noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)1114 noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)1115 1116 if self.do_classifier_free_guidance and self.guidance_rescale > 0.0:1117 # Based on 3.4. in https://huggingface.co/papers/2305.088911118 noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale)1119 1120 # compute the previous noisy sample x_t -> x_t-11121 latents_dtype = latents.dtype1122 latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]1123 if latents.dtype != latents_dtype:1124 if torch.backends.mps.is_available():1125 # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/992721126 latents = latents.to(latents_dtype)1127 1128 if callback_on_step_end is not None:1129 callback_kwargs = {}1130 for k in callback_on_step_end_tensor_inputs:1131 callback_kwargs[k] = locals()[k]1132 callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)1133 1134 latents = callback_outputs.pop("latents", latents)1135 prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)1136 negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)1137 add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds)1138 negative_pooled_prompt_embeds = callback_outputs.pop(1139 "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds1140 )1141 add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids)1142 negative_add_time_ids = callback_outputs.pop("negative_add_time_ids", negative_add_time_ids)1143 1144 # call the callback, if provided1145 if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):1146 progress_bar.update()1147 if callback is not None and i % callback_steps == 0:1148 step_idx = i // getattr(self.scheduler, "order", 1)1149 callback(step_idx, t, latents)1150 1151 if XLA_AVAILABLE:1152 xm.mark_step()1153 1154 if not output_type == "latent":1155 # make sure the VAE is in float32 mode, as it overflows in float161156 needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast1157 1158 if needs_upcasting:1159 self.upcast_vae()1160 latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)1161 elif latents.dtype != self.vae.dtype:1162 if torch.backends.mps.is_available():1163 # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/992721164 self.vae = self.vae.to(latents.dtype)1165 1166 # unscale/denormalize the latents1167 # denormalize with the mean and std if available and not None1168 has_latents_mean = hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None1169 has_latents_std = hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None1170 if has_latents_mean and has_latents_std:1171 latents_mean = (1172 torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype)1173 )1174 latents_std = (1175 torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype)1176 )1177 latents = latents * latents_std / self.vae.config.scaling_factor + latents_mean1178 else:1179 latents = latents / self.vae.config.scaling_factor1180 1181 image = self.vae.decode(latents, return_dict=False)[0]1182 1183 # cast back to fp16 if needed1184 if needs_upcasting:1185 self.vae.to(dtype=torch.float16)1186 else:1187 image = latents1188 1189 if not output_type == "latent":1190 # apply watermark if available1191 if self.watermark is not None:1192 image = self.watermark.apply_watermark(image)1193 1194 image = self.image_processor.postprocess(image, output_type=output_type)1195 1196 # Offload all models1197 self.maybe_free_model_hooks()1198 1199 if not return_dict:1200 return (image,)