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