diffusers 0.23.1__py3-none-any.whl → 0.25.0__py3-none-any.whl

Sign up to get free protection for your applications and to get access to all the features.
Files changed (238) hide show
  1. diffusers/__init__.py +26 -2
  2. diffusers/commands/fp16_safetensors.py +10 -11
  3. diffusers/configuration_utils.py +13 -8
  4. diffusers/dependency_versions_check.py +0 -1
  5. diffusers/dependency_versions_table.py +5 -5
  6. diffusers/experimental/rl/value_guided_sampling.py +1 -1
  7. diffusers/image_processor.py +463 -51
  8. diffusers/loaders/__init__.py +82 -0
  9. diffusers/loaders/ip_adapter.py +159 -0
  10. diffusers/loaders/lora.py +1553 -0
  11. diffusers/loaders/lora_conversion_utils.py +284 -0
  12. diffusers/loaders/single_file.py +637 -0
  13. diffusers/loaders/textual_inversion.py +455 -0
  14. diffusers/loaders/unet.py +828 -0
  15. diffusers/loaders/utils.py +59 -0
  16. diffusers/models/__init__.py +26 -9
  17. diffusers/models/activations.py +9 -6
  18. diffusers/models/attention.py +301 -29
  19. diffusers/models/attention_flax.py +9 -1
  20. diffusers/models/attention_processor.py +378 -6
  21. diffusers/models/autoencoders/__init__.py +5 -0
  22. diffusers/models/{autoencoder_asym_kl.py → autoencoders/autoencoder_asym_kl.py} +17 -12
  23. diffusers/models/{autoencoder_kl.py → autoencoders/autoencoder_kl.py} +47 -23
  24. diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py +402 -0
  25. diffusers/models/{autoencoder_tiny.py → autoencoders/autoencoder_tiny.py} +24 -28
  26. diffusers/models/{consistency_decoder_vae.py → autoencoders/consistency_decoder_vae.py} +51 -44
  27. diffusers/models/{vae.py → autoencoders/vae.py} +71 -17
  28. diffusers/models/controlnet.py +59 -39
  29. diffusers/models/controlnet_flax.py +19 -18
  30. diffusers/models/downsampling.py +338 -0
  31. diffusers/models/embeddings.py +112 -29
  32. diffusers/models/embeddings_flax.py +2 -0
  33. diffusers/models/lora.py +131 -1
  34. diffusers/models/modeling_flax_utils.py +14 -8
  35. diffusers/models/modeling_outputs.py +17 -0
  36. diffusers/models/modeling_utils.py +37 -29
  37. diffusers/models/normalization.py +110 -4
  38. diffusers/models/resnet.py +299 -652
  39. diffusers/models/transformer_2d.py +22 -5
  40. diffusers/models/transformer_temporal.py +183 -1
  41. diffusers/models/unet_2d_blocks_flax.py +5 -0
  42. diffusers/models/unet_2d_condition.py +46 -0
  43. diffusers/models/unet_2d_condition_flax.py +13 -13
  44. diffusers/models/unet_3d_blocks.py +957 -173
  45. diffusers/models/unet_3d_condition.py +16 -8
  46. diffusers/models/unet_kandinsky3.py +535 -0
  47. diffusers/models/unet_motion_model.py +48 -33
  48. diffusers/models/unet_spatio_temporal_condition.py +489 -0
  49. diffusers/models/upsampling.py +454 -0
  50. diffusers/models/uvit_2d.py +471 -0
  51. diffusers/models/vae_flax.py +7 -0
  52. diffusers/models/vq_model.py +12 -3
  53. diffusers/optimization.py +16 -9
  54. diffusers/pipelines/__init__.py +137 -76
  55. diffusers/pipelines/amused/__init__.py +62 -0
  56. diffusers/pipelines/amused/pipeline_amused.py +328 -0
  57. diffusers/pipelines/amused/pipeline_amused_img2img.py +347 -0
  58. diffusers/pipelines/amused/pipeline_amused_inpaint.py +378 -0
  59. diffusers/pipelines/animatediff/pipeline_animatediff.py +66 -8
  60. diffusers/pipelines/audioldm/pipeline_audioldm.py +1 -0
  61. diffusers/pipelines/auto_pipeline.py +23 -13
  62. diffusers/pipelines/consistency_models/pipeline_consistency_models.py +1 -0
  63. diffusers/pipelines/controlnet/pipeline_controlnet.py +238 -35
  64. diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py +148 -37
  65. diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py +155 -41
  66. diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py +123 -43
  67. diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py +216 -39
  68. diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py +106 -34
  69. diffusers/pipelines/dance_diffusion/pipeline_dance_diffusion.py +1 -0
  70. diffusers/pipelines/ddim/pipeline_ddim.py +1 -0
  71. diffusers/pipelines/ddpm/pipeline_ddpm.py +1 -0
  72. diffusers/pipelines/deepfloyd_if/pipeline_if.py +13 -1
  73. diffusers/pipelines/deepfloyd_if/pipeline_if_img2img.py +13 -1
  74. diffusers/pipelines/deepfloyd_if/pipeline_if_img2img_superresolution.py +13 -1
  75. diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting.py +13 -1
  76. diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting_superresolution.py +13 -1
  77. diffusers/pipelines/deepfloyd_if/pipeline_if_superresolution.py +13 -1
  78. diffusers/pipelines/deprecated/__init__.py +153 -0
  79. diffusers/pipelines/{alt_diffusion → deprecated/alt_diffusion}/__init__.py +3 -3
  80. diffusers/pipelines/{alt_diffusion → deprecated/alt_diffusion}/pipeline_alt_diffusion.py +177 -34
  81. diffusers/pipelines/{alt_diffusion → deprecated/alt_diffusion}/pipeline_alt_diffusion_img2img.py +182 -37
  82. diffusers/pipelines/{alt_diffusion → deprecated/alt_diffusion}/pipeline_output.py +1 -1
  83. diffusers/pipelines/{audio_diffusion → deprecated/audio_diffusion}/__init__.py +1 -1
  84. diffusers/pipelines/{audio_diffusion → deprecated/audio_diffusion}/mel.py +2 -2
  85. diffusers/pipelines/{audio_diffusion → deprecated/audio_diffusion}/pipeline_audio_diffusion.py +4 -4
  86. diffusers/pipelines/{latent_diffusion_uncond → deprecated/latent_diffusion_uncond}/__init__.py +1 -1
  87. diffusers/pipelines/{latent_diffusion_uncond → deprecated/latent_diffusion_uncond}/pipeline_latent_diffusion_uncond.py +4 -4
  88. diffusers/pipelines/{pndm → deprecated/pndm}/__init__.py +1 -1
  89. diffusers/pipelines/{pndm → deprecated/pndm}/pipeline_pndm.py +4 -4
  90. diffusers/pipelines/{repaint → deprecated/repaint}/__init__.py +1 -1
  91. diffusers/pipelines/{repaint → deprecated/repaint}/pipeline_repaint.py +5 -5
  92. diffusers/pipelines/{score_sde_ve → deprecated/score_sde_ve}/__init__.py +1 -1
  93. diffusers/pipelines/{score_sde_ve → deprecated/score_sde_ve}/pipeline_score_sde_ve.py +5 -4
  94. diffusers/pipelines/{spectrogram_diffusion → deprecated/spectrogram_diffusion}/__init__.py +6 -6
  95. diffusers/pipelines/{spectrogram_diffusion/continous_encoder.py → deprecated/spectrogram_diffusion/continuous_encoder.py} +2 -2
  96. diffusers/pipelines/{spectrogram_diffusion → deprecated/spectrogram_diffusion}/midi_utils.py +1 -1
  97. diffusers/pipelines/{spectrogram_diffusion → deprecated/spectrogram_diffusion}/notes_encoder.py +2 -2
  98. diffusers/pipelines/{spectrogram_diffusion → deprecated/spectrogram_diffusion}/pipeline_spectrogram_diffusion.py +8 -7
  99. diffusers/pipelines/deprecated/stable_diffusion_variants/__init__.py +55 -0
  100. diffusers/pipelines/{stable_diffusion → deprecated/stable_diffusion_variants}/pipeline_cycle_diffusion.py +34 -13
  101. diffusers/pipelines/{stable_diffusion → deprecated/stable_diffusion_variants}/pipeline_onnx_stable_diffusion_inpaint_legacy.py +7 -6
  102. diffusers/pipelines/{stable_diffusion → deprecated/stable_diffusion_variants}/pipeline_stable_diffusion_inpaint_legacy.py +12 -11
  103. diffusers/pipelines/{stable_diffusion → deprecated/stable_diffusion_variants}/pipeline_stable_diffusion_model_editing.py +17 -11
  104. diffusers/pipelines/{stable_diffusion → deprecated/stable_diffusion_variants}/pipeline_stable_diffusion_paradigms.py +11 -10
  105. diffusers/pipelines/{stable_diffusion → deprecated/stable_diffusion_variants}/pipeline_stable_diffusion_pix2pix_zero.py +14 -13
  106. diffusers/pipelines/{stochastic_karras_ve → deprecated/stochastic_karras_ve}/__init__.py +1 -1
  107. diffusers/pipelines/{stochastic_karras_ve → deprecated/stochastic_karras_ve}/pipeline_stochastic_karras_ve.py +4 -4
  108. diffusers/pipelines/{versatile_diffusion → deprecated/versatile_diffusion}/__init__.py +3 -3
  109. diffusers/pipelines/{versatile_diffusion → deprecated/versatile_diffusion}/modeling_text_unet.py +83 -51
  110. diffusers/pipelines/{versatile_diffusion → deprecated/versatile_diffusion}/pipeline_versatile_diffusion.py +4 -4
  111. diffusers/pipelines/{versatile_diffusion → deprecated/versatile_diffusion}/pipeline_versatile_diffusion_dual_guided.py +7 -6
  112. diffusers/pipelines/{versatile_diffusion → deprecated/versatile_diffusion}/pipeline_versatile_diffusion_image_variation.py +7 -6
  113. diffusers/pipelines/{versatile_diffusion → deprecated/versatile_diffusion}/pipeline_versatile_diffusion_text_to_image.py +7 -6
  114. diffusers/pipelines/{vq_diffusion → deprecated/vq_diffusion}/__init__.py +3 -3
  115. diffusers/pipelines/{vq_diffusion → deprecated/vq_diffusion}/pipeline_vq_diffusion.py +5 -5
  116. diffusers/pipelines/dit/pipeline_dit.py +1 -0
  117. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py +1 -1
  118. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py +3 -3
  119. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py +1 -1
  120. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpainting.py +1 -1
  121. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py +1 -1
  122. diffusers/pipelines/kandinsky3/__init__.py +49 -0
  123. diffusers/pipelines/kandinsky3/convert_kandinsky3_unet.py +98 -0
  124. diffusers/pipelines/kandinsky3/pipeline_kandinsky3.py +589 -0
  125. diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py +654 -0
  126. diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_img2img.py +111 -11
  127. diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_text2img.py +102 -9
  128. diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py +1 -0
  129. diffusers/pipelines/musicldm/pipeline_musicldm.py +1 -1
  130. diffusers/pipelines/onnx_utils.py +8 -5
  131. diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py +7 -2
  132. diffusers/pipelines/pipeline_flax_utils.py +11 -8
  133. diffusers/pipelines/pipeline_utils.py +63 -42
  134. diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py +247 -38
  135. diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py +3 -3
  136. diffusers/pipelines/stable_diffusion/__init__.py +37 -65
  137. diffusers/pipelines/stable_diffusion/convert_from_ckpt.py +75 -78
  138. diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py +2 -2
  139. diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_img2img.py +2 -4
  140. diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_inpaint.py +1 -0
  141. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +174 -11
  142. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_depth2img.py +8 -3
  143. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py +1 -0
  144. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py +178 -11
  145. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py +224 -13
  146. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py +74 -20
  147. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_latent_upscale.py +4 -0
  148. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py +7 -0
  149. diffusers/pipelines/stable_diffusion/pipeline_stable_unclip.py +5 -0
  150. diffusers/pipelines/stable_diffusion/pipeline_stable_unclip_img2img.py +5 -0
  151. diffusers/pipelines/stable_diffusion_attend_and_excite/__init__.py +48 -0
  152. diffusers/pipelines/{stable_diffusion → stable_diffusion_attend_and_excite}/pipeline_stable_diffusion_attend_and_excite.py +6 -2
  153. diffusers/pipelines/stable_diffusion_diffedit/__init__.py +48 -0
  154. diffusers/pipelines/{stable_diffusion → stable_diffusion_diffedit}/pipeline_stable_diffusion_diffedit.py +3 -3
  155. diffusers/pipelines/stable_diffusion_gligen/__init__.py +50 -0
  156. diffusers/pipelines/{stable_diffusion → stable_diffusion_gligen}/pipeline_stable_diffusion_gligen.py +3 -2
  157. diffusers/pipelines/{stable_diffusion → stable_diffusion_gligen}/pipeline_stable_diffusion_gligen_text_image.py +4 -3
  158. diffusers/pipelines/stable_diffusion_k_diffusion/__init__.py +60 -0
  159. diffusers/pipelines/{stable_diffusion → stable_diffusion_k_diffusion}/pipeline_stable_diffusion_k_diffusion.py +7 -1
  160. diffusers/pipelines/stable_diffusion_ldm3d/__init__.py +48 -0
  161. diffusers/pipelines/{stable_diffusion → stable_diffusion_ldm3d}/pipeline_stable_diffusion_ldm3d.py +51 -7
  162. diffusers/pipelines/stable_diffusion_panorama/__init__.py +48 -0
  163. diffusers/pipelines/{stable_diffusion → stable_diffusion_panorama}/pipeline_stable_diffusion_panorama.py +57 -8
  164. diffusers/pipelines/stable_diffusion_safe/pipeline_stable_diffusion_safe.py +58 -6
  165. diffusers/pipelines/stable_diffusion_sag/__init__.py +48 -0
  166. diffusers/pipelines/{stable_diffusion → stable_diffusion_sag}/pipeline_stable_diffusion_sag.py +68 -10
  167. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py +194 -17
  168. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py +205 -16
  169. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py +206 -17
  170. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_instruct_pix2pix.py +23 -17
  171. diffusers/pipelines/stable_video_diffusion/__init__.py +58 -0
  172. diffusers/pipelines/stable_video_diffusion/pipeline_stable_video_diffusion.py +652 -0
  173. diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_adapter.py +108 -12
  174. diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_xl_adapter.py +115 -14
  175. diffusers/pipelines/text_to_video_synthesis/__init__.py +2 -0
  176. diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth.py +6 -0
  177. diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth_img2img.py +23 -3
  178. diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero.py +334 -10
  179. diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero_sdxl.py +1331 -0
  180. diffusers/pipelines/unclip/pipeline_unclip.py +2 -1
  181. diffusers/pipelines/unclip/pipeline_unclip_image_variation.py +1 -0
  182. diffusers/pipelines/wuerstchen/modeling_paella_vq_model.py +1 -1
  183. diffusers/pipelines/wuerstchen/modeling_wuerstchen_common.py +14 -4
  184. diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py +9 -5
  185. diffusers/pipelines/wuerstchen/pipeline_wuerstchen.py +1 -1
  186. diffusers/pipelines/wuerstchen/pipeline_wuerstchen_combined.py +2 -2
  187. diffusers/pipelines/wuerstchen/pipeline_wuerstchen_prior.py +5 -1
  188. diffusers/schedulers/__init__.py +4 -4
  189. diffusers/schedulers/deprecated/__init__.py +50 -0
  190. diffusers/schedulers/{scheduling_karras_ve.py → deprecated/scheduling_karras_ve.py} +4 -4
  191. diffusers/schedulers/{scheduling_sde_vp.py → deprecated/scheduling_sde_vp.py} +4 -6
  192. diffusers/schedulers/scheduling_amused.py +162 -0
  193. diffusers/schedulers/scheduling_consistency_models.py +2 -0
  194. diffusers/schedulers/scheduling_ddim.py +1 -3
  195. diffusers/schedulers/scheduling_ddim_inverse.py +2 -7
  196. diffusers/schedulers/scheduling_ddim_parallel.py +1 -3
  197. diffusers/schedulers/scheduling_ddpm.py +47 -3
  198. diffusers/schedulers/scheduling_ddpm_parallel.py +47 -3
  199. diffusers/schedulers/scheduling_deis_multistep.py +28 -6
  200. diffusers/schedulers/scheduling_dpmsolver_multistep.py +28 -6
  201. diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py +28 -6
  202. diffusers/schedulers/scheduling_dpmsolver_sde.py +3 -3
  203. diffusers/schedulers/scheduling_dpmsolver_singlestep.py +28 -6
  204. diffusers/schedulers/scheduling_euler_ancestral_discrete.py +59 -3
  205. diffusers/schedulers/scheduling_euler_discrete.py +102 -16
  206. diffusers/schedulers/scheduling_heun_discrete.py +17 -5
  207. diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py +17 -5
  208. diffusers/schedulers/scheduling_k_dpm_2_discrete.py +17 -5
  209. diffusers/schedulers/scheduling_lcm.py +123 -29
  210. diffusers/schedulers/scheduling_lms_discrete.py +3 -3
  211. diffusers/schedulers/scheduling_pndm.py +1 -3
  212. diffusers/schedulers/scheduling_repaint.py +1 -3
  213. diffusers/schedulers/scheduling_unipc_multistep.py +28 -6
  214. diffusers/schedulers/scheduling_utils.py +3 -1
  215. diffusers/schedulers/scheduling_utils_flax.py +3 -1
  216. diffusers/training_utils.py +1 -1
  217. diffusers/utils/__init__.py +1 -2
  218. diffusers/utils/constants.py +10 -12
  219. diffusers/utils/dummy_pt_objects.py +75 -0
  220. diffusers/utils/dummy_torch_and_transformers_objects.py +105 -0
  221. diffusers/utils/dynamic_modules_utils.py +18 -22
  222. diffusers/utils/export_utils.py +8 -3
  223. diffusers/utils/hub_utils.py +24 -36
  224. diffusers/utils/logging.py +11 -11
  225. diffusers/utils/outputs.py +5 -5
  226. diffusers/utils/peft_utils.py +88 -44
  227. diffusers/utils/state_dict_utils.py +8 -0
  228. diffusers/utils/testing_utils.py +199 -1
  229. diffusers/utils/torch_utils.py +4 -4
  230. {diffusers-0.23.1.dist-info → diffusers-0.25.0.dist-info}/METADATA +86 -69
  231. diffusers-0.25.0.dist-info/RECORD +360 -0
  232. {diffusers-0.23.1.dist-info → diffusers-0.25.0.dist-info}/WHEEL +1 -1
  233. {diffusers-0.23.1.dist-info → diffusers-0.25.0.dist-info}/entry_points.txt +0 -1
  234. diffusers/loaders.py +0 -3336
  235. diffusers-0.23.1.dist-info/RECORD +0 -323
  236. /diffusers/pipelines/{alt_diffusion → deprecated/alt_diffusion}/modeling_roberta_series.py +0 -0
  237. {diffusers-0.23.1.dist-info → diffusers-0.25.0.dist-info}/LICENSE +0 -0
  238. {diffusers-0.23.1.dist-info → diffusers-0.25.0.dist-info}/top_level.txt +0 -0
