SunderAli17 commited on
Commit
5d58002
·
verified ·
1 Parent(s): d9c1219

Create pipelines/pipeline_controlnet_xl_SAK_img2img.py

Browse files
SAK/pipelines/pipeline_controlnet_xl_SAK_img2img.py ADDED
@@ -0,0 +1,1358 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2024 The HuggingFace Team. All rights reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+
15
+
16
+ import inspect
17
+ from typing import Any, Callable, Dict, List, Optional, Tuple, Union
18
+
19
+ import numpy as np
20
+ import PIL.Image
21
+ import torch
22
+ import torch.nn.functional as F
23
+ from transformers import (
24
+ CLIPImageProcessor,
25
+ CLIPTextModel,
26
+ CLIPTextModelWithProjection,
27
+ CLIPTokenizer,
28
+ CLIPVisionModelWithProjection,
29
+ )
30
+
31
+ from diffusers.utils.import_utils import is_invisible_watermark_available
32
+
33
+ from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
34
+ from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
35
+ from diffusers.loaders import (
36
+ FromSingleFileMixin,
37
+ IPAdapterMixin,
38
+ StableDiffusionXLLoraLoaderMixin,
39
+ TextualInversionLoaderMixin,
40
+ )
41
+ from diffusers.models import AutoencoderKL, ImageProjection, UNet2DConditionModel
42
+ from diffusers.models.attention_processor import (
43
+ AttnProcessor2_0,
44
+ XFormersAttnProcessor,
45
+ )
46
+ from diffusers.models.lora import adjust_lora_scale_text_encoder
47
+ from diffusers.schedulers import KarrasDiffusionSchedulers
48
+ from diffusers.utils import (
49
+ USE_PEFT_BACKEND,
50
+ deprecate,
51
+ logging,
52
+ replace_example_docstring,
53
+ scale_lora_layers,
54
+ unscale_lora_layers,
55
+ )
56
+ from diffusers.utils.torch_utils import is_compiled_module, randn_tensor
57
+ from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusionMixin
58
+ from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput
59
+ from diffusers.pipelines.controlnet import MultiControlNetModel
60
+
61
+ from ..models.controlnet import ControlNetModel
62
+
63
+
64
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
65
+
66
+
67
+
68
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
69
+ def retrieve_latents(
70
+ encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
71
+ ):
72
+ if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
73
+ return encoder_output.latent_dist.sample(generator)
74
+ elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
75
+ return encoder_output.latent_dist.mode()
76
+ elif hasattr(encoder_output, "latents"):
77
+ return encoder_output.latents
78
+ else:
79
+ raise AttributeError("Could not access latents of provided encoder_output")
80
+
81
+
82
+ class StableDiffusionXLControlNetImg2ImgPipeline(
83
+ DiffusionPipeline,
84
+ StableDiffusionMixin,
85
+ TextualInversionLoaderMixin,
86
+ StableDiffusionXLLoraLoaderMixin,
87
+ FromSingleFileMixin,
88
+ IPAdapterMixin,
89
+ ):
90
+ r"""
91
+ Pipeline for image-to-image generation using Stable Diffusion XL with ControlNet guidance.
92
+ This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
93
+ library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
94
+ The pipeline also inherits the following loading methods:
95
+ - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings
96
+ - [`~loaders.StableDiffusionXLLoraLoaderMixin.load_lora_weights`] for loading LoRA weights
97
+ - [`~loaders.StableDiffusionXLLoraLoaderMixin.save_lora_weights`] for saving LoRA weights
98
+ - [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters
99
+ Args:
100
+ vae ([`AutoencoderKL`]):
101
+ Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
102
+ text_encoder ([`CLIPTextModel`]):
103
+ Frozen text-encoder. Stable Diffusion uses the text portion of
104
+ [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
105
+ the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
106
+ tokenizer (`CLIPTokenizer`):
107
+ Tokenizer of class
108
+ [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
109
+ unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.
110
+ controlnet ([`ControlNetModel`] or `List[ControlNetModel]`):
111
+ Provides additional conditioning to the unet during the denoising process. If you set multiple ControlNets
112
+ as a list, the outputs from each ControlNet are added together to create one combined additional
113
+ conditioning.
114
+ scheduler ([`SchedulerMixin`]):
115
+ A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
116
+ [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
117
+ requires_aesthetics_score (`bool`, *optional*, defaults to `"False"`):
118
+ Whether the `unet` requires an `aesthetic_score` condition to be passed during inference. Also see the
119
+ config of `stabilityai/stable-diffusion-xl-refiner-1-0`.
120
+ force_zeros_for_empty_prompt (`bool`, *optional*, defaults to `"True"`):
121
+ Whether the negative prompt embeddings shall be forced to always be set to 0. Also see the config of
122
+ `stabilityai/stable-diffusion-xl-base-1-0`.
123
+ add_watermarker (`bool`, *optional*):
124
+ Whether to use the [invisible_watermark library](https://github.com/ShieldMnt/invisible-watermark/) to
125
+ watermark output images. If not defined, it will default to True if the package is installed, otherwise no
126
+ watermarker will be used.
127
+ feature_extractor ([`~transformers.CLIPImageProcessor`]):
128
+ A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`.
129
+ """
130
+
131
+ model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae"
132
+ _optional_components = [
133
+ "tokenizer",
134
+ "text_encoder",
135
+ "feature_extractor",
136
+ "image_encoder",
137
+ ]
138
+ _callback_tensor_inputs = [
139
+ "latents",
140
+ "prompt_embeds",
141
+ "negative_prompt_embeds",
142
+ "add_text_embeds",
143
+ "add_time_ids",
144
+ "negative_pooled_prompt_embeds",
145
+ "add_neg_time_ids",
146
+ ]
147
+
148
+ def __init__(
149
+ self,
150
+ vae: AutoencoderKL,
151
+ text_encoder: CLIPTextModel,
152
+ tokenizer: CLIPTokenizer,
153
+ unet: UNet2DConditionModel,
154
+ controlnet: Union[ControlNetModel, List[ControlNetModel], Tuple[ControlNetModel], MultiControlNetModel],
155
+ scheduler: KarrasDiffusionSchedulers,
156
+ requires_aesthetics_score: bool = False,
157
+ force_zeros_for_empty_prompt: bool = True,
158
+ feature_extractor: CLIPImageProcessor = None,
159
+ image_encoder: CLIPVisionModelWithProjection = None,
160
+ ):
161
+ super().__init__()
162
+
163
+ if isinstance(controlnet, (list, tuple)):
164
+ controlnet = MultiControlNetModel(controlnet)
165
+
166
+ self.register_modules(
167
+ vae=vae,
168
+ text_encoder=text_encoder,
169
+ tokenizer=tokenizer,
170
+ unet=unet,
171
+ controlnet=controlnet,
172
+ scheduler=scheduler,
173
+ feature_extractor=feature_extractor,
174
+ image_encoder=image_encoder,
175
+ )
176
+ self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
177
+ self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True)
178
+ self.control_image_processor = VaeImageProcessor(
179
+ vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True, do_normalize=False
180
+ )
181
+
182
+ self.watermark = None
183
+
184
+ self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt)
185
+ self.register_to_config(requires_aesthetics_score=requires_aesthetics_score)
186
+
187
+
188
+ def encode_prompt(
189
+ self,
190
+ prompt,
191
+ device: Optional[torch.device] = None,
192
+ num_images_per_prompt: int = 1,
193
+ do_classifier_free_guidance: bool = True,
194
+ negative_prompt=None,
195
+ prompt_embeds: Optional[torch.FloatTensor] = None,
196
+ negative_prompt_embeds: Optional[torch.FloatTensor] = None,
197
+ pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
198
+ negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
199
+ lora_scale: Optional[float] = None,
200
+ ):
201
+ r"""
202
+ Encodes the prompt into text encoder hidden states.
203
+ Args:
204
+ prompt (`str` or `List[str]`, *optional*):
205
+ prompt to be encoded
206
+ device: (`torch.device`):
207
+ torch device
208
+ num_images_per_prompt (`int`):
209
+ number of images that should be generated per prompt
210
+ do_classifier_free_guidance (`bool`):
211
+ whether to use classifier free guidance or not
212
+ negative_prompt (`str` or `List[str]`, *optional*):
213
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass
214
+ `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
215
+ less than `1`).
216
+ prompt_embeds (`torch.FloatTensor`, *optional*):
217
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
218
+ provided, text embeddings will be generated from `prompt` input argument.
219
+ negative_prompt_embeds (`torch.FloatTensor`, *optional*):
220
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
221
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
222
+ argument.
223
+ pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
224
+ Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
225
+ If not provided, pooled text embeddings will be generated from `prompt` input argument.
226
+ negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
227
+ Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
228
+ weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
229
+ input argument.
230
+ lora_scale (`float`, *optional*):
231
+ A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
232
+ """
233
+ # from IPython import embed; embed(); exit()
234
+ device = device or self._execution_device
235
+
236
+ # set lora scale so that monkey patched LoRA
237
+ # function of text encoder can correctly access it
238
+ if lora_scale is not None and isinstance(self, LoraLoaderMixin):
239
+ self._lora_scale = lora_scale
240
+
241
+ if prompt is not None and isinstance(prompt, str):
242
+ batch_size = 1
243
+ elif prompt is not None and isinstance(prompt, list):
244
+ batch_size = len(prompt)
245
+ else:
246
+ batch_size = prompt_embeds.shape[0]
247
+
248
+ # Define tokenizers and text encoders
249
+ tokenizers = [self.tokenizer]
250
+ text_encoders = [self.text_encoder]
251
+
252
+ if prompt_embeds is None:
253
+ # textual inversion: procecss multi-vector tokens if necessary
254
+ prompt_embeds_list = []
255
+ for tokenizer, text_encoder in zip(tokenizers, text_encoders):
256
+ if isinstance(self, TextualInversionLoaderMixin):
257
+ prompt = self.maybe_convert_prompt(prompt, tokenizer)
258
+
259
+ text_inputs = tokenizer(
260
+ prompt,
261
+ padding="max_length",
262
+ max_length=256,
263
+ truncation=True,
264
+ return_tensors="pt",
265
+ ).to('cuda')
266
+ output = text_encoder(
267
+ input_ids=text_inputs['input_ids'] ,
268
+ attention_mask=text_inputs['attention_mask'],
269
+ position_ids=text_inputs['position_ids'],
270
+ output_hidden_states=True)
271
+ prompt_embeds = output.hidden_states[-2].permute(1, 0, 2).clone()
272
+ pooled_prompt_embeds = output.hidden_states[-1][-1, :, :].clone() # [batch_size, 4096]
273
+ bs_embed, seq_len, _ = prompt_embeds.shape
274
+ prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
275
+ prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
276
+
277
+ prompt_embeds_list.append(prompt_embeds)
278
+
279
+ # prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
280
+ prompt_embeds = prompt_embeds_list[0]
281
+
282
+ # get unconditional embeddings for classifier free guidance
283
+ zero_out_negative_prompt = negative_prompt is None and self.config.force_zeros_for_empty_prompt
284
+ if do_classifier_free_guidance and negative_prompt_embeds is None and zero_out_negative_prompt:
285
+ negative_prompt_embeds = torch.zeros_like(prompt_embeds)
286
+ negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds)
287
+ elif do_classifier_free_guidance and negative_prompt_embeds is None:
288
+ # negative_prompt = negative_prompt or ""
289
+ uncond_tokens: List[str]
290
+ if negative_prompt is None:
291
+ uncond_tokens = [""] * batch_size
292
+ elif prompt is not None and type(prompt) is not type(negative_prompt):
293
+ raise TypeError(
294
+ f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
295
+ f" {type(prompt)}."
296
+ )
297
+ elif isinstance(negative_prompt, str):
298
+ uncond_tokens = [negative_prompt]
299
+ elif batch_size != len(negative_prompt):
300
+ raise ValueError(
301
+ f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
302
+ f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
303
+ " the batch size of `prompt`."
304
+ )
305
+ else:
306
+ uncond_tokens = negative_prompt
307
+
308
+ negative_prompt_embeds_list = []
309
+ for tokenizer, text_encoder in zip(tokenizers, text_encoders):
310
+ # textual inversion: procecss multi-vector tokens if necessary
311
+ if isinstance(self, TextualInversionLoaderMixin):
312
+ uncond_tokens = self.maybe_convert_prompt(uncond_tokens, tokenizer)
313
+
314
+ max_length = prompt_embeds.shape[1]
315
+ uncond_input = tokenizer(
316
+ uncond_tokens,
317
+ padding="max_length",
318
+ max_length=max_length,
319
+ truncation=True,
320
+ return_tensors="pt",
321
+ ).to('cuda')
322
+ output = text_encoder(
323
+ input_ids=uncond_input['input_ids'] ,
324
+ attention_mask=uncond_input['attention_mask'],
325
+ position_ids=uncond_input['position_ids'],
326
+ output_hidden_states=True)
327
+ negative_prompt_embeds = output.hidden_states[-2].permute(1, 0, 2).clone()
328
+ negative_pooled_prompt_embeds = output.hidden_states[-1][-1, :, :].clone() # [batch_size, 4096]
329
+
330
+ if do_classifier_free_guidance:
331
+ # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
332
+ seq_len = negative_prompt_embeds.shape[1]
333
+
334
+ negative_prompt_embeds = negative_prompt_embeds.to(dtype=text_encoder.dtype, device=device)
335
+
336
+ negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
337
+ negative_prompt_embeds = negative_prompt_embeds.view(
338
+ batch_size * num_images_per_prompt, seq_len, -1
339
+ )
340
+
341
+ # For classifier free guidance, we need to do two forward passes.
342
+ # Here we concatenate the unconditional and text embeddings into a single batch
343
+ # to avoid doing two forward passes
344
+
345
+ negative_prompt_embeds_list.append(negative_prompt_embeds)
346
+
347
+ # negative_prompt_embeds = torch.concat(negative_prompt_embeds_list, dim=-1)
348
+ negative_prompt_embeds = negative_prompt_embeds_list[0]
349
+
350
+ bs_embed = pooled_prompt_embeds.shape[0]
351
+ pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
352
+ bs_embed * num_images_per_prompt, -1
353
+ )
354
+ if do_classifier_free_guidance:
355
+ negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
356
+ bs_embed * num_images_per_prompt, -1
357
+ )
358
+
359
+ return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds
360
+
361
+
362
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image
363
+ def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None):
364
+ dtype = next(self.image_encoder.parameters()).dtype
365
+
366
+ if not isinstance(image, torch.Tensor):
367
+ image = self.feature_extractor(image, return_tensors="pt").pixel_values
368
+
369
+ image = image.to(device=device, dtype=dtype)
370
+ if output_hidden_states:
371
+ image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
372
+ image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
373
+ uncond_image_enc_hidden_states = self.image_encoder(
374
+ torch.zeros_like(image), output_hidden_states=True
375
+ ).hidden_states[-2]
376
+ uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave(
377
+ num_images_per_prompt, dim=0
378
+ )
379
+ return image_enc_hidden_states, uncond_image_enc_hidden_states
380
+ else:
381
+ image_embeds = self.image_encoder(image).image_embeds
382
+ image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
383
+ uncond_image_embeds = torch.zeros_like(image_embeds)
384
+
385
+ return image_embeds, uncond_image_embeds
386
+
387
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds
388
+ def prepare_ip_adapter_image_embeds(
389
+ self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance
390
+ ):
391
+ image_embeds = []
392
+ if do_classifier_free_guidance:
393
+ negative_image_embeds = []
394
+ if ip_adapter_image_embeds is None:
395
+ if not isinstance(ip_adapter_image, list):
396
+ ip_adapter_image = [ip_adapter_image]
397
+
398
+ if len(ip_adapter_image) != len(self.unet.encoder_hid_proj.image_projection_layers):
399
+ raise ValueError(
400
+ f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters."
401
+ )
402
+
403
+ for single_ip_adapter_image, image_proj_layer in zip(
404
+ ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers
405
+ ):
406
+ output_hidden_state = not isinstance(image_proj_layer, ImageProjection)
407
+ single_image_embeds, single_negative_image_embeds = self.encode_image(
408
+ single_ip_adapter_image, device, 1, output_hidden_state
409
+ )
410
+
411
+ image_embeds.append(single_image_embeds[None, :])
412
+ if do_classifier_free_guidance:
413
+ negative_image_embeds.append(single_negative_image_embeds[None, :])
414
+ else:
415
+ for single_image_embeds in ip_adapter_image_embeds:
416
+ if do_classifier_free_guidance:
417
+ single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2)
418
+ negative_image_embeds.append(single_negative_image_embeds)
419
+ image_embeds.append(single_image_embeds)
420
+
421
+ ip_adapter_image_embeds = []
422
+ for i, single_image_embeds in enumerate(image_embeds):
423
+ single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0)
424
+ if do_classifier_free_guidance:
425
+ single_negative_image_embeds = torch.cat([negative_image_embeds[i]] * num_images_per_prompt, dim=0)
426
+ single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds], dim=0)
427
+
428
+ single_image_embeds = single_image_embeds.to(device=device)
429
+ ip_adapter_image_embeds.append(single_image_embeds)
430
+
431
+ return ip_adapter_image_embeds
432
+
433
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
434
+ def prepare_extra_step_kwargs(self, generator, eta):
435
+ # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
436
+ # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
437
+ # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
438
+ # and should be between [0, 1]
439
+
440
+ accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
441
+ extra_step_kwargs = {}
442
+ if accepts_eta:
443
+ extra_step_kwargs["eta"] = eta
444
+
445
+ # check if the scheduler accepts generator
446
+ accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
447
+ if accepts_generator:
448
+ extra_step_kwargs["generator"] = generator
449
+ return extra_step_kwargs
450
+
451
+ def check_inputs(
452
+ self,
453
+ prompt,
454
+ image,
455
+ strength,
456
+ num_inference_steps,
457
+ callback_steps,
458
+ negative_prompt=None,
459
+ prompt_embeds=None,
460
+ negative_prompt_embeds=None,
461
+ pooled_prompt_embeds=None,
462
+ negative_pooled_prompt_embeds=None,
463
+ ip_adapter_image=None,
464
+ ip_adapter_image_embeds=None,
465
+ controlnet_conditioning_scale=1.0,
466
+ control_guidance_start=0.0,
467
+ control_guidance_end=1.0,
468
+ callback_on_step_end_tensor_inputs=None,
469
+ ):
470
+ if strength < 0 or strength > 1:
471
+ raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}")
472
+ if num_inference_steps is None:
473
+ raise ValueError("`num_inference_steps` cannot be None.")
474
+ elif not isinstance(num_inference_steps, int) or num_inference_steps <= 0:
475
+ raise ValueError(
476
+ f"`num_inference_steps` has to be a positive integer but is {num_inference_steps} of type"
477
+ f" {type(num_inference_steps)}."
478
+ )
479
+
480
+ if callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0):
481
+ raise ValueError(
482
+ f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
483
+ f" {type(callback_steps)}."
484
+ )
485
+
486
+ if callback_on_step_end_tensor_inputs is not None and not all(
487
+ k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
488
+ ):
489
+ raise ValueError(
490
+ 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]}"
491
+ )
492
+
493
+ if prompt is not None and prompt_embeds is not None:
494
+ raise ValueError(
495
+ f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
496
+ " only forward one of the two."
497
+ )
498
+ elif prompt is None and prompt_embeds is None:
499
+ raise ValueError(
500
+ "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
501
+ )
502
+ elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
503
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
504
+
505
+ if negative_prompt is not None and negative_prompt_embeds is not None:
506
+ raise ValueError(
507
+ f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
508
+ f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
509
+ )
510
+
511
+ if prompt_embeds is not None and negative_prompt_embeds is not None:
512
+ if prompt_embeds.shape != negative_prompt_embeds.shape:
513
+ raise ValueError(
514
+ "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
515
+ f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
516
+ f" {negative_prompt_embeds.shape}."
517
+ )
518
+
519
+ if prompt_embeds is not None and pooled_prompt_embeds is None:
520
+ raise ValueError(
521
+ "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`."
522
+ )
523
+
524
+ if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is None:
525
+ raise ValueError(
526
+ "If `negative_prompt_embeds` are provided, `negative_pooled_prompt_embeds` also have to be passed. Make sure to generate `negative_pooled_prompt_embeds` from the same text encoder that was used to generate `negative_prompt_embeds`."
527
+ )
528
+
529
+ # `prompt` needs more sophisticated handling when there are multiple
530
+ # conditionings.
531
+ if isinstance(self.controlnet, MultiControlNetModel):
532
+ if isinstance(prompt, list):
533
+ logger.warning(
534
+ f"You have {len(self.controlnet.nets)} ControlNets and you have passed {len(prompt)}"
535
+ " prompts. The conditionings will be fixed across the prompts."
536
+ )
537
+
538
+ # Check `image`
539
+ is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance(
540
+ self.controlnet, torch._dynamo.eval_frame.OptimizedModule
541
+ )
542
+ if (
543
+ isinstance(self.controlnet, ControlNetModel)
544
+ or is_compiled
545
+ and isinstance(self.controlnet._orig_mod, ControlNetModel)
546
+ ):
547
+ self.check_image(image, prompt, prompt_embeds)
548
+ elif (
549
+ isinstance(self.controlnet, MultiControlNetModel)
550
+ or is_compiled
551
+ and isinstance(self.controlnet._orig_mod, MultiControlNetModel)
552
+ ):
553
+ if not isinstance(image, list):
554
+ raise TypeError("For multiple controlnets: `image` must be type `list`")
555
+
556
+ # When `image` is a nested list:
557
+ # (e.g. [[canny_image_1, pose_image_1], [canny_image_2, pose_image_2]])
558
+ elif any(isinstance(i, list) for i in image):
559
+ raise ValueError("A single batch of multiple conditionings are supported at the moment.")
560
+ elif len(image) != len(self.controlnet.nets):
561
+ raise ValueError(
562
+ f"For multiple controlnets: `image` must have the same length as the number of controlnets, but got {len(image)} images and {len(self.controlnet.nets)} ControlNets."
563
+ )
564
+
565
+ for image_ in image:
566
+ self.check_image(image_, prompt, prompt_embeds)
567
+ else:
568
+ assert False
569
+
570
+ # Check `controlnet_conditioning_scale`
571
+ if (
572
+ isinstance(self.controlnet, ControlNetModel)
573
+ or is_compiled
574
+ and isinstance(self.controlnet._orig_mod, ControlNetModel)
575
+ ):
576
+ if not isinstance(controlnet_conditioning_scale, float):
577
+ raise TypeError("For single controlnet: `controlnet_conditioning_scale` must be type `float`.")
578
+ elif (
579
+ isinstance(self.controlnet, MultiControlNetModel)
580
+ or is_compiled
581
+ and isinstance(self.controlnet._orig_mod, MultiControlNetModel)
582
+ ):
583
+ if isinstance(controlnet_conditioning_scale, list):
584
+ if any(isinstance(i, list) for i in controlnet_conditioning_scale):
585
+ raise ValueError("A single batch of multiple conditionings are supported at the moment.")
586
+ elif isinstance(controlnet_conditioning_scale, list) and len(controlnet_conditioning_scale) != len(
587
+ self.controlnet.nets
588
+ ):
589
+ raise ValueError(
590
+ "For multiple controlnets: When `controlnet_conditioning_scale` is specified as `list`, it must have"
591
+ " the same length as the number of controlnets"
592
+ )
593
+ else:
594
+ assert False
595
+
596
+ if not isinstance(control_guidance_start, (tuple, list)):
597
+ control_guidance_start = [control_guidance_start]
598
+
599
+ if not isinstance(control_guidance_end, (tuple, list)):
600
+ control_guidance_end = [control_guidance_end]
601
+
602
+ if len(control_guidance_start) != len(control_guidance_end):
603
+ raise ValueError(
604
+ f"`control_guidance_start` has {len(control_guidance_start)} elements, but `control_guidance_end` has {len(control_guidance_end)} elements. Make sure to provide the same number of elements to each list."
605
+ )
606
+
607
+ if isinstance(self.controlnet, MultiControlNetModel):
608
+ if len(control_guidance_start) != len(self.controlnet.nets):
609
+ raise ValueError(
610
+ f"`control_guidance_start`: {control_guidance_start} has {len(control_guidance_start)} elements but there are {len(self.controlnet.nets)} controlnets available. Make sure to provide {len(self.controlnet.nets)}."
611
+ )
612
+
613
+ for start, end in zip(control_guidance_start, control_guidance_end):
614
+ if start >= end:
615
+ raise ValueError(
616
+ f"control guidance start: {start} cannot be larger or equal to control guidance end: {end}."
617
+ )
618
+ if start < 0.0:
619
+ raise ValueError(f"control guidance start: {start} can't be smaller than 0.")
620
+ if end > 1.0:
621
+ raise ValueError(f"control guidance end: {end} can't be larger than 1.0.")
622
+
623
+ if ip_adapter_image is not None and ip_adapter_image_embeds is not None:
624
+ raise ValueError(
625
+ "Provide either `ip_adapter_image` or `ip_adapter_image_embeds`. Cannot leave both `ip_adapter_image` and `ip_adapter_image_embeds` defined."
626
+ )
627
+
628
+ if ip_adapter_image_embeds is not None:
629
+ if not isinstance(ip_adapter_image_embeds, list):
630
+ raise ValueError(
631
+ f"`ip_adapter_image_embeds` has to be of type `list` but is {type(ip_adapter_image_embeds)}"
632
+ )
633
+ elif ip_adapter_image_embeds[0].ndim not in [3, 4]:
634
+ raise ValueError(
635
+ f"`ip_adapter_image_embeds` has to be a list of 3D or 4D tensors but is {ip_adapter_image_embeds[0].ndim}D"
636
+ )
637
+
638
+ # Copied from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl.StableDiffusionXLControlNetPipeline.check_image
639
+ def check_image(self, image, prompt, prompt_embeds):
640
+ image_is_pil = isinstance(image, PIL.Image.Image)
641
+ image_is_tensor = isinstance(image, torch.Tensor)
642
+ image_is_np = isinstance(image, np.ndarray)
643
+ image_is_pil_list = isinstance(image, list) and isinstance(image[0], PIL.Image.Image)
644
+ image_is_tensor_list = isinstance(image, list) and isinstance(image[0], torch.Tensor)
645
+ image_is_np_list = isinstance(image, list) and isinstance(image[0], np.ndarray)
646
+
647
+ if (
648
+ not image_is_pil
649
+ and not image_is_tensor
650
+ and not image_is_np
651
+ and not image_is_pil_list
652
+ and not image_is_tensor_list
653
+ and not image_is_np_list
654
+ ):
655
+ raise TypeError(
656
+ f"image must be passed and be one of PIL image, numpy array, torch tensor, list of PIL images, list of numpy arrays or list of torch tensors, but is {type(image)}"
657
+ )
658
+
659
+ if image_is_pil:
660
+ image_batch_size = 1
661
+ else:
662
+ image_batch_size = len(image)
663
+
664
+ if prompt is not None and isinstance(prompt, str):
665
+ prompt_batch_size = 1
666
+ elif prompt is not None and isinstance(prompt, list):
667
+ prompt_batch_size = len(prompt)
668
+ elif prompt_embeds is not None:
669
+ prompt_batch_size = prompt_embeds.shape[0]
670
+
671
+ if image_batch_size != 1 and image_batch_size != prompt_batch_size:
672
+ raise ValueError(
673
+ f"If image batch size is not 1, image batch size must be same as prompt batch size. image batch size: {image_batch_size}, prompt batch size: {prompt_batch_size}"
674
+ )
675
+
676
+ # Copied from diffusers.pipelines.controlnet.pipeline_controlnet_sd_xl.StableDiffusionXLControlNetPipeline.prepare_image
677
+ def prepare_control_image(
678
+ self,
679
+ image,
680
+ width,
681
+ height,
682
+ batch_size,
683
+ num_images_per_prompt,
684
+ device,
685
+ dtype,
686
+ do_classifier_free_guidance=False,
687
+ guess_mode=False,
688
+ ):
689
+ image = self.control_image_processor.preprocess(image, height=height, width=width).to(dtype=torch.float32)
690
+ image_batch_size = image.shape[0]
691
+
692
+ if image_batch_size == 1:
693
+ repeat_by = batch_size
694
+ else:
695
+ # image batch size is the same as prompt batch size
696
+ repeat_by = num_images_per_prompt
697
+
698
+ image = image.repeat_interleave(repeat_by, dim=0)
699
+
700
+ image = image.to(device=device, dtype=dtype)
701
+
702
+ if do_classifier_free_guidance and not guess_mode:
703
+ image = torch.cat([image] * 2)
704
+
705
+ return image
706
+
707
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.StableDiffusionImg2ImgPipeline.get_timesteps
708
+ def get_timesteps(self, num_inference_steps, strength, device):
709
+ # get the original timestep using init_timestep
710
+ init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
711
+
712
+ t_start = max(num_inference_steps - init_timestep, 0)
713
+ timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :]
714
+ if hasattr(self.scheduler, "set_begin_index"):
715
+ self.scheduler.set_begin_index(t_start * self.scheduler.order)
716
+
717
+ return timesteps, num_inference_steps - t_start
718
+
719
+ # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_img2img.StableDiffusionXLImg2ImgPipeline.prepare_latents
720
+ def prepare_latents(
721
+ self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None, add_noise=True
722
+ ):
723
+ if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)):
724
+ raise ValueError(
725
+ f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}"
726
+ )
727
+
728
+ # Offload text encoder if `enable_model_cpu_offload` was enabled
729
+ if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
730
+ torch.cuda.empty_cache()
731
+
732
+ image = image.to(device=device, dtype=dtype)
733
+
734
+ batch_size = batch_size * num_images_per_prompt
735
+
736
+ if image.shape[1] == 4:
737
+ init_latents = image
738
+
739
+ else:
740
+ # make sure the VAE is in float32 mode, as it overflows in float16
741
+ if self.vae.config.force_upcast:
742
+ image = image.float()
743
+ self.vae.to(dtype=torch.float32)
744
+
745
+ if isinstance(generator, list) and len(generator) != batch_size:
746
+ raise ValueError(
747
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
748
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
749
+ )
750
+
751
+ elif isinstance(generator, list):
752
+ init_latents = [
753
+ retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i])
754
+ for i in range(batch_size)
755
+ ]
756
+ init_latents = torch.cat(init_latents, dim=0)
757
+ else:
758
+ init_latents = retrieve_latents(self.vae.encode(image), generator=generator)
759
+
760
+ if self.vae.config.force_upcast:
761
+ self.vae.to(dtype)
762
+
763
+ init_latents = init_latents.to(dtype)
764
+
765
+ init_latents = self.vae.config.scaling_factor * init_latents
766
+
767
+ if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0:
768
+ # expand init_latents for batch_size
769
+ additional_image_per_prompt = batch_size // init_latents.shape[0]
770
+ init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0)
771
+ elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0:
772
+ raise ValueError(
773
+ f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts."
774
+ )
775
+ else:
776
+ init_latents = torch.cat([init_latents], dim=0)
777
+
778
+ if add_noise:
779
+ shape = init_latents.shape
780
+ noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
781
+ # get latents
782
+ init_latents = self.scheduler.add_noise(init_latents, noise, timestep)
783
+
784
+ latents = init_latents
785
+
786
+ return latents
787
+
788
+
789
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents
790
+ def prepare_latents_t2i(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
791
+ shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
792
+ if isinstance(generator, list) and len(generator) != batch_size:
793
+ raise ValueError(
794
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
795
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
796
+ )
797
+
798
+ if latents is None:
799
+ latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
800
+ else:
801
+ latents = latents.to(device)
802
+
803
+ # scale the initial noise by the standard deviation required by the scheduler
804
+ latents = latents * self.scheduler.init_noise_sigma
805
+ return latents
806
+
807
+
808
+
809
+ def _get_add_time_ids(self, original_size, crops_coords_top_left, target_size, dtype):
810
+ add_time_ids = list(original_size + crops_coords_top_left + target_size)
811
+
812
+ passed_add_embed_dim = (
813
+ self.unet.config.addition_time_embed_dim * len(add_time_ids) + 4096
814
+ )
815
+ expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features
816
+
817
+ if expected_add_embed_dim != passed_add_embed_dim:
818
+ raise ValueError(
819
+ f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`."
820
+ )
821
+
822
+ add_time_ids = torch.tensor([add_time_ids], dtype=dtype)
823
+ return add_time_ids
824
+
825
+
826
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae
827
+ def upcast_vae(self):
828
+ dtype = self.vae.dtype
829
+ self.vae.to(dtype=torch.float32)
830
+ use_torch_2_0_or_xformers = isinstance(
831
+ self.vae.decoder.mid_block.attentions[0].processor,
832
+ (
833
+ AttnProcessor2_0,
834
+ XFormersAttnProcessor,
835
+ ),
836
+ )
837
+ # if xformers or torch_2_0 is used attention block does not need
838
+ # to be in float32 which can save lots of memory
839
+ if use_torch_2_0_or_xformers:
840
+ self.vae.post_quant_conv.to(dtype)
841
+ self.vae.decoder.conv_in.to(dtype)
842
+ self.vae.decoder.mid_block.to(dtype)
843
+
844
+ @property
845
+ def guidance_scale(self):
846
+ return self._guidance_scale
847
+
848
+ @property
849
+ def clip_skip(self):
850
+ return self._clip_skip
851
+
852
+ # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
853
+ # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
854
+ # corresponds to doing no classifier free guidance.
855
+ @property
856
+ def do_classifier_free_guidance(self):
857
+ return self._guidance_scale > 1
858
+
859
+ @property
860
+ def cross_attention_kwargs(self):
861
+ return self._cross_attention_kwargs
862
+
863
+ @property
864
+ def num_timesteps(self):
865
+ return self._num_timesteps
866
+
867
+ @torch.no_grad()
868
+ def __call__(
869
+ self,
870
+ prompt: Union[str, List[str]] = None,
871
+ image: PipelineImageInput = None,
872
+ control_image: PipelineImageInput = None,
873
+ height: Optional[int] = None,
874
+ width: Optional[int] = None,
875
+ strength: float = 0.8,
876
+ num_inference_steps: int = 50,
877
+ guidance_scale: float = 5.0,
878
+ negative_prompt: Optional[Union[str, List[str]]] = None,
879
+ num_images_per_prompt: Optional[int] = 1,
880
+ eta: float = 0.0,
881
+ guess_mode: bool = False,
882
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
883
+ latents: Optional[torch.Tensor] = None,
884
+ prompt_embeds: Optional[torch.Tensor] = None,
885
+ negative_prompt_embeds: Optional[torch.Tensor] = None,
886
+ pooled_prompt_embeds: Optional[torch.Tensor] = None,
887
+ negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
888
+ ip_adapter_image: Optional[PipelineImageInput] = None,
889
+ ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
890
+ output_type: Optional[str] = "pil",
891
+ return_dict: bool = True,
892
+ cross_attention_kwargs: Optional[Dict[str, Any]] = None,
893
+ controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
894
+ control_guidance_start: Union[float, List[float]] = 0.0,
895
+ control_guidance_end: Union[float, List[float]] = 1.0,
896
+ original_size: Tuple[int, int] = None,
897
+ crops_coords_top_left: Tuple[int, int] = (0, 0),
898
+ target_size: Tuple[int, int] = None,
899
+ clip_skip: Optional[int] = None,
900
+ callback_on_step_end: Optional[
901
+ Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
902
+ ] = None,
903
+ callback_on_step_end_tensor_inputs: List[str] = ["latents"],
904
+ **kwargs,
905
+ ):
906
+ r"""
907
+ Function invoked when calling the pipeline for generation.
908
+ Args:
909
+ prompt (`str` or `List[str]`, *optional*):
910
+ The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
911
+ instead.
912
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
913
+ `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
914
+ The initial image will be used as the starting point for the image generation process. Can also accept
915
+ image latents as `image`, if passing latents directly, it will not be encoded again.
916
+ control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
917
+ `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
918
+ The ControlNet input condition. ControlNet uses this input condition to generate guidance to Unet. If
919
+ the type is specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also
920
+ be accepted as an image. The dimensions of the output image defaults to `image`'s dimensions. If height
921
+ and/or width are passed, `image` is resized according to them. If multiple ControlNets are specified in
922
+ init, images must be passed as a list such that each element of the list can be correctly batched for
923
+ input to a single controlnet.
924
+ height (`int`, *optional*, defaults to the size of control_image):
925
+ The height in pixels of the generated image. Anything below 512 pixels won't work well for
926
+ [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
927
+ and checkpoints that are not specifically fine-tuned on low resolutions.
928
+ width (`int`, *optional*, defaults to the size of control_image):
929
+ The width in pixels of the generated image. Anything below 512 pixels won't work well for
930
+ [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
931
+ and checkpoints that are not specifically fine-tuned on low resolutions.
932
+ strength (`float`, *optional*, defaults to 0.8):
933
+ Indicates extent to transform the reference `image`. Must be between 0 and 1. `image` is used as a
934
+ starting point and more noise is added the higher the `strength`. The number of denoising steps depends
935
+ on the amount of noise initially added. When `strength` is 1, added noise is maximum and the denoising
936
+ process runs for the full number of iterations specified in `num_inference_steps`. A value of 1
937
+ essentially ignores `image`.
938
+ num_inference_steps (`int`, *optional*, defaults to 50):
939
+ The number of denoising steps. More denoising steps usually lead to a higher quality image at the
940
+ expense of slower inference.
941
+ guidance_scale (`float`, *optional*, defaults to 7.5):
942
+ Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
943
+ `guidance_scale` is defined as `w` of equation 2. of [Imagen
944
+ Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
945
+ 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
946
+ usually at the expense of lower image quality.
947
+ negative_prompt (`str` or `List[str]`, *optional*):
948
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass
949
+ `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
950
+ less than `1`).
951
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
952
+ The number of images to generate per prompt.
953
+ eta (`float`, *optional*, defaults to 0.0):
954
+ Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
955
+ [`schedulers.DDIMScheduler`], will be ignored for others.
956
+ generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
957
+ One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
958
+ to make generation deterministic.
959
+ latents (`torch.Tensor`, *optional*):
960
+ Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
961
+ generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
962
+ tensor will ge generated by sampling using the supplied random `generator`.
963
+ prompt_embeds (`torch.Tensor`, *optional*):
964
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
965
+ provided, text embeddings will be generated from `prompt` input argument.
966
+ negative_prompt_embeds (`torch.Tensor`, *optional*):
967
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
968
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
969
+ argument.
970
+ pooled_prompt_embeds (`torch.Tensor`, *optional*):
971
+ Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
972
+ If not provided, pooled text embeddings will be generated from `prompt` input argument.
973
+ negative_pooled_prompt_embeds (`torch.Tensor`, *optional*):
974
+ Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
975
+ weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
976
+ input argument.
977
+ ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
978
+ ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
979
+ Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
980
+ IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
981
+ contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
982
+ provided, embeddings are computed from the `ip_adapter_image` input argument.
983
+ output_type (`str`, *optional*, defaults to `"pil"`):
984
+ The output format of the generate image. Choose between
985
+ [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
986
+ return_dict (`bool`, *optional*, defaults to `True`):
987
+ Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
988
+ plain tuple.
989
+ cross_attention_kwargs (`dict`, *optional*):
990
+ A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
991
+ `self.processor` in
992
+ [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
993
+ controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
994
+ The outputs of the controlnet are multiplied by `controlnet_conditioning_scale` before they are added
995
+ to the residual in the original unet. If multiple ControlNets are specified in init, you can set the
996
+ corresponding scale as a list.
997
+ control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
998
+ The percentage of total steps at which the controlnet starts applying.
999
+ control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
1000
+ The percentage of total steps at which the controlnet stops applying.
1001
+ original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
1002
+ If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
1003
+ `original_size` defaults to `(height, width)` if not specified. Part of SDXL's micro-conditioning as
1004
+ explained in section 2.2 of
1005
+ [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
1006
+ crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
1007
+ `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
1008
+ `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
1009
+ `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
1010
+ [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
1011
+ target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
1012
+ For most cases, `target_size` should be set to the desired height and width of the generated image. If
1013
+ not specified it will default to `(height, width)`. Part of SDXL's micro-conditioning as explained in
1014
+ section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
1015
+ clip_skip (`int`, *optional*):
1016
+ Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
1017
+ the output of the pre-final layer will be used for computing the prompt embeddings.
1018
+ callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
1019
+ A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of
1020
+ each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
1021
+ DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
1022
+ list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
1023
+ callback_on_step_end_tensor_inputs (`List`, *optional*):
1024
+ The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
1025
+ will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
1026
+ `._callback_tensor_inputs` attribute of your pipeline class.
1027
+ Examples:
1028
+ Returns:
1029
+ [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
1030
+ [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple`
1031
+ containing the output images.
1032
+ """
1033
+
1034
+ callback = kwargs.pop("callback", None)
1035
+ callback_steps = kwargs.pop("callback_steps", None)
1036
+
1037
+ if callback is not None:
1038
+ deprecate(
1039
+ "callback",
1040
+ "1.0.0",
1041
+ "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
1042
+ )
1043
+ if callback_steps is not None:
1044
+ deprecate(
1045
+ "callback_steps",
1046
+ "1.0.0",
1047
+ "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
1048
+ )
1049
+
1050
+ if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
1051
+ callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
1052
+
1053
+ controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet
1054
+
1055
+ # align format for control guidance
1056
+ if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list):
1057
+ control_guidance_start = len(control_guidance_end) * [control_guidance_start]
1058
+ elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list):
1059
+ control_guidance_end = len(control_guidance_start) * [control_guidance_end]
1060
+ elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list):
1061
+ mult = len(controlnet.nets) if isinstance(controlnet, MultiControlNetModel) else 1
1062
+ control_guidance_start, control_guidance_end = (
1063
+ mult * [control_guidance_start],
1064
+ mult * [control_guidance_end],
1065
+ )
1066
+
1067
+ # from IPython import embed; embed()
1068
+ # 1. Check inputs. Raise error if not correct
1069
+ self.check_inputs(
1070
+ prompt,
1071
+ control_image,
1072
+ strength,
1073
+ num_inference_steps,
1074
+ callback_steps,
1075
+ negative_prompt,
1076
+ prompt_embeds,
1077
+ negative_prompt_embeds,
1078
+ pooled_prompt_embeds,
1079
+ negative_pooled_prompt_embeds,
1080
+ ip_adapter_image,
1081
+ ip_adapter_image_embeds,
1082
+ controlnet_conditioning_scale,
1083
+ control_guidance_start,
1084
+ control_guidance_end,
1085
+ callback_on_step_end_tensor_inputs,
1086
+ )
1087
+
1088
+ self._guidance_scale = guidance_scale
1089
+ self._clip_skip = clip_skip
1090
+ self._cross_attention_kwargs = cross_attention_kwargs
1091
+
1092
+ # 2. Define call parameters
1093
+ if prompt is not None and isinstance(prompt, str):
1094
+ batch_size = 1
1095
+ elif prompt is not None and isinstance(prompt, list):
1096
+ batch_size = len(prompt)
1097
+ else:
1098
+ batch_size = prompt_embeds.shape[0]
1099
+
1100
+ device = self._execution_device
1101
+
1102
+ if isinstance(controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float):
1103
+ controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(controlnet.nets)
1104
+
1105
+ # 3.1. Encode input prompt
1106
+ text_encoder_lora_scale = (
1107
+ self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None
1108
+ )
1109
+ (
1110
+ prompt_embeds,
1111
+ negative_prompt_embeds,
1112
+ pooled_prompt_embeds,
1113
+ negative_pooled_prompt_embeds,
1114
+ ) = self.encode_prompt(
1115
+ prompt,
1116
+ device,
1117
+ num_images_per_prompt,
1118
+ self.do_classifier_free_guidance,
1119
+ negative_prompt,
1120
+ prompt_embeds=prompt_embeds,
1121
+ negative_prompt_embeds=negative_prompt_embeds,
1122
+ pooled_prompt_embeds=pooled_prompt_embeds,
1123
+ negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
1124
+ lora_scale=text_encoder_lora_scale,
1125
+ )
1126
+
1127
+ # 3.2 Encode ip_adapter_image
1128
+ if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
1129
+ image_embeds = self.prepare_ip_adapter_image_embeds(
1130
+ ip_adapter_image,
1131
+ ip_adapter_image_embeds,
1132
+ device,
1133
+ batch_size * num_images_per_prompt,
1134
+ self.do_classifier_free_guidance,
1135
+ )
1136
+
1137
+ # 4. Prepare image and controlnet_conditioning_image
1138
+ image = self.image_processor.preprocess(image, height=height, width=width).to(dtype=torch.float32)
1139
+
1140
+ if isinstance(controlnet, ControlNetModel):
1141
+ control_image = self.prepare_control_image(
1142
+ image=control_image,
1143
+ width=width,
1144
+ height=height,
1145
+ batch_size=batch_size * num_images_per_prompt,
1146
+ num_images_per_prompt=num_images_per_prompt,
1147
+ device=device,
1148
+ dtype=controlnet.dtype,
1149
+ do_classifier_free_guidance=self.do_classifier_free_guidance,
1150
+ guess_mode=guess_mode,
1151
+ )
1152
+ height, width = control_image.shape[-2:]
1153
+ elif isinstance(controlnet, MultiControlNetModel):
1154
+ control_images = []
1155
+
1156
+ for control_image_ in control_image:
1157
+ control_image_ = self.prepare_control_image(
1158
+ image=control_image_,
1159
+ width=width,
1160
+ height=height,
1161
+ batch_size=batch_size * num_images_per_prompt,
1162
+ num_images_per_prompt=num_images_per_prompt,
1163
+ device=device,
1164
+ dtype=controlnet.dtype,
1165
+ do_classifier_free_guidance=self.do_classifier_free_guidance,
1166
+ guess_mode=guess_mode,
1167
+ )
1168
+
1169
+ control_images.append(control_image_)
1170
+
1171
+ control_image = control_images
1172
+ height, width = control_image[0].shape[-2:]
1173
+ else:
1174
+ assert False
1175
+
1176
+ # 5. Prepare timesteps
1177
+ self.scheduler.set_timesteps(num_inference_steps, device=device)
1178
+ timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device)
1179
+ latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
1180
+ self._num_timesteps = len(timesteps)
1181
+
1182
+ # 6. Prepare latent variables
1183
+
1184
+ num_channels_latents = self.unet.config.in_channels
1185
+ if latents is None:
1186
+ if strength >= 1.0:
1187
+ latents = self.prepare_latents_t2i(
1188
+ batch_size * num_images_per_prompt,
1189
+ num_channels_latents,
1190
+ height,
1191
+ width,
1192
+ prompt_embeds.dtype,
1193
+ device,
1194
+ generator,
1195
+ latents,
1196
+ )
1197
+ else:
1198
+ latents = self.prepare_latents(
1199
+ image,
1200
+ latent_timestep,
1201
+ batch_size,
1202
+ num_images_per_prompt,
1203
+ prompt_embeds.dtype,
1204
+ device,
1205
+ generator,
1206
+ True,
1207
+ )
1208
+
1209
+
1210
+ # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
1211
+ extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
1212
+
1213
+ # 7.1 Create tensor stating which controlnets to keep
1214
+ controlnet_keep = []
1215
+ for i in range(len(timesteps)):
1216
+ keeps = [
1217
+ 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e)
1218
+ for s, e in zip(control_guidance_start, control_guidance_end)
1219
+ ]
1220
+ controlnet_keep.append(keeps[0] if isinstance(controlnet, ControlNetModel) else keeps)
1221
+
1222
+ # 7.2 Prepare added time ids & embeddings
1223
+ if isinstance(control_image, list):
1224
+ original_size = original_size or control_image[0].shape[-2:]
1225
+ else:
1226
+ original_size = original_size or control_image.shape[-2:]
1227
+ target_size = target_size or (height, width)
1228
+
1229
+ # 7. Prepare added time ids & embeddings
1230
+ add_text_embeds = pooled_prompt_embeds
1231
+ add_time_ids = self._get_add_time_ids(
1232
+ original_size, crops_coords_top_left, target_size, dtype=prompt_embeds.dtype
1233
+ )
1234
+
1235
+ if self.do_classifier_free_guidance:
1236
+ prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
1237
+ add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)
1238
+ add_time_ids = torch.cat([add_time_ids, add_time_ids], dim=0)
1239
+
1240
+ prompt_embeds = prompt_embeds.to(device)
1241
+ add_text_embeds = add_text_embeds.to(device)
1242
+ add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1)
1243
+
1244
+ # 8. Denoising loop
1245
+ num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
1246
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
1247
+ for i, t in enumerate(timesteps):
1248
+ # expand the latents if we are doing classifier free guidance
1249
+ latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
1250
+ latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
1251
+
1252
+ added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
1253
+
1254
+ # controlnet(s) inference
1255
+ if guess_mode and self.do_classifier_free_guidance:
1256
+ # Infer ControlNet only for the conditional batch.
1257
+ control_model_input = latents
1258
+ control_model_input = self.scheduler.scale_model_input(control_model_input, t)
1259
+ controlnet_prompt_embeds = prompt_embeds.chunk(2)[1]
1260
+ controlnet_added_cond_kwargs = {
1261
+ "text_embeds": add_text_embeds.chunk(2)[1],
1262
+ "time_ids": add_time_ids.chunk(2)[1],
1263
+ }
1264
+ else:
1265
+ control_model_input = latent_model_input
1266
+ controlnet_prompt_embeds = prompt_embeds
1267
+ controlnet_added_cond_kwargs = added_cond_kwargs
1268
+
1269
+ if isinstance(controlnet_keep[i], list):
1270
+ cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])]
1271
+ else:
1272
+ controlnet_cond_scale = controlnet_conditioning_scale
1273
+ if isinstance(controlnet_cond_scale, list):
1274
+ controlnet_cond_scale = controlnet_cond_scale[0]
1275
+ cond_scale = controlnet_cond_scale * controlnet_keep[i]
1276
+
1277
+ down_block_res_samples, mid_block_res_sample = self.controlnet(
1278
+ control_model_input,
1279
+ t,
1280
+ encoder_hidden_states=controlnet_prompt_embeds,
1281
+ controlnet_cond=control_image,
1282
+ conditioning_scale=cond_scale,
1283
+ guess_mode=guess_mode,
1284
+ added_cond_kwargs=controlnet_added_cond_kwargs,
1285
+ return_dict=False,
1286
+ )
1287
+
1288
+ if guess_mode and self.do_classifier_free_guidance:
1289
+ # Infered ControlNet only for the conditional batch.
1290
+ # To apply the output of ControlNet to both the unconditional and conditional batches,
1291
+ # add 0 to the unconditional batch to keep it unchanged.
1292
+ down_block_res_samples = [torch.cat([torch.zeros_like(d), d]) for d in down_block_res_samples]
1293
+ mid_block_res_sample = torch.cat([torch.zeros_like(mid_block_res_sample), mid_block_res_sample])
1294
+
1295
+ if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
1296
+ added_cond_kwargs["image_embeds"] = image_embeds
1297
+
1298
+ # predict the noise residual
1299
+ noise_pred = self.unet(
1300
+ latent_model_input,
1301
+ t,
1302
+ encoder_hidden_states=prompt_embeds,
1303
+ cross_attention_kwargs=self.cross_attention_kwargs,
1304
+ down_block_additional_residuals=down_block_res_samples,
1305
+ mid_block_additional_residual=mid_block_res_sample,
1306
+ added_cond_kwargs=added_cond_kwargs,
1307
+ return_dict=False,
1308
+ )[0]
1309
+
1310
+ # perform guidance
1311
+ if self.do_classifier_free_guidance:
1312
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
1313
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
1314
+
1315
+ # compute the previous noisy sample x_t -> x_t-1
1316
+ latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
1317
+
1318
+ # call the callback, if provided
1319
+ if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
1320
+ progress_bar.update()
1321
+ if callback is not None and i % callback_steps == 0:
1322
+ step_idx = i // getattr(self.scheduler, "order", 1)
1323
+ callback(step_idx, t, latents)
1324
+
1325
+ # If we do sequential model offloading, let's offload unet and controlnet
1326
+ # manually for max memory savings
1327
+ if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
1328
+ self.unet.to("cpu")
1329
+ self.controlnet.to("cpu")
1330
+ torch.cuda.empty_cache()
1331
+
1332
+ if not output_type == "latent":
1333
+ # make sure the VAE is in float32 mode, as it overflows in float16
1334
+ needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast
1335
+
1336
+ if needs_upcasting:
1337
+ self.upcast_vae()
1338
+ latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)
1339
+
1340
+ latents = latents / self.vae.config.scaling_factor
1341
+ image = self.vae.decode(latents, return_dict=False)[0]
1342
+
1343
+ # cast back to fp16 if needed
1344
+ if needs_upcasting:
1345
+ self.vae.to(dtype=torch.float16)
1346
+ else:
1347
+ image = latents
1348
+ return StableDiffusionXLPipelineOutput(images=image)
1349
+
1350
+ image = self.image_processor.postprocess(image, output_type=output_type)
1351
+
1352
+ # Offload all models
1353
+ self.maybe_free_model_hooks()
1354
+
1355
+ if not return_dict:
1356
+ return (image,)
1357
+
1358
+ return StableDiffusionXLPipelineOutput(images=image)