codeysun commited on
Commit
bcadd8d
·
verified ·
1 Parent(s): 438ec7a

Upload pipeline.py

Browse files
Files changed (1) hide show
  1. pipeline.py +1547 -0
pipeline.py ADDED
@@ -0,0 +1,1547 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ CLIPTokenizer,
27
+ CLIPVisionModelWithProjection,
28
+ )
29
+
30
+ from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
31
+ from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
32
+ from diffusers.loaders import (
33
+ FromSingleFileMixin,
34
+ IPAdapterMixin,
35
+ StableDiffusionLoraLoaderMixin,
36
+ TextualInversionLoaderMixin,
37
+ )
38
+ from diffusers.models import (
39
+ AutoencoderKL,
40
+ ControlNetModel,
41
+ ImageProjection,
42
+ UNet2DConditionModel,
43
+ )
44
+ from diffusers.models.lora import adjust_lora_scale_text_encoder
45
+ from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel
46
+ from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusionMixin
47
+ from diffusers.pipelines.stable_diffusion.pipeline_output import (
48
+ StableDiffusionPipelineOutput,
49
+ )
50
+ from diffusers.pipelines.stable_diffusion.safety_checker import (
51
+ StableDiffusionSafetyChecker,
52
+ )
53
+ from diffusers.schedulers import KarrasDiffusionSchedulers
54
+ from diffusers.utils import (
55
+ USE_PEFT_BACKEND,
56
+ deprecate,
57
+ logging,
58
+ replace_example_docstring,
59
+ scale_lora_layers,
60
+ unscale_lora_layers,
61
+ )
62
+ from diffusers.utils.torch_utils import (
63
+ is_compiled_module,
64
+ is_torch_version,
65
+ randn_tensor,
66
+ )
67
+
68
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
69
+
70
+
71
+ EXAMPLE_DOC_STRING = """
72
+ Examples:
73
+ ```py
74
+ >>> # !pip install opencv-python transformers accelerate
75
+ >>> from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
76
+ >>> from diffusers.utils import load_image
77
+ >>> import numpy as np
78
+ >>> import torch
79
+
80
+ >>> import cv2
81
+ >>> from PIL import Image
82
+
83
+ >>> # download an image
84
+ >>> image = load_image(
85
+ ... "https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png"
86
+ ... )
87
+ >>> image = np.array(image)
88
+
89
+ >>> # get canny image
90
+ >>> image = cv2.Canny(image, 100, 200)
91
+ >>> image = image[:, :, None]
92
+ >>> image = np.concatenate([image, image, image], axis=2)
93
+ >>> canny_image = Image.fromarray(image)
94
+
95
+ >>> # load control net and stable diffusion v1-5
96
+ >>> controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
97
+ >>> pipe = StableDiffusionControlNetPipeline.from_pretrained(
98
+ ... "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
99
+ ... )
100
+
101
+ >>> # speed up diffusion process with faster scheduler and memory optimization
102
+ >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
103
+ >>> # remove following line if xformers is not installed
104
+ >>> pipe.enable_xformers_memory_efficient_attention()
105
+
106
+ >>> pipe.enable_model_cpu_offload()
107
+
108
+ >>> # generate image
109
+ >>> generator = torch.manual_seed(0)
110
+ >>> image = pipe(
111
+ ... "futuristic-looking woman", num_inference_steps=20, generator=generator, image=canny_image
112
+ ... ).images[0]
113
+ ```
114
+ """
115
+
116
+
117
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
118
+ def retrieve_timesteps(
119
+ scheduler,
120
+ num_inference_steps: Optional[int] = None,
121
+ device: Optional[Union[str, torch.device]] = None,
122
+ timesteps: Optional[List[int]] = None,
123
+ sigmas: Optional[List[float]] = None,
124
+ **kwargs,
125
+ ):
126
+ r"""
127
+ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
128
+ custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
129
+
130
+ Args:
131
+ scheduler (`SchedulerMixin`):
132
+ The scheduler to get timesteps from.
133
+ num_inference_steps (`int`):
134
+ The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
135
+ must be `None`.
136
+ device (`str` or `torch.device`, *optional*):
137
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
138
+ timesteps (`List[int]`, *optional*):
139
+ Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
140
+ `num_inference_steps` and `sigmas` must be `None`.
141
+ sigmas (`List[float]`, *optional*):
142
+ Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
143
+ `num_inference_steps` and `timesteps` must be `None`.
144
+
145
+ Returns:
146
+ `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
147
+ second element is the number of inference steps.
148
+ """
149
+ if timesteps is not None and sigmas is not None:
150
+ raise ValueError(
151
+ "Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values"
152
+ )
153
+ if timesteps is not None:
154
+ accepts_timesteps = "timesteps" in set(
155
+ inspect.signature(scheduler.set_timesteps).parameters.keys()
156
+ )
157
+ if not accepts_timesteps:
158
+ raise ValueError(
159
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
160
+ f" timestep schedules. Please check whether you are using the correct scheduler."
161
+ )
162
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
163
+ timesteps = scheduler.timesteps
164
+ num_inference_steps = len(timesteps)
165
+ elif sigmas is not None:
166
+ accept_sigmas = "sigmas" in set(
167
+ inspect.signature(scheduler.set_timesteps).parameters.keys()
168
+ )
169
+ if not accept_sigmas:
170
+ raise ValueError(
171
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
172
+ f" sigmas schedules. Please check whether you are using the correct scheduler."
173
+ )
174
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
175
+ timesteps = scheduler.timesteps
176
+ num_inference_steps = len(timesteps)
177
+ else:
178
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
179
+ timesteps = scheduler.timesteps
180
+ return timesteps, num_inference_steps
181
+
182
+
183
+ class StableDiffusionControlNetPipeline(
184
+ DiffusionPipeline,
185
+ StableDiffusionMixin,
186
+ TextualInversionLoaderMixin,
187
+ StableDiffusionLoraLoaderMixin,
188
+ IPAdapterMixin,
189
+ FromSingleFileMixin,
190
+ ):
191
+ r"""
192
+ Pipeline for text-to-image generation using Stable Diffusion with ControlNet guidance.
193
+
194
+ This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
195
+ implemented for all pipelines (downloading, saving, running on a particular device, etc.).
196
+
197
+ The pipeline also inherits the following loading methods:
198
+ - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings
199
+ - [`~loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`] for loading LoRA weights
200
+ - [`~loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`] for saving LoRA weights
201
+ - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files
202
+ - [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters
203
+
204
+ Args:
205
+ vae ([`AutoencoderKL`]):
206
+ Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations.
207
+ text_encoder ([`~transformers.CLIPTextModel`]):
208
+ Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)).
209
+ tokenizer ([`~transformers.CLIPTokenizer`]):
210
+ A `CLIPTokenizer` to tokenize text.
211
+ unet ([`UNet2DConditionModel`]):
212
+ A `UNet2DConditionModel` to denoise the encoded image latents.
213
+ controlnet ([`ControlNetModel`] or `List[ControlNetModel]`):
214
+ Provides additional conditioning to the `unet` during the denoising process. If you set multiple
215
+ ControlNets as a list, the outputs from each ControlNet are added together to create one combined
216
+ additional conditioning.
217
+ scheduler ([`SchedulerMixin`]):
218
+ A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
219
+ [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
220
+ safety_checker ([`StableDiffusionSafetyChecker`]):
221
+ Classification module that estimates whether generated images could be considered offensive or harmful.
222
+ Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details
223
+ about a model's potential harms.
224
+ feature_extractor ([`~transformers.CLIPImageProcessor`]):
225
+ A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`.
226
+ """
227
+
228
+ model_cpu_offload_seq = "text_encoder->image_encoder->unet->vae"
229
+ _optional_components = ["safety_checker", "feature_extractor", "image_encoder"]
230
+ _exclude_from_cpu_offload = ["safety_checker"]
231
+ _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
232
+
233
+ def __init__(
234
+ self,
235
+ vae: AutoencoderKL,
236
+ text_encoder: CLIPTextModel,
237
+ tokenizer: CLIPTokenizer,
238
+ unet: UNet2DConditionModel,
239
+ controlnet: Union[
240
+ ControlNetModel,
241
+ List[ControlNetModel],
242
+ Tuple[ControlNetModel],
243
+ MultiControlNetModel,
244
+ ],
245
+ scheduler: KarrasDiffusionSchedulers,
246
+ safety_checker: StableDiffusionSafetyChecker,
247
+ feature_extractor: CLIPImageProcessor,
248
+ image_encoder: CLIPVisionModelWithProjection = None,
249
+ requires_safety_checker: bool = True,
250
+ ):
251
+ super().__init__()
252
+
253
+ if safety_checker is None and requires_safety_checker:
254
+ logger.warning(
255
+ f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"
256
+ " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"
257
+ " results in services or applications open to the public. Both the diffusers team and Hugging Face"
258
+ " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
259
+ " it only for use-cases that involve analyzing network behavior or auditing its results. For more"
260
+ " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."
261
+ )
262
+
263
+ if safety_checker is not None and feature_extractor is None:
264
+ raise ValueError(
265
+ "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"
266
+ " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."
267
+ )
268
+
269
+ if isinstance(controlnet, (list, tuple)):
270
+ controlnet = MultiControlNetModel(controlnet)
271
+
272
+ self.register_modules(
273
+ vae=vae,
274
+ text_encoder=text_encoder,
275
+ tokenizer=tokenizer,
276
+ unet=unet,
277
+ controlnet=controlnet,
278
+ scheduler=scheduler,
279
+ safety_checker=safety_checker,
280
+ feature_extractor=feature_extractor,
281
+ image_encoder=image_encoder,
282
+ )
283
+ self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
284
+ self.image_processor = VaeImageProcessor(
285
+ vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True
286
+ )
287
+ self.control_image_processor = VaeImageProcessor(
288
+ vae_scale_factor=self.vae_scale_factor,
289
+ do_convert_rgb=True,
290
+ do_normalize=False,
291
+ )
292
+ self.register_to_config(requires_safety_checker=requires_safety_checker)
293
+
294
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._encode_prompt
295
+ def _encode_prompt(
296
+ self,
297
+ prompt,
298
+ device,
299
+ num_images_per_prompt,
300
+ do_classifier_free_guidance,
301
+ negative_prompt=None,
302
+ prompt_embeds: Optional[torch.Tensor] = None,
303
+ negative_prompt_embeds: Optional[torch.Tensor] = None,
304
+ lora_scale: Optional[float] = None,
305
+ **kwargs,
306
+ ):
307
+ deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
308
+ deprecate("_encode_prompt()", "1.0.0", deprecation_message, standard_warn=False)
309
+
310
+ prompt_embeds_tuple = self.encode_prompt(
311
+ prompt=prompt,
312
+ device=device,
313
+ num_images_per_prompt=num_images_per_prompt,
314
+ do_classifier_free_guidance=do_classifier_free_guidance,
315
+ negative_prompt=negative_prompt,
316
+ prompt_embeds=prompt_embeds,
317
+ negative_prompt_embeds=negative_prompt_embeds,
318
+ lora_scale=lora_scale,
319
+ **kwargs,
320
+ )
321
+
322
+ # concatenate for backwards comp
323
+ prompt_embeds = torch.cat([prompt_embeds_tuple[1], prompt_embeds_tuple[0]])
324
+
325
+ return prompt_embeds
326
+
327
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_prompt
328
+ def encode_prompt(
329
+ self,
330
+ prompt,
331
+ device,
332
+ num_images_per_prompt,
333
+ do_classifier_free_guidance,
334
+ negative_prompt=None,
335
+ prompt_embeds: Optional[torch.Tensor] = None,
336
+ negative_prompt_embeds: Optional[torch.Tensor] = None,
337
+ lora_scale: Optional[float] = None,
338
+ clip_skip: Optional[int] = None,
339
+ ):
340
+ r"""
341
+ Encodes the prompt into text encoder hidden states.
342
+
343
+ Args:
344
+ prompt (`str` or `List[str]`, *optional*):
345
+ prompt to be encoded
346
+ device: (`torch.device`):
347
+ torch device
348
+ num_images_per_prompt (`int`):
349
+ number of images that should be generated per prompt
350
+ do_classifier_free_guidance (`bool`):
351
+ whether to use classifier free guidance or not
352
+ negative_prompt (`str` or `List[str]`, *optional*):
353
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass
354
+ `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
355
+ less than `1`).
356
+ prompt_embeds (`torch.Tensor`, *optional*):
357
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
358
+ provided, text embeddings will be generated from `prompt` input argument.
359
+ negative_prompt_embeds (`torch.Tensor`, *optional*):
360
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
361
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
362
+ argument.
363
+ lora_scale (`float`, *optional*):
364
+ A LoRA scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
365
+ clip_skip (`int`, *optional*):
366
+ Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
367
+ the output of the pre-final layer will be used for computing the prompt embeddings.
368
+ """
369
+ # set lora scale so that monkey patched LoRA
370
+ # function of text encoder can correctly access it
371
+ if lora_scale is not None and isinstance(self, StableDiffusionLoraLoaderMixin):
372
+ self._lora_scale = lora_scale
373
+
374
+ # dynamically adjust the LoRA scale
375
+ if not USE_PEFT_BACKEND:
376
+ adjust_lora_scale_text_encoder(self.text_encoder, lora_scale)
377
+ else:
378
+ scale_lora_layers(self.text_encoder, lora_scale)
379
+
380
+ if prompt is not None and isinstance(prompt, str):
381
+ batch_size = 1
382
+ elif prompt is not None and isinstance(prompt, list):
383
+ batch_size = len(prompt)
384
+ else:
385
+ batch_size = prompt_embeds.shape[0]
386
+
387
+ if prompt_embeds is None:
388
+ # textual inversion: process multi-vector tokens if necessary
389
+ if isinstance(self, TextualInversionLoaderMixin):
390
+ prompt = self.maybe_convert_prompt(prompt, self.tokenizer)
391
+
392
+ text_inputs = self.tokenizer(
393
+ prompt,
394
+ padding="max_length",
395
+ max_length=self.tokenizer.model_max_length,
396
+ truncation=True,
397
+ return_tensors="pt",
398
+ )
399
+ text_input_ids = text_inputs.input_ids
400
+ untruncated_ids = self.tokenizer(
401
+ prompt, padding="longest", return_tensors="pt"
402
+ ).input_ids
403
+
404
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[
405
+ -1
406
+ ] and not torch.equal(text_input_ids, untruncated_ids):
407
+ removed_text = self.tokenizer.batch_decode(
408
+ untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
409
+ )
410
+ logger.warning(
411
+ "The following part of your input was truncated because CLIP can only handle sequences up to"
412
+ f" {self.tokenizer.model_max_length} tokens: {removed_text}"
413
+ )
414
+
415
+ if (
416
+ hasattr(self.text_encoder.config, "use_attention_mask")
417
+ and self.text_encoder.config.use_attention_mask
418
+ ):
419
+ attention_mask = text_inputs.attention_mask.to(device)
420
+ else:
421
+ attention_mask = None
422
+
423
+ if clip_skip is None:
424
+ prompt_embeds = self.text_encoder(
425
+ text_input_ids.to(device), attention_mask=attention_mask
426
+ )
427
+ prompt_embeds = prompt_embeds[0]
428
+ else:
429
+ prompt_embeds = self.text_encoder(
430
+ text_input_ids.to(device),
431
+ attention_mask=attention_mask,
432
+ output_hidden_states=True,
433
+ )
434
+ # Access the `hidden_states` first, that contains a tuple of
435
+ # all the hidden states from the encoder layers. Then index into
436
+ # the tuple to access the hidden states from the desired layer.
437
+ prompt_embeds = prompt_embeds[-1][-(clip_skip + 1)]
438
+ # We also need to apply the final LayerNorm here to not mess with the
439
+ # representations. The `last_hidden_states` that we typically use for
440
+ # obtaining the final prompt representations passes through the LayerNorm
441
+ # layer.
442
+ prompt_embeds = self.text_encoder.text_model.final_layer_norm(
443
+ prompt_embeds
444
+ )
445
+
446
+ if self.text_encoder is not None:
447
+ prompt_embeds_dtype = self.text_encoder.dtype
448
+ elif self.unet is not None:
449
+ prompt_embeds_dtype = self.unet.dtype
450
+ else:
451
+ prompt_embeds_dtype = prompt_embeds.dtype
452
+
453
+ prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
454
+
455
+ bs_embed, seq_len, _ = prompt_embeds.shape
456
+ # duplicate text embeddings for each generation per prompt, using mps friendly method
457
+ prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
458
+ prompt_embeds = prompt_embeds.view(
459
+ bs_embed * num_images_per_prompt, seq_len, -1
460
+ )
461
+
462
+ # get unconditional embeddings for classifier free guidance
463
+ if do_classifier_free_guidance and negative_prompt_embeds is None:
464
+ uncond_tokens: List[str]
465
+ if negative_prompt is None:
466
+ uncond_tokens = [""] * batch_size
467
+ elif prompt is not None and type(prompt) is not type(negative_prompt):
468
+ raise TypeError(
469
+ f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
470
+ f" {type(prompt)}."
471
+ )
472
+ elif isinstance(negative_prompt, str):
473
+ uncond_tokens = [negative_prompt]
474
+ elif batch_size != len(negative_prompt):
475
+ raise ValueError(
476
+ f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
477
+ f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
478
+ " the batch size of `prompt`."
479
+ )
480
+ else:
481
+ uncond_tokens = negative_prompt
482
+
483
+ # textual inversion: process multi-vector tokens if necessary
484
+ if isinstance(self, TextualInversionLoaderMixin):
485
+ uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)
486
+
487
+ max_length = prompt_embeds.shape[1]
488
+ uncond_input = self.tokenizer(
489
+ uncond_tokens,
490
+ padding="max_length",
491
+ max_length=max_length,
492
+ truncation=True,
493
+ return_tensors="pt",
494
+ )
495
+
496
+ if (
497
+ hasattr(self.text_encoder.config, "use_attention_mask")
498
+ and self.text_encoder.config.use_attention_mask
499
+ ):
500
+ attention_mask = uncond_input.attention_mask.to(device)
501
+ else:
502
+ attention_mask = None
503
+
504
+ negative_prompt_embeds = self.text_encoder(
505
+ uncond_input.input_ids.to(device),
506
+ attention_mask=attention_mask,
507
+ )
508
+ negative_prompt_embeds = negative_prompt_embeds[0]
509
+
510
+ if do_classifier_free_guidance:
511
+ # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
512
+ seq_len = negative_prompt_embeds.shape[1]
513
+
514
+ negative_prompt_embeds = negative_prompt_embeds.to(
515
+ dtype=prompt_embeds_dtype, device=device
516
+ )
517
+
518
+ negative_prompt_embeds = negative_prompt_embeds.repeat(
519
+ 1, num_images_per_prompt, 1
520
+ )
521
+ negative_prompt_embeds = negative_prompt_embeds.view(
522
+ batch_size * num_images_per_prompt, seq_len, -1
523
+ )
524
+
525
+ if self.text_encoder is not None:
526
+ if isinstance(self, StableDiffusionLoraLoaderMixin) and USE_PEFT_BACKEND:
527
+ # Retrieve the original scale by scaling back the LoRA layers
528
+ unscale_lora_layers(self.text_encoder, lora_scale)
529
+
530
+ return prompt_embeds, negative_prompt_embeds
531
+
532
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image
533
+ def encode_image(
534
+ self, image, device, num_images_per_prompt, output_hidden_states=None
535
+ ):
536
+ dtype = next(self.image_encoder.parameters()).dtype
537
+
538
+ if not isinstance(image, torch.Tensor):
539
+ image = self.feature_extractor(image, return_tensors="pt").pixel_values
540
+
541
+ image = image.to(device=device, dtype=dtype)
542
+ if output_hidden_states:
543
+ image_enc_hidden_states = self.image_encoder(
544
+ image, output_hidden_states=True
545
+ ).hidden_states[-2]
546
+ image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(
547
+ num_images_per_prompt, dim=0
548
+ )
549
+ uncond_image_enc_hidden_states = self.image_encoder(
550
+ torch.zeros_like(image), output_hidden_states=True
551
+ ).hidden_states[-2]
552
+ uncond_image_enc_hidden_states = (
553
+ uncond_image_enc_hidden_states.repeat_interleave(
554
+ num_images_per_prompt, dim=0
555
+ )
556
+ )
557
+ return image_enc_hidden_states, uncond_image_enc_hidden_states
558
+ else:
559
+ image_embeds = self.image_encoder(image).image_embeds
560
+ image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
561
+ uncond_image_embeds = torch.zeros_like(image_embeds)
562
+
563
+ return image_embeds, uncond_image_embeds
564
+
565
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds
566
+ def prepare_ip_adapter_image_embeds(
567
+ self,
568
+ ip_adapter_image,
569
+ ip_adapter_image_embeds,
570
+ device,
571
+ num_images_per_prompt,
572
+ do_classifier_free_guidance,
573
+ ):
574
+ image_embeds = []
575
+ if do_classifier_free_guidance:
576
+ negative_image_embeds = []
577
+ if ip_adapter_image_embeds is None:
578
+ if not isinstance(ip_adapter_image, list):
579
+ ip_adapter_image = [ip_adapter_image]
580
+
581
+ if len(ip_adapter_image) != len(
582
+ self.unet.encoder_hid_proj.image_projection_layers
583
+ ):
584
+ raise ValueError(
585
+ 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."
586
+ )
587
+
588
+ for single_ip_adapter_image, image_proj_layer in zip(
589
+ ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers
590
+ ):
591
+ output_hidden_state = not isinstance(image_proj_layer, ImageProjection)
592
+ single_image_embeds, single_negative_image_embeds = self.encode_image(
593
+ single_ip_adapter_image, device, 1, output_hidden_state
594
+ )
595
+
596
+ image_embeds.append(single_image_embeds[None, :])
597
+ if do_classifier_free_guidance:
598
+ negative_image_embeds.append(single_negative_image_embeds[None, :])
599
+ else:
600
+ for single_image_embeds in ip_adapter_image_embeds:
601
+ if do_classifier_free_guidance:
602
+ (
603
+ single_negative_image_embeds,
604
+ single_image_embeds,
605
+ ) = single_image_embeds.chunk(2)
606
+ negative_image_embeds.append(single_negative_image_embeds)
607
+ image_embeds.append(single_image_embeds)
608
+
609
+ ip_adapter_image_embeds = []
610
+ for i, single_image_embeds in enumerate(image_embeds):
611
+ single_image_embeds = torch.cat(
612
+ [single_image_embeds] * num_images_per_prompt, dim=0
613
+ )
614
+ if do_classifier_free_guidance:
615
+ single_negative_image_embeds = torch.cat(
616
+ [negative_image_embeds[i]] * num_images_per_prompt, dim=0
617
+ )
618
+ single_image_embeds = torch.cat(
619
+ [single_negative_image_embeds, single_image_embeds], dim=0
620
+ )
621
+
622
+ single_image_embeds = single_image_embeds.to(device=device)
623
+ ip_adapter_image_embeds.append(single_image_embeds)
624
+
625
+ return ip_adapter_image_embeds
626
+
627
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker
628
+ def run_safety_checker(self, image, device, dtype):
629
+ if self.safety_checker is None:
630
+ has_nsfw_concept = None
631
+ else:
632
+ if torch.is_tensor(image):
633
+ feature_extractor_input = self.image_processor.postprocess(
634
+ image, output_type="pil"
635
+ )
636
+ else:
637
+ feature_extractor_input = self.image_processor.numpy_to_pil(image)
638
+ safety_checker_input = self.feature_extractor(
639
+ feature_extractor_input, return_tensors="pt"
640
+ ).to(device)
641
+ image, has_nsfw_concept = self.safety_checker(
642
+ images=image, clip_input=safety_checker_input.pixel_values.to(dtype)
643
+ )
644
+ return image, has_nsfw_concept
645
+
646
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.decode_latents
647
+ def decode_latents(self, latents):
648
+ deprecation_message = "The decode_latents method is deprecated and will be removed in 1.0.0. Please use VaeImageProcessor.postprocess(...) instead"
649
+ deprecate("decode_latents", "1.0.0", deprecation_message, standard_warn=False)
650
+
651
+ latents = 1 / self.vae.config.scaling_factor * latents
652
+ image = self.vae.decode(latents, return_dict=False)[0]
653
+ image = (image / 2 + 0.5).clamp(0, 1)
654
+ # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
655
+ image = image.cpu().permute(0, 2, 3, 1).float().numpy()
656
+ return image
657
+
658
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
659
+ def prepare_extra_step_kwargs(self, generator, eta):
660
+ # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
661
+ # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
662
+ # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
663
+ # and should be between [0, 1]
664
+
665
+ accepts_eta = "eta" in set(
666
+ inspect.signature(self.scheduler.step).parameters.keys()
667
+ )
668
+ extra_step_kwargs = {}
669
+ if accepts_eta:
670
+ extra_step_kwargs["eta"] = eta
671
+
672
+ # check if the scheduler accepts generator
673
+ accepts_generator = "generator" in set(
674
+ inspect.signature(self.scheduler.step).parameters.keys()
675
+ )
676
+ if accepts_generator:
677
+ extra_step_kwargs["generator"] = generator
678
+ return extra_step_kwargs
679
+
680
+ def check_inputs(
681
+ self,
682
+ prompt,
683
+ image,
684
+ callback_steps,
685
+ negative_prompt=None,
686
+ prompt_embeds=None,
687
+ negative_prompt_embeds=None,
688
+ ip_adapter_image=None,
689
+ ip_adapter_image_embeds=None,
690
+ controlnet_conditioning_scale=1.0,
691
+ control_guidance_start=0.0,
692
+ control_guidance_end=1.0,
693
+ callback_on_step_end_tensor_inputs=None,
694
+ ):
695
+ if callback_steps is not None and (
696
+ not isinstance(callback_steps, int) or callback_steps <= 0
697
+ ):
698
+ raise ValueError(
699
+ f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
700
+ f" {type(callback_steps)}."
701
+ )
702
+
703
+ if callback_on_step_end_tensor_inputs is not None and not all(
704
+ k in self._callback_tensor_inputs
705
+ for k in callback_on_step_end_tensor_inputs
706
+ ):
707
+ raise ValueError(
708
+ 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]}"
709
+ )
710
+
711
+ if prompt is not None and prompt_embeds is not None:
712
+ raise ValueError(
713
+ f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
714
+ " only forward one of the two."
715
+ )
716
+ elif prompt is None and prompt_embeds is None:
717
+ raise ValueError(
718
+ "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
719
+ )
720
+ elif prompt is not None and (
721
+ not isinstance(prompt, str) and not isinstance(prompt, list)
722
+ ):
723
+ raise ValueError(
724
+ f"`prompt` has to be of type `str` or `list` but is {type(prompt)}"
725
+ )
726
+
727
+ if negative_prompt is not None and negative_prompt_embeds is not None:
728
+ raise ValueError(
729
+ f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
730
+ f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
731
+ )
732
+
733
+ if prompt_embeds is not None and negative_prompt_embeds is not None:
734
+ if prompt_embeds.shape != negative_prompt_embeds.shape:
735
+ raise ValueError(
736
+ "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
737
+ f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
738
+ f" {negative_prompt_embeds.shape}."
739
+ )
740
+
741
+ # Check `image`
742
+ is_compiled = hasattr(F, "scaled_dot_product_attention") and isinstance(
743
+ self.controlnet, torch._dynamo.eval_frame.OptimizedModule
744
+ )
745
+ if (
746
+ isinstance(self.controlnet, ControlNetModel)
747
+ or is_compiled
748
+ and isinstance(self.controlnet._orig_mod, ControlNetModel)
749
+ ):
750
+ self.check_image(image, prompt, prompt_embeds)
751
+ elif (
752
+ isinstance(self.controlnet, MultiControlNetModel)
753
+ or is_compiled
754
+ and isinstance(self.controlnet._orig_mod, MultiControlNetModel)
755
+ ):
756
+ if not isinstance(image, list):
757
+ raise TypeError("For multiple controlnets: `image` must be type `list`")
758
+
759
+ # When `image` is a nested list:
760
+ # (e.g. [[canny_image_1, pose_image_1], [canny_image_2, pose_image_2]])
761
+ elif any(isinstance(i, list) for i in image):
762
+ transposed_image = [list(t) for t in zip(*image)]
763
+ if len(transposed_image) != len(self.controlnet.nets):
764
+ raise ValueError(
765
+ f"For multiple controlnets: if you pass`image` as a list of list, each sublist must have the same length as the number of controlnets, but the sublists in `image` got {len(transposed_image)} images and {len(self.controlnet.nets)} ControlNets."
766
+ )
767
+ for image_ in transposed_image:
768
+ self.check_image(image_, prompt, prompt_embeds)
769
+ elif len(image) != len(self.controlnet.nets):
770
+ raise ValueError(
771
+ 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."
772
+ )
773
+ else:
774
+ for image_ in image:
775
+ self.check_image(image_, prompt, prompt_embeds)
776
+ else:
777
+ assert False
778
+
779
+ # Check `controlnet_conditioning_scale`
780
+ if (
781
+ isinstance(self.controlnet, ControlNetModel)
782
+ or is_compiled
783
+ and isinstance(self.controlnet._orig_mod, ControlNetModel)
784
+ ):
785
+ if not isinstance(controlnet_conditioning_scale, float):
786
+ raise TypeError(
787
+ "For single controlnet: `controlnet_conditioning_scale` must be type `float`."
788
+ )
789
+ elif (
790
+ isinstance(self.controlnet, MultiControlNetModel)
791
+ or is_compiled
792
+ and isinstance(self.controlnet._orig_mod, MultiControlNetModel)
793
+ ):
794
+ if isinstance(controlnet_conditioning_scale, list):
795
+ if any(isinstance(i, list) for i in controlnet_conditioning_scale):
796
+ raise ValueError(
797
+ "A single batch of varying conditioning scale settings (e.g. [[1.0, 0.5], [0.2, 0.8]]) is not supported at the moment. "
798
+ "The conditioning scale must be fixed across the batch."
799
+ )
800
+ elif isinstance(controlnet_conditioning_scale, list) and len(
801
+ controlnet_conditioning_scale
802
+ ) != len(self.controlnet.nets):
803
+ raise ValueError(
804
+ "For multiple controlnets: When `controlnet_conditioning_scale` is specified as `list`, it must have"
805
+ " the same length as the number of controlnets"
806
+ )
807
+ else:
808
+ assert False
809
+
810
+ if not isinstance(control_guidance_start, (tuple, list)):
811
+ control_guidance_start = [control_guidance_start]
812
+
813
+ if not isinstance(control_guidance_end, (tuple, list)):
814
+ control_guidance_end = [control_guidance_end]
815
+
816
+ if len(control_guidance_start) != len(control_guidance_end):
817
+ raise ValueError(
818
+ 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."
819
+ )
820
+
821
+ if isinstance(self.controlnet, MultiControlNetModel):
822
+ if len(control_guidance_start) != len(self.controlnet.nets):
823
+ raise ValueError(
824
+ 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)}."
825
+ )
826
+
827
+ for start, end in zip(control_guidance_start, control_guidance_end):
828
+ if start >= end:
829
+ raise ValueError(
830
+ f"control guidance start: {start} cannot be larger or equal to control guidance end: {end}."
831
+ )
832
+ if start < 0.0:
833
+ raise ValueError(
834
+ f"control guidance start: {start} can't be smaller than 0."
835
+ )
836
+ if end > 1.0:
837
+ raise ValueError(
838
+ f"control guidance end: {end} can't be larger than 1.0."
839
+ )
840
+
841
+ if ip_adapter_image is not None and ip_adapter_image_embeds is not None:
842
+ raise ValueError(
843
+ "Provide either `ip_adapter_image` or `ip_adapter_image_embeds`. Cannot leave both `ip_adapter_image` and `ip_adapter_image_embeds` defined."
844
+ )
845
+
846
+ if ip_adapter_image_embeds is not None:
847
+ if not isinstance(ip_adapter_image_embeds, list):
848
+ raise ValueError(
849
+ f"`ip_adapter_image_embeds` has to be of type `list` but is {type(ip_adapter_image_embeds)}"
850
+ )
851
+ elif ip_adapter_image_embeds[0].ndim not in [3, 4]:
852
+ raise ValueError(
853
+ f"`ip_adapter_image_embeds` has to be a list of 3D or 4D tensors but is {ip_adapter_image_embeds[0].ndim}D"
854
+ )
855
+
856
+ def check_image(self, image, prompt, prompt_embeds):
857
+ image_is_pil = isinstance(image, PIL.Image.Image)
858
+ image_is_tensor = isinstance(image, torch.Tensor)
859
+ image_is_np = isinstance(image, np.ndarray)
860
+ image_is_pil_list = isinstance(image, list) and isinstance(
861
+ image[0], PIL.Image.Image
862
+ )
863
+ image_is_tensor_list = isinstance(image, list) and isinstance(
864
+ image[0], torch.Tensor
865
+ )
866
+ image_is_np_list = isinstance(image, list) and isinstance(image[0], np.ndarray)
867
+
868
+ if (
869
+ not image_is_pil
870
+ and not image_is_tensor
871
+ and not image_is_np
872
+ and not image_is_pil_list
873
+ and not image_is_tensor_list
874
+ and not image_is_np_list
875
+ ):
876
+ raise TypeError(
877
+ 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)}"
878
+ )
879
+
880
+ if image_is_pil:
881
+ image_batch_size = 1
882
+ else:
883
+ image_batch_size = len(image)
884
+
885
+ if prompt is not None and isinstance(prompt, str):
886
+ prompt_batch_size = 1
887
+ elif prompt is not None and isinstance(prompt, list):
888
+ prompt_batch_size = len(prompt)
889
+ elif prompt_embeds is not None:
890
+ prompt_batch_size = prompt_embeds.shape[0]
891
+
892
+ if image_batch_size != 1 and image_batch_size != prompt_batch_size:
893
+ raise ValueError(
894
+ 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}"
895
+ )
896
+
897
+ def prepare_image(
898
+ self,
899
+ image,
900
+ width,
901
+ height,
902
+ batch_size,
903
+ num_images_per_prompt,
904
+ device,
905
+ dtype,
906
+ do_classifier_free_guidance=False,
907
+ guess_mode=False,
908
+ ):
909
+ image = self.control_image_processor.preprocess(
910
+ image, height=height, width=width
911
+ ).to(dtype=torch.float32)
912
+ image_batch_size = image.shape[0]
913
+
914
+ if image_batch_size == 1:
915
+ repeat_by = batch_size
916
+ else:
917
+ # image batch size is the same as prompt batch size
918
+ repeat_by = num_images_per_prompt
919
+
920
+ image = image.repeat_interleave(repeat_by, dim=0)
921
+
922
+ image = image.to(device=device, dtype=dtype)
923
+
924
+ if do_classifier_free_guidance and not guess_mode:
925
+ image = torch.cat([image] * 2)
926
+
927
+ return image
928
+
929
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents
930
+ def prepare_latents(
931
+ self,
932
+ batch_size,
933
+ num_channels_latents,
934
+ height,
935
+ width,
936
+ dtype,
937
+ device,
938
+ generator,
939
+ latents=None,
940
+ ):
941
+ shape = (
942
+ batch_size,
943
+ num_channels_latents,
944
+ int(height) // self.vae_scale_factor,
945
+ int(width) // self.vae_scale_factor,
946
+ )
947
+ if isinstance(generator, list) and len(generator) != batch_size:
948
+ raise ValueError(
949
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
950
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
951
+ )
952
+
953
+ if latents is None:
954
+ latents = randn_tensor(
955
+ shape, generator=generator, device=device, dtype=dtype
956
+ )
957
+ else:
958
+ latents = latents.to(device)
959
+
960
+ # scale the initial noise by the standard deviation required by the scheduler
961
+ latents = latents * self.scheduler.init_noise_sigma
962
+ return latents
963
+
964
+ # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding
965
+ def get_guidance_scale_embedding(
966
+ self,
967
+ w: torch.Tensor,
968
+ embedding_dim: int = 512,
969
+ dtype: torch.dtype = torch.float32,
970
+ ) -> torch.Tensor:
971
+ """
972
+ See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298
973
+
974
+ Args:
975
+ w (`torch.Tensor`):
976
+ Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings.
977
+ embedding_dim (`int`, *optional*, defaults to 512):
978
+ Dimension of the embeddings to generate.
979
+ dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):
980
+ Data type of the generated embeddings.
981
+
982
+ Returns:
983
+ `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`.
984
+ """
985
+ assert len(w.shape) == 1
986
+ w = w * 1000.0
987
+
988
+ half_dim = embedding_dim // 2
989
+ emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
990
+ emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)
991
+ emb = w.to(dtype)[:, None] * emb[None, :]
992
+ emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
993
+ if embedding_dim % 2 == 1: # zero pad
994
+ emb = torch.nn.functional.pad(emb, (0, 1))
995
+ assert emb.shape == (w.shape[0], embedding_dim)
996
+ return emb
997
+
998
+ @property
999
+ def guidance_scale(self):
1000
+ return self._guidance_scale
1001
+
1002
+ @property
1003
+ def clip_skip(self):
1004
+ return self._clip_skip
1005
+
1006
+ # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
1007
+ # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
1008
+ # corresponds to doing no classifier free guidance.
1009
+ @property
1010
+ def do_classifier_free_guidance(self):
1011
+ return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None
1012
+
1013
+ @property
1014
+ def cross_attention_kwargs(self):
1015
+ return self._cross_attention_kwargs
1016
+
1017
+ @property
1018
+ def num_timesteps(self):
1019
+ return self._num_timesteps
1020
+
1021
+ @property
1022
+ def interrupt(self):
1023
+ return self._interrupt
1024
+
1025
+ @torch.no_grad()
1026
+ @replace_example_docstring(EXAMPLE_DOC_STRING)
1027
+ def __call__(
1028
+ self,
1029
+ prompt: Union[str, List[str]] = None,
1030
+ image: PipelineImageInput = None,
1031
+ height: Optional[int] = None,
1032
+ width: Optional[int] = None,
1033
+ num_inference_steps: int = 50,
1034
+ timesteps: List[int] = None,
1035
+ sigmas: List[float] = None,
1036
+ guidance_scale: float = 7.5,
1037
+ negative_prompt: Optional[Union[str, List[str]]] = None,
1038
+ num_images_per_prompt: Optional[int] = 1,
1039
+ eta: float = 0.0,
1040
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
1041
+ latents: Optional[torch.Tensor] = None,
1042
+ prompt_embeds: Optional[torch.Tensor] = None,
1043
+ negative_prompt_embeds: Optional[torch.Tensor] = None,
1044
+ ip_adapter_image: Optional[PipelineImageInput] = None,
1045
+ ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
1046
+ output_type: Optional[str] = "pil",
1047
+ return_dict: bool = True,
1048
+ cross_attention_kwargs: Optional[Dict[str, Any]] = None,
1049
+ controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
1050
+ guess_mode: bool = False,
1051
+ control_guidance_start: Union[float, List[float]] = 0.0,
1052
+ control_guidance_end: Union[float, List[float]] = 1.0,
1053
+ clip_skip: Optional[int] = None,
1054
+ callback_on_step_end: Optional[
1055
+ Union[
1056
+ Callable[[int, int, Dict], None],
1057
+ PipelineCallback,
1058
+ MultiPipelineCallbacks,
1059
+ ]
1060
+ ] = None,
1061
+ callback_on_step_end_tensor_inputs: List[str] = ["latents"],
1062
+ **kwargs,
1063
+ ):
1064
+ r"""
1065
+ The call function to the pipeline for generation.
1066
+
1067
+ Args:
1068
+ prompt (`str` or `List[str]`, *optional*):
1069
+ The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
1070
+ image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
1071
+ `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
1072
+ The ControlNet input condition to provide guidance to the `unet` for generation. If the type is
1073
+ specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted
1074
+ as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or
1075
+ width are passed, `image` is resized accordingly. If multiple ControlNets are specified in `init`,
1076
+ images must be passed as a list such that each element of the list can be correctly batched for input
1077
+ to a single ControlNet. When `prompt` is a list, and if a list of images is passed for a single
1078
+ ControlNet, each will be paired with each prompt in the `prompt` list. This also applies to multiple
1079
+ ControlNets, where a list of image lists can be passed to batch for each prompt and each ControlNet.
1080
+ height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
1081
+ The height in pixels of the generated image.
1082
+ width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
1083
+ The width in pixels of the generated image.
1084
+ num_inference_steps (`int`, *optional*, defaults to 50):
1085
+ The number of denoising steps. More denoising steps usually lead to a higher quality image at the
1086
+ expense of slower inference.
1087
+ timesteps (`List[int]`, *optional*):
1088
+ Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
1089
+ in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
1090
+ passed will be used. Must be in descending order.
1091
+ sigmas (`List[float]`, *optional*):
1092
+ Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
1093
+ their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
1094
+ will be used.
1095
+ guidance_scale (`float`, *optional*, defaults to 7.5):
1096
+ A higher guidance scale value encourages the model to generate images closely linked to the text
1097
+ `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
1098
+ negative_prompt (`str` or `List[str]`, *optional*):
1099
+ The prompt or prompts to guide what to not include in image generation. If not defined, you need to
1100
+ pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
1101
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
1102
+ The number of images to generate per prompt.
1103
+ eta (`float`, *optional*, defaults to 0.0):
1104
+ Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies
1105
+ to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
1106
+ generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
1107
+ A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
1108
+ generation deterministic.
1109
+ latents (`torch.Tensor`, *optional*):
1110
+ Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
1111
+ generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
1112
+ tensor is generated by sampling using the supplied random `generator`.
1113
+ prompt_embeds (`torch.Tensor`, *optional*):
1114
+ Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
1115
+ provided, text embeddings are generated from the `prompt` input argument.
1116
+ negative_prompt_embeds (`torch.Tensor`, *optional*):
1117
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
1118
+ not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
1119
+ ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
1120
+ ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
1121
+ Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
1122
+ IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should
1123
+ contain the negative image embedding if `do_classifier_free_guidance` is set to `True`. If not
1124
+ provided, embeddings are computed from the `ip_adapter_image` input argument.
1125
+ output_type (`str`, *optional*, defaults to `"pil"`):
1126
+ The output format of the generated image. Choose between `PIL.Image` or `np.array`.
1127
+ return_dict (`bool`, *optional*, defaults to `True`):
1128
+ Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
1129
+ plain tuple.
1130
+ callback (`Callable`, *optional*):
1131
+ A function that calls every `callback_steps` steps during inference. The function is called with the
1132
+ following arguments: `callback(step: int, timestep: int, latents: torch.Tensor)`.
1133
+ callback_steps (`int`, *optional*, defaults to 1):
1134
+ The frequency at which the `callback` function is called. If not specified, the callback is called at
1135
+ every step.
1136
+ cross_attention_kwargs (`dict`, *optional*):
1137
+ A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
1138
+ [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
1139
+ controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
1140
+ The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
1141
+ to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set
1142
+ the corresponding scale as a list.
1143
+ guess_mode (`bool`, *optional*, defaults to `False`):
1144
+ The ControlNet encoder tries to recognize the content of the input image even if you remove all
1145
+ prompts. A `guidance_scale` value between 3.0 and 5.0 is recommended.
1146
+ control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
1147
+ The percentage of total steps at which the ControlNet starts applying.
1148
+ control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
1149
+ The percentage of total steps at which the ControlNet stops applying.
1150
+ clip_skip (`int`, *optional*):
1151
+ Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
1152
+ the output of the pre-final layer will be used for computing the prompt embeddings.
1153
+ callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
1154
+ A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of
1155
+ each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
1156
+ DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
1157
+ list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
1158
+ callback_on_step_end_tensor_inputs (`List`, *optional*):
1159
+ The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
1160
+ will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
1161
+ `._callback_tensor_inputs` attribute of your pipeline class.
1162
+
1163
+ Examples:
1164
+
1165
+ Returns:
1166
+ [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
1167
+ If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned,
1168
+ otherwise a `tuple` is returned where the first element is a list with the generated images and the
1169
+ second element is a list of `bool`s indicating whether the corresponding generated image contains
1170
+ "not-safe-for-work" (nsfw) content.
1171
+ """
1172
+
1173
+ callback = kwargs.pop("callback", None)
1174
+ callback_steps = kwargs.pop("callback_steps", None)
1175
+
1176
+ if callback is not None:
1177
+ deprecate(
1178
+ "callback",
1179
+ "1.0.0",
1180
+ "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
1181
+ )
1182
+ if callback_steps is not None:
1183
+ deprecate(
1184
+ "callback_steps",
1185
+ "1.0.0",
1186
+ "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
1187
+ )
1188
+
1189
+ if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
1190
+ callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
1191
+
1192
+ controlnet = (
1193
+ self.controlnet._orig_mod
1194
+ if is_compiled_module(self.controlnet)
1195
+ else self.controlnet
1196
+ )
1197
+
1198
+ # align format for control guidance
1199
+ if not isinstance(control_guidance_start, list) and isinstance(
1200
+ control_guidance_end, list
1201
+ ):
1202
+ control_guidance_start = len(control_guidance_end) * [
1203
+ control_guidance_start
1204
+ ]
1205
+ elif not isinstance(control_guidance_end, list) and isinstance(
1206
+ control_guidance_start, list
1207
+ ):
1208
+ control_guidance_end = len(control_guidance_start) * [control_guidance_end]
1209
+ elif not isinstance(control_guidance_start, list) and not isinstance(
1210
+ control_guidance_end, list
1211
+ ):
1212
+ mult = (
1213
+ len(controlnet.nets)
1214
+ if isinstance(controlnet, MultiControlNetModel)
1215
+ else 1
1216
+ )
1217
+ control_guidance_start, control_guidance_end = (
1218
+ mult * [control_guidance_start],
1219
+ mult * [control_guidance_end],
1220
+ )
1221
+
1222
+ # 1. Check inputs. Raise error if not correct
1223
+ self.check_inputs(
1224
+ prompt,
1225
+ image,
1226
+ callback_steps,
1227
+ negative_prompt,
1228
+ prompt_embeds,
1229
+ negative_prompt_embeds,
1230
+ ip_adapter_image,
1231
+ ip_adapter_image_embeds,
1232
+ controlnet_conditioning_scale,
1233
+ control_guidance_start,
1234
+ control_guidance_end,
1235
+ callback_on_step_end_tensor_inputs,
1236
+ )
1237
+
1238
+ self._guidance_scale = guidance_scale
1239
+ self._clip_skip = clip_skip
1240
+ self._cross_attention_kwargs = cross_attention_kwargs
1241
+ self._interrupt = False
1242
+
1243
+ # 2. Define call parameters
1244
+ if prompt is not None and isinstance(prompt, str):
1245
+ batch_size = 1
1246
+ elif prompt is not None and isinstance(prompt, list):
1247
+ batch_size = len(prompt)
1248
+ else:
1249
+ batch_size = prompt_embeds.shape[0]
1250
+
1251
+ device = self._execution_device
1252
+
1253
+ if isinstance(controlnet, MultiControlNetModel) and isinstance(
1254
+ controlnet_conditioning_scale, float
1255
+ ):
1256
+ controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(
1257
+ controlnet.nets
1258
+ )
1259
+
1260
+ global_pool_conditions = (
1261
+ controlnet.config.global_pool_conditions
1262
+ if isinstance(controlnet, ControlNetModel)
1263
+ else controlnet.nets[0].config.global_pool_conditions
1264
+ )
1265
+ guess_mode = guess_mode or global_pool_conditions
1266
+
1267
+ # 3. Encode input prompt
1268
+ text_encoder_lora_scale = (
1269
+ self.cross_attention_kwargs.get("scale", None)
1270
+ if self.cross_attention_kwargs is not None
1271
+ else None
1272
+ )
1273
+ prompt_embeds, negative_prompt_embeds = self.encode_prompt(
1274
+ prompt,
1275
+ device,
1276
+ num_images_per_prompt,
1277
+ self.do_classifier_free_guidance,
1278
+ negative_prompt,
1279
+ prompt_embeds=prompt_embeds,
1280
+ negative_prompt_embeds=negative_prompt_embeds,
1281
+ lora_scale=text_encoder_lora_scale,
1282
+ clip_skip=self.clip_skip,
1283
+ )
1284
+ # For classifier free guidance, we need to do two forward passes.
1285
+ # Here we concatenate the unconditional and text embeddings into a single batch
1286
+ # to avoid doing two forward passes
1287
+ if self.do_classifier_free_guidance:
1288
+ prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
1289
+
1290
+ if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
1291
+ image_embeds = self.prepare_ip_adapter_image_embeds(
1292
+ ip_adapter_image,
1293
+ ip_adapter_image_embeds,
1294
+ device,
1295
+ batch_size * num_images_per_prompt,
1296
+ self.do_classifier_free_guidance,
1297
+ )
1298
+
1299
+ # 4. Prepare image
1300
+ if isinstance(controlnet, ControlNetModel):
1301
+ image = self.prepare_image(
1302
+ image=image,
1303
+ width=width,
1304
+ height=height,
1305
+ batch_size=batch_size * num_images_per_prompt,
1306
+ num_images_per_prompt=num_images_per_prompt,
1307
+ device=device,
1308
+ dtype=controlnet.dtype,
1309
+ do_classifier_free_guidance=self.do_classifier_free_guidance,
1310
+ guess_mode=guess_mode,
1311
+ )
1312
+ height, width = image.shape[-2:]
1313
+ elif isinstance(controlnet, MultiControlNetModel):
1314
+ images = []
1315
+
1316
+ # Nested lists as ControlNet condition
1317
+ if isinstance(image[0], list):
1318
+ # Transpose the nested image list
1319
+ image = [list(t) for t in zip(*image)]
1320
+
1321
+ for image_ in image:
1322
+ image_ = self.prepare_image(
1323
+ image=image_,
1324
+ width=width,
1325
+ height=height,
1326
+ batch_size=batch_size * num_images_per_prompt,
1327
+ num_images_per_prompt=num_images_per_prompt,
1328
+ device=device,
1329
+ dtype=controlnet.dtype,
1330
+ do_classifier_free_guidance=self.do_classifier_free_guidance,
1331
+ guess_mode=guess_mode,
1332
+ )
1333
+
1334
+ images.append(image_)
1335
+
1336
+ image = images
1337
+ height, width = image[0].shape[-2:]
1338
+ else:
1339
+ assert False
1340
+
1341
+ # 5. Prepare timesteps
1342
+ timesteps, num_inference_steps = retrieve_timesteps(
1343
+ self.scheduler, num_inference_steps, device, timesteps, sigmas
1344
+ )
1345
+ self._num_timesteps = len(timesteps)
1346
+
1347
+ # 6. Prepare latent variables
1348
+ num_channels_latents = self.unet.config.in_channels
1349
+ latents = self.prepare_latents(
1350
+ batch_size * num_images_per_prompt,
1351
+ num_channels_latents,
1352
+ height,
1353
+ width,
1354
+ prompt_embeds.dtype,
1355
+ device,
1356
+ generator,
1357
+ latents,
1358
+ )
1359
+
1360
+ # 6.5 Optionally get Guidance Scale Embedding
1361
+ timestep_cond = None
1362
+ if self.unet.config.time_cond_proj_dim is not None:
1363
+ guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(
1364
+ batch_size * num_images_per_prompt
1365
+ )
1366
+ timestep_cond = self.get_guidance_scale_embedding(
1367
+ guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
1368
+ ).to(device=device, dtype=latents.dtype)
1369
+
1370
+ # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
1371
+ extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
1372
+
1373
+ # 7.1 Add image embeds for IP-Adapter
1374
+ added_cond_kwargs = (
1375
+ {"image_embeds": image_embeds}
1376
+ if ip_adapter_image is not None or ip_adapter_image_embeds is not None
1377
+ else None
1378
+ )
1379
+
1380
+ # 7.2 Create tensor stating which controlnets to keep
1381
+ controlnet_keep = []
1382
+ for i in range(len(timesteps)):
1383
+ keeps = [
1384
+ 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e)
1385
+ for s, e in zip(control_guidance_start, control_guidance_end)
1386
+ ]
1387
+ controlnet_keep.append(
1388
+ keeps[0] if isinstance(controlnet, ControlNetModel) else keeps
1389
+ )
1390
+
1391
+ # 8. Denoising loop
1392
+ num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
1393
+ is_unet_compiled = is_compiled_module(self.unet)
1394
+ is_controlnet_compiled = is_compiled_module(self.controlnet)
1395
+ is_torch_higher_equal_2_1 = is_torch_version(">=", "2.1")
1396
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
1397
+ for i, t in enumerate(timesteps):
1398
+ if self.interrupt:
1399
+ continue
1400
+
1401
+ # Relevant thread:
1402
+ # https://dev-discuss.pytorch.org/t/cudagraphs-in-pytorch-2-0/1428
1403
+ if (
1404
+ is_unet_compiled and is_controlnet_compiled
1405
+ ) and is_torch_higher_equal_2_1:
1406
+ torch._inductor.cudagraph_mark_step_begin()
1407
+ # expand the latents if we are doing classifier free guidance
1408
+ latent_model_input = (
1409
+ torch.cat([latents] * 2)
1410
+ if self.do_classifier_free_guidance
1411
+ else latents
1412
+ )
1413
+ latent_model_input = self.scheduler.scale_model_input(
1414
+ latent_model_input, t
1415
+ )
1416
+
1417
+ # controlnet(s) inference
1418
+ if guess_mode and self.do_classifier_free_guidance:
1419
+ # Infer ControlNet only for the conditional batch.
1420
+ control_model_input = latents
1421
+ control_model_input = self.scheduler.scale_model_input(
1422
+ control_model_input, t
1423
+ )
1424
+ controlnet_prompt_embeds = prompt_embeds.chunk(2)[1]
1425
+ else:
1426
+ control_model_input = latent_model_input
1427
+ controlnet_prompt_embeds = prompt_embeds
1428
+
1429
+ if isinstance(controlnet_keep[i], list):
1430
+ cond_scale = [
1431
+ c * s
1432
+ for c, s in zip(
1433
+ controlnet_conditioning_scale, controlnet_keep[i]
1434
+ )
1435
+ ]
1436
+ else:
1437
+ controlnet_cond_scale = controlnet_conditioning_scale
1438
+ if isinstance(controlnet_cond_scale, list):
1439
+ controlnet_cond_scale = controlnet_cond_scale[0]
1440
+ cond_scale = controlnet_cond_scale * controlnet_keep[i]
1441
+
1442
+ down_block_res_samples, mid_block_res_sample = self.controlnet(
1443
+ control_model_input,
1444
+ t,
1445
+ encoder_hidden_states=controlnet_prompt_embeds,
1446
+ controlnet_cond=image,
1447
+ conditioning_scale=cond_scale,
1448
+ guess_mode=guess_mode,
1449
+ return_dict=False,
1450
+ )
1451
+
1452
+ if guess_mode and self.do_classifier_free_guidance:
1453
+ # Inferred ControlNet only for the conditional batch.
1454
+ # To apply the output of ControlNet to both the unconditional and conditional batches,
1455
+ # add 0 to the unconditional batch to keep it unchanged.
1456
+ down_block_res_samples = [
1457
+ torch.cat([torch.zeros_like(d), d])
1458
+ for d in down_block_res_samples
1459
+ ]
1460
+ mid_block_res_sample = torch.cat(
1461
+ [torch.zeros_like(mid_block_res_sample), mid_block_res_sample]
1462
+ )
1463
+
1464
+ # predict the noise residual
1465
+ noise_pred = self.unet(
1466
+ latent_model_input,
1467
+ t,
1468
+ encoder_hidden_states=prompt_embeds,
1469
+ timestep_cond=timestep_cond,
1470
+ cross_attention_kwargs=self.cross_attention_kwargs,
1471
+ down_block_additional_residuals=down_block_res_samples,
1472
+ mid_block_additional_residual=mid_block_res_sample,
1473
+ added_cond_kwargs=added_cond_kwargs,
1474
+ return_dict=False,
1475
+ )[0]
1476
+
1477
+ # perform guidance
1478
+ if self.do_classifier_free_guidance:
1479
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
1480
+ noise_pred = noise_pred_uncond + self.guidance_scale * (
1481
+ noise_pred_text - noise_pred_uncond
1482
+ )
1483
+
1484
+ # compute the previous noisy sample x_t -> x_t-1
1485
+ latents = self.scheduler.step(
1486
+ noise_pred, t, latents, **extra_step_kwargs, return_dict=False
1487
+ )[0]
1488
+
1489
+ if callback_on_step_end is not None:
1490
+ callback_kwargs = {}
1491
+ for k in callback_on_step_end_tensor_inputs:
1492
+ callback_kwargs[k] = locals()[k]
1493
+ callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
1494
+
1495
+ latents = callback_outputs.pop("latents", latents)
1496
+ prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
1497
+ negative_prompt_embeds = callback_outputs.pop(
1498
+ "negative_prompt_embeds", negative_prompt_embeds
1499
+ )
1500
+
1501
+ # call the callback, if provided
1502
+ if i == len(timesteps) - 1 or (
1503
+ (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
1504
+ ):
1505
+ progress_bar.update()
1506
+ if callback is not None and i % callback_steps == 0:
1507
+ step_idx = i // getattr(self.scheduler, "order", 1)
1508
+ callback(step_idx, t, latents)
1509
+
1510
+ # If we do sequential model offloading, let's offload unet and controlnet
1511
+ # manually for max memory savings
1512
+ if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
1513
+ self.unet.to("cpu")
1514
+ self.controlnet.to("cpu")
1515
+ torch.cuda.empty_cache()
1516
+
1517
+ if not output_type == "latent":
1518
+ image = self.vae.decode(
1519
+ latents / self.vae.config.scaling_factor,
1520
+ return_dict=False,
1521
+ generator=generator,
1522
+ )[0]
1523
+ image, has_nsfw_concept = self.run_safety_checker(
1524
+ image, device, prompt_embeds.dtype
1525
+ )
1526
+ else:
1527
+ image = latents
1528
+ has_nsfw_concept = None
1529
+
1530
+ if has_nsfw_concept is None:
1531
+ do_denormalize = [True] * image.shape[0]
1532
+ else:
1533
+ do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
1534
+
1535
+ image = self.image_processor.postprocess(
1536
+ image, output_type=output_type, do_denormalize=do_denormalize
1537
+ )
1538
+
1539
+ # Offload all models
1540
+ self.maybe_free_model_hooks()
1541
+
1542
+ if not return_dict:
1543
+ return (image, has_nsfw_concept)
1544
+
1545
+ return StableDiffusionPipelineOutput(
1546
+ images=image, nsfw_content_detected=has_nsfw_concept
1547
+ )