diffusers 0.30.3__py3-none-any.whl → 0.31.0__py3-none-any.whl

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