@@ -19,6 +19,7 @@ import urllib.parse as ul
19
19
  from typing import Callable, List, Optional, Tuple, Union
20
20
 
21
21
  import torch
22
+ import torch.nn.functional as F
22
23
  from transformers import T5EncoderModel, T5Tokenizer
23
24
 
24
25
  from ...image_processor import VaeImageProcessor
@@ -26,6 +27,7 @@ from ...models import AutoencoderKL, Transformer2DModel
26
27
  from ...schedulers import DPMSolverMultistepScheduler
27
28
  from ...utils import (
28
29
  BACKENDS_MAPPING,
30
+ deprecate,
29
31
  is_bs4_available,
30
32
  is_ftfy_available,
31
33
  logging,
@@ -43,7 +45,6 @@ if is_bs4_available():
43
45
  if is_ftfy_available():
44
46
  import ftfy
45
47
 
46
-
47
48
  EXAMPLE_DOC_STRING = """
48
49
  Examples:
49
50
  ```py
@@ -60,6 +61,123 @@ EXAMPLE_DOC_STRING = """
60
61
  ```
61
62
  """
62
63
 
64
+ ASPECT_RATIO_1024_BIN = {
65
+ "0.25": [512.0, 2048.0],
66
+ "0.28": [512.0, 1856.0],
67
+ "0.32": [576.0, 1792.0],
68
+ "0.33": [576.0, 1728.0],
69
+ "0.35": [576.0, 1664.0],
70
+ "0.4": [640.0, 1600.0],
71
+ "0.42": [640.0, 1536.0],
72
+ "0.48": [704.0, 1472.0],
73
+ "0.5": [704.0, 1408.0],
74
+ "0.52": [704.0, 1344.0],
75
+ "0.57": [768.0, 1344.0],
76
+ "0.6": [768.0, 1280.0],
77
+ "0.68": [832.0, 1216.0],
78
+ "0.72": [832.0, 1152.0],
79
+ "0.78": [896.0, 1152.0],
80
+ "0.82": [896.0, 1088.0],
81
+ "0.88": [960.0, 1088.0],
82
+ "0.94": [960.0, 1024.0],
83
+ "1.0": [1024.0, 1024.0],
84
+ "1.07": [1024.0, 960.0],
85
+ "1.13": [1088.0, 960.0],
86
+ "1.21": [1088.0, 896.0],
87
+ "1.29": [1152.0, 896.0],
88
+ "1.38": [1152.0, 832.0],
89
+ "1.46": [1216.0, 832.0],
90
+ "1.67": [1280.0, 768.0],
91
+ "1.75": [1344.0, 768.0],
92
+ "2.0": [1408.0, 704.0],
93
+ "2.09": [1472.0, 704.0],
94
+ "2.4": [1536.0, 640.0],
95
+ "2.5": [1600.0, 640.0],
96
+ "3.0": [1728.0, 576.0],
97
+ "4.0": [2048.0, 512.0],
98
+ }
99
+
100
+ ASPECT_RATIO_512_BIN = {
101
+ "0.25": [256.0, 1024.0],
102
+ "0.28": [256.0, 928.0],
103
+ "0.32": [288.0, 896.0],
104
+ "0.33": [288.0, 864.0],
105
+ "0.35": [288.0, 832.0],
106
+ "0.4": [320.0, 800.0],
107
+ "0.42": [320.0, 768.0],
108
+ "0.48": [352.0, 736.0],
109
+ "0.5": [352.0, 704.0],
110
+ "0.52": [352.0, 672.0],
111
+ "0.57": [384.0, 672.0],
112
+ "0.6": [384.0, 640.0],
113
+ "0.68": [416.0, 608.0],
114
+ "0.72": [416.0, 576.0],
115
+ "0.78": [448.0, 576.0],
116
+ "0.82": [448.0, 544.0],
117
+ "0.88": [480.0, 544.0],
118
+ "0.94": [480.0, 512.0],
119
+ "1.0": [512.0, 512.0],
120
+ "1.07": [512.0, 480.0],
121
+ "1.13": [544.0, 480.0],
122
+ "1.21": [544.0, 448.0],
123
+ "1.29": [576.0, 448.0],
124
+ "1.38": [576.0, 416.0],
125
+ "1.46": [608.0, 416.0],
126
+ "1.67": [640.0, 384.0],
127
+ "1.75": [672.0, 384.0],
128
+ "2.0": [704.0, 352.0],
129
+ "2.09": [736.0, 352.0],
130
+ "2.4": [768.0, 320.0],
131
+ "2.5": [800.0, 320.0],
132
+ "3.0": [864.0, 288.0],
133
+ "4.0": [1024.0, 256.0],
134
+ }
135
+
136
+
137
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
138
+ def retrieve_timesteps(
139
+ scheduler,
140
+ num_inference_steps: Optional[int] = None,
141
+ device: Optional[Union[str, torch.device]] = None,
142
+ timesteps: Optional[List[int]] = None,
143
+ **kwargs,
144
+ ):
145
+ """
146
+ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
147
+ custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
148
+
149
+ Args:
150
+ scheduler (`SchedulerMixin`):
151
+ The scheduler to get timesteps from.
152
+ num_inference_steps (`int`):
153
+ The number of diffusion steps used when generating samples with a pre-trained model. If used,
154
+ `timesteps` must be `None`.
155
+ device (`str` or `torch.device`, *optional*):
156
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
157
+ timesteps (`List[int]`, *optional*):
158
+ Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
159
+ timestep spacing strategy of the scheduler is used. If `timesteps` is passed, `num_inference_steps`
160
+ must be `None`.
161
+
162
+ Returns:
163
+ `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
164
+ second element is the number of inference steps.
165
+ """
166
+ if timesteps is not None:
167
+ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
168
+ if not accepts_timesteps:
169
+ raise ValueError(
170
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
171
+ f" timestep schedules. Please check whether you are using the correct scheduler."
172
+ )
173
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
174
+ timesteps = scheduler.timesteps
175
+ num_inference_steps = len(timesteps)
176
+ else:
177
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
178
+ timesteps = scheduler.timesteps
179
+ return timesteps, num_inference_steps
180
+
63
181
 
64
182
  class PixArtAlphaPipeline(DiffusionPipeline):
65
183
  r"""
@@ -83,8 +201,21 @@ class PixArtAlphaPipeline(DiffusionPipeline):
83
201
  scheduler ([`SchedulerMixin`]):
84
202
  A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
85
203
  """
204
+
86
205
  bad_punct_regex = re.compile(
87
- r"[" + "#®•©™&@·º½¾¿¡§~" + "\)" + "\(" + "\]" + "\[" + "\}" + "\{" + "\|" + "\\" + "\/" + "\*" + r"]{1,}"
206
+ r"["
207
+ + "#®•©™&@·º½¾¿¡§~"
208
+ + r"\)"
209
+ + r"\("
210
+ + r"\]"
211
+ + r"\["
212
+ + r"\}"
213
+ + r"\{"
214
+ + r"\|"
215
+ + "\\"
216
+ + r"\/"
217
+ + r"\*"
218
+ + r"]{1,}"
88
219
  ) # noqa
89
220
 
90
221
  _optional_components = ["tokenizer", "text_encoder"]
@@ -126,8 +257,10 @@ class PixArtAlphaPipeline(DiffusionPipeline):
126
257
  device: Optional[torch.device] = None,
127
258
  prompt_embeds: Optional[torch.FloatTensor] = None,
128
259
  negative_prompt_embeds: Optional[torch.FloatTensor] = None,
260
+ prompt_attention_mask: Optional[torch.FloatTensor] = None,
261
+ negative_prompt_attention_mask: Optional[torch.FloatTensor] = None,
129
262
  clean_caption: bool = False,
130
- mask_feature: bool = True,
263
+ **kwargs,
131
264
  ):
132
265
  r"""
133
266
  Encodes the prompt into text encoder hidden states.
@@ -153,10 +286,11 @@ class PixArtAlphaPipeline(DiffusionPipeline):
153
286
  string.
154
287
  clean_caption (bool, defaults to `False`):
155
288
  If `True`, the function will preprocess and clean the provided caption before encoding.
156
- mask_feature: (bool, defaults to `True`):
157
- If `True`, the function will mask the text embeddings.
158
289
  """
159
- embeds_initially_provided = prompt_embeds is not None and negative_prompt_embeds is not None
290
+
291
+ if "mask_feature" in kwargs:
292
+ deprecation_message = "The use of `mask_feature` is deprecated. It is no longer used in any computation and that doesn't affect the end results. It will be removed in a future version."
293
+ deprecate("mask_feature", "1.0.0", deprecation_message, standard_warn=False)
160
294
 
161
295
  if device is None:
162
296
  device = self._execution_device
@@ -193,13 +327,11 @@ class PixArtAlphaPipeline(DiffusionPipeline):
193
327
  f" {max_length} tokens: {removed_text}"
194
328
  )
195
329
 
196
- attention_mask = text_inputs.attention_mask.to(device)
197
- prompt_embeds_attention_mask = attention_mask
330
+ prompt_attention_mask = text_inputs.attention_mask
331
+ prompt_attention_mask = prompt_attention_mask.to(device)
198
332
 
199
- prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=attention_mask)
333
+ prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask)
200
334
  prompt_embeds = prompt_embeds[0]
201
- else:
202
- prompt_embeds_attention_mask = torch.ones_like(prompt_embeds)
203
335
 
204
336
  if self.text_encoder is not None:
205
337
  dtype = self.text_encoder.dtype
@@ -214,8 +346,8 @@ class PixArtAlphaPipeline(DiffusionPipeline):
214
346
  # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
215
347
  prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
216
348
  prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
217
- prompt_embeds_attention_mask = prompt_embeds_attention_mask.view(bs_embed, -1)
218
- prompt_embeds_attention_mask = prompt_embeds_attention_mask.repeat(num_images_per_prompt, 1)
349
+ prompt_attention_mask = prompt_attention_mask.view(bs_embed, -1)
350
+ prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1)
219
351
 
220
352
  # get unconditional embeddings for classifier free guidance
221
353
  if do_classifier_free_guidance and negative_prompt_embeds is None:
@@ -231,11 +363,11 @@ class PixArtAlphaPipeline(DiffusionPipeline):
231
363
  add_special_tokens=True,
232
364
  return_tensors="pt",
233
365
  )
234
- attention_mask = uncond_input.attention_mask.to(device)
366
+ negative_prompt_attention_mask = uncond_input.attention_mask
367
+ negative_prompt_attention_mask = negative_prompt_attention_mask.to(device)
235
368
 
236
369
  negative_prompt_embeds = self.text_encoder(
237
- uncond_input.input_ids.to(device),
238
- attention_mask=attention_mask,
370
+ uncond_input.input_ids.to(device), attention_mask=negative_prompt_attention_mask
239
371
  )
240
372
  negative_prompt_embeds = negative_prompt_embeds[0]
241
373
 
@@ -248,23 +380,13 @@ class PixArtAlphaPipeline(DiffusionPipeline):
248
380
  negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
249
381
  negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
250
382
 
251
- # For classifier free guidance, we need to do two forward passes.
252
- # Here we concatenate the unconditional and text embeddings into a single batch
253
- # to avoid doing two forward passes
383
+ negative_prompt_attention_mask = negative_prompt_attention_mask.view(bs_embed, -1)
384
+ negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 1)
254
385
  else:
255
386
  negative_prompt_embeds = None
387
+ negative_prompt_attention_mask = None
256
388
 
257
- # Perform additional masking.
258
- if mask_feature and not embeds_initially_provided:
259
- prompt_embeds = prompt_embeds.unsqueeze(1)
260
- masked_prompt_embeds, keep_indices = self.mask_text_embeddings(prompt_embeds, prompt_embeds_attention_mask)
261
- masked_prompt_embeds = masked_prompt_embeds.squeeze(1)
262
- masked_negative_prompt_embeds = (
263
- negative_prompt_embeds[:, :keep_indices, :] if negative_prompt_embeds is not None else None
264
- )
265
- return masked_prompt_embeds, masked_negative_prompt_embeds
266
-
267
- return prompt_embeds, negative_prompt_embeds
389
+ return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask
268
390
 
269
391
  # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
270
392
  def prepare_extra_step_kwargs(self, generator, eta):
@@ -293,6 +415,8 @@ class PixArtAlphaPipeline(DiffusionPipeline):
293
415
  callback_steps,
294
416
  prompt_embeds=None,
295
417
  negative_prompt_embeds=None,
418
+ prompt_attention_mask=None,
419
+ negative_prompt_attention_mask=None,
296
420
  ):
297
421
  if height % 8 != 0 or width % 8 != 0:
298
422
  raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
@@ -329,6 +453,12 @@ class PixArtAlphaPipeline(DiffusionPipeline):
329
453
  f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
330
454
  )
331
455
 
456
+ if prompt_embeds is not None and prompt_attention_mask is None:
457
+ raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.")
458
+
459
+ if negative_prompt_embeds is not None and negative_prompt_attention_mask is None:
460
+ raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.")
461
+
332
462
  if prompt_embeds is not None and negative_prompt_embeds is not None:
333
463
  if prompt_embeds.shape != negative_prompt_embeds.shape:
334
464
  raise ValueError(
@@ -336,6 +466,12 @@ class PixArtAlphaPipeline(DiffusionPipeline):
336
466
  f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
337
467
  f" {negative_prompt_embeds.shape}."
338
468
  )
469
+ if prompt_attention_mask.shape != negative_prompt_attention_mask.shape:
470
+ raise ValueError(
471
+ "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but"
472
+ f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`"
473
+ f" {negative_prompt_attention_mask.shape}."
474
+ )
339
475
 
340
476
  # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing
341
477
  def _text_preprocessing(self, text, clean_caption=False):
@@ -495,6 +631,38 @@ class PixArtAlphaPipeline(DiffusionPipeline):
495
631
  latents = latents * self.scheduler.init_noise_sigma
496
632
  return latents
497
633
 
634
+ @staticmethod
635
+ def classify_height_width_bin(height: int, width: int, ratios: dict) -> Tuple[int, int]:
636
+ """Returns binned height and width."""
637
+ ar = float(height / width)
638
+ closest_ratio = min(ratios.keys(), key=lambda ratio: abs(float(ratio) - ar))
639
+ default_hw = ratios[closest_ratio]
640
+ return int(default_hw[0]), int(default_hw[1])
641
+
642
+ @staticmethod
643
+ def resize_and_crop_tensor(samples: torch.Tensor, new_width: int, new_height: int) -> torch.Tensor:
644
+ orig_height, orig_width = samples.shape[2], samples.shape[3]
645
+
646
+ # Check if resizing is needed
647
+ if orig_height != new_height or orig_width != new_width:
648
+ ratio = max(new_height / orig_height, new_width / orig_width)
649
+ resized_width = int(orig_width * ratio)
650
+ resized_height = int(orig_height * ratio)
651
+
652
+ # Resize
653
+ samples = F.interpolate(
654
+ samples, size=(resized_height, resized_width), mode="bilinear", align_corners=False
655
+ )
656
+
657
+ # Center Crop
658
+ start_x = (resized_width - new_width) // 2
659
+ end_x = start_x + new_width
660
+ start_y = (resized_height - new_height) // 2
661
+ end_y = start_y + new_height
662
+ samples = samples[:, :, start_y:end_y, start_x:end_x]
663
+
664
+ return samples
665
+
498
666
  @torch.no_grad()
499
667
  @replace_example_docstring(EXAMPLE_DOC_STRING)
500
668
  def __call__(
@@ -511,13 +679,16 @@ class PixArtAlphaPipeline(DiffusionPipeline):
511
679
  generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
512
680
  latents: Optional[torch.FloatTensor] = None,
513
681
  prompt_embeds: Optional[torch.FloatTensor] = None,
682
+ prompt_attention_mask: Optional[torch.FloatTensor] = None,
514
683
  negative_prompt_embeds: Optional[torch.FloatTensor] = None,
684
+ negative_prompt_attention_mask: Optional[torch.FloatTensor] = None,
515
685
  output_type: Optional[str] = "pil",
516
686
  return_dict: bool = True,
517
687
  callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
518
688
  callback_steps: int = 1,
519
689
  clean_caption: bool = True,
520
- mask_feature: bool = True,
690
+ use_resolution_binning: bool = True,
691
+ **kwargs,
521
692
  ) -> Union[ImagePipelineOutput, Tuple]:
522
693
  """
523
694
  Function invoked when calling the pipeline for generation.
@@ -536,7 +707,7 @@ class PixArtAlphaPipeline(DiffusionPipeline):
536
707
  timesteps (`List[int]`, *optional*):
537
708
  Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps`
538
709
  timesteps are used. Must be in descending order.
539
- guidance_scale (`float`, *optional*, defaults to 7.0):
710
+ guidance_scale (`float`, *optional*, defaults to 4.5):
540
711
  Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
541
712
  `guidance_scale` is defined as `w` of equation 2. of [Imagen
542
713
  Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
@@ -561,9 +732,12 @@ class PixArtAlphaPipeline(DiffusionPipeline):
561
732
  prompt_embeds (`torch.FloatTensor`, *optional*):
562
733
  Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
563
734
  provided, text embeddings will be generated from `prompt` input argument.
735
+ prompt_attention_mask (`torch.FloatTensor`, *optional*): Pre-generated attention mask for text embeddings.
564
736
  negative_prompt_embeds (`torch.FloatTensor`, *optional*):
565
737
  Pre-generated negative text embeddings. For PixArt-Alpha this negative prompt should be "". If not
566
738
  provided, negative_prompt_embeds will be generated from `negative_prompt` input argument.
739
+ negative_prompt_attention_mask (`torch.FloatTensor`, *optional*):
740
+ Pre-generated attention mask for negative text embeddings.
567
741
  output_type (`str`, *optional*, defaults to `"pil"`):
568
742
  The output format of the generate image. Choose between
569
743
  [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
@@ -579,7 +753,10 @@ class PixArtAlphaPipeline(DiffusionPipeline):
579
753
  Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to
580
754
  be installed. If the dependencies are not installed, the embeddings will be created from the raw
581
755
  prompt.
582
- mask_feature (`bool` defaults to `True`): If set to `True`, the text embeddings will be masked.
756
+ use_resolution_binning (`bool` defaults to `True`):
757
+ If set to `True`, the requested height and width are first mapped to the closest resolutions using
758
+ `ASPECT_RATIO_1024_BIN`. After the produced latents are decoded into images, they are resized back to
759
+ the requested resolution. Useful for generating non-square images.
583
760
 
584
761
  Examples:
585
762
 
@@ -588,11 +765,29 @@ class PixArtAlphaPipeline(DiffusionPipeline):
588
765
  If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is
589
766
  returned where the first element is a list with the generated images
590
767
  """
768
+ if "mask_feature" in kwargs:
769
+ deprecation_message = "The use of `mask_feature` is deprecated. It is no longer used in any computation and that doesn't affect the end results. It will be removed in a future version."
770
+ deprecate("mask_feature", "1.0.0", deprecation_message, standard_warn=False)
591
771
  # 1. Check inputs. Raise error if not correct
592
772
  height = height or self.transformer.config.sample_size * self.vae_scale_factor
593
773
  width = width or self.transformer.config.sample_size * self.vae_scale_factor
774
+ if use_resolution_binning:
775
+ aspect_ratio_bin = (
776
+ ASPECT_RATIO_1024_BIN if self.transformer.config.sample_size == 128 else ASPECT_RATIO_512_BIN
777
+ )
778
+ orig_height, orig_width = height, width
779
+ height, width = self.classify_height_width_bin(height, width, ratios=aspect_ratio_bin)
780
+
594
781
  self.check_inputs(
595
- prompt, height, width, negative_prompt, callback_steps, prompt_embeds, negative_prompt_embeds
782
+ prompt,
783
+ height,
784
+ width,
785
+ negative_prompt,
786
+ callback_steps,
787
+ prompt_embeds,
788
+ negative_prompt_embeds,
789
+ prompt_attention_mask,
790
+ negative_prompt_attention_mask,
596
791
  )
597
792
 
598
793
  # 2. Default height and width to transformer
@@ -611,7 +806,12 @@ class PixArtAlphaPipeline(DiffusionPipeline):
611
806
  do_classifier_free_guidance = guidance_scale > 1.0
612
807
 
613
808
  # 3. Encode input prompt
614
- prompt_embeds, negative_prompt_embeds = self.encode_prompt(
809
+ (
810
+ prompt_embeds,
811
+ prompt_attention_mask,
812
+ negative_prompt_embeds,
813
+ negative_prompt_attention_mask,
814
+ ) = self.encode_prompt(
615
815
  prompt,
616
816
  do_classifier_free_guidance,
617
817
  negative_prompt=negative_prompt,
@@ -619,15 +819,16 @@ class PixArtAlphaPipeline(DiffusionPipeline):
619
819
  device=device,
620
820
  prompt_embeds=prompt_embeds,
621
821
  negative_prompt_embeds=negative_prompt_embeds,
822
+ prompt_attention_mask=prompt_attention_mask,
823
+ negative_prompt_attention_mask=negative_prompt_attention_mask,
622
824
  clean_caption=clean_caption,
623
- mask_feature=mask_feature,
624
825
  )
625
826
  if do_classifier_free_guidance:
626
827
  prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
828
+ prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
627
829
 
628
830
  # 4. Prepare timesteps
629
- self.scheduler.set_timesteps(num_inference_steps, device=device)
630
- timesteps = self.scheduler.timesteps
831
+ timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
631
832
 
632
833
  # 5. Prepare latents.
633
834
  latent_channels = self.transformer.config.in_channels
@@ -652,6 +853,11 @@ class PixArtAlphaPipeline(DiffusionPipeline):
652
853
  aspect_ratio = torch.tensor([float(height / width)]).repeat(batch_size * num_images_per_prompt, 1)
653
854
  resolution = resolution.to(dtype=prompt_embeds.dtype, device=device)
654
855
  aspect_ratio = aspect_ratio.to(dtype=prompt_embeds.dtype, device=device)
856
+
857
+ if do_classifier_free_guidance:
858
+ resolution = torch.cat([resolution, resolution], dim=0)
859
+ aspect_ratio = torch.cat([aspect_ratio, aspect_ratio], dim=0)
860
+
655
861
  added_cond_kwargs = {"resolution": resolution, "aspect_ratio": aspect_ratio}
656
862
 
657
863
  # 7. Denoising loop
@@ -681,6 +887,7 @@ class PixArtAlphaPipeline(DiffusionPipeline):
681
887
  noise_pred = self.transformer(
682
888
  latent_model_input,
683
889
  encoder_hidden_states=prompt_embeds,
890
+ encoder_attention_mask=prompt_attention_mask,
684
891
  timestep=current_timestep,
685
892
  added_cond_kwargs=added_cond_kwargs,
686
893
  return_dict=False,
@@ -709,6 +916,8 @@ class PixArtAlphaPipeline(DiffusionPipeline):
709
916
 
710
917
  if not output_type == "latent":
711
918
  image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
919
+ if use_resolution_binning:
920
+ image = self.resize_and_crop_tensor(image, orig_width, orig_height)
712
921
  else:
713
922
  image = latents
714
923
 
@@ -283,6 +283,9 @@ class ShapEImg2ImgPipeline(DiffusionPipeline):
283
283
  f"Only the output types `pil`, `np`, `latent` and `mesh` are supported not output_type={output_type}"
284
284
  )
285
285
 
286
+ # Offload all models
287
+ self.maybe_free_model_hooks()
288
+
286
289
  if output_type == "latent":
287
290
  return ShapEPipelineOutput(images=latents)
288
291
 
@@ -312,9 +315,6 @@ class ShapEImg2ImgPipeline(DiffusionPipeline):
312
315
  if output_type == "pil":
313
316
  images = [self.numpy_to_pil(image) for image in images]
314
317
 
315
- # Offload all models
316
- self.maybe_free_model_hooks()
317
-
318
318
  if not return_dict:
319
319
  return (images,)
320
320