diffusers 0.27.1__py3-none-any.whl → 0.28.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 (270) hide show
  1. diffusers/__init__.py +18 -1
  2. diffusers/callbacks.py +156 -0
  3. diffusers/commands/env.py +110 -6
  4. diffusers/configuration_utils.py +16 -11
  5. diffusers/dependency_versions_table.py +2 -1
  6. diffusers/image_processor.py +158 -45
  7. diffusers/loaders/__init__.py +2 -5
  8. diffusers/loaders/autoencoder.py +4 -4
  9. diffusers/loaders/controlnet.py +4 -4
  10. diffusers/loaders/ip_adapter.py +80 -22
  11. diffusers/loaders/lora.py +134 -20
  12. diffusers/loaders/lora_conversion_utils.py +46 -43
  13. diffusers/loaders/peft.py +4 -3
  14. diffusers/loaders/single_file.py +401 -170
  15. diffusers/loaders/single_file_model.py +290 -0
  16. diffusers/loaders/single_file_utils.py +616 -672
  17. diffusers/loaders/textual_inversion.py +41 -20
  18. diffusers/loaders/unet.py +168 -115
  19. diffusers/loaders/unet_loader_utils.py +163 -0
  20. diffusers/models/__init__.py +2 -0
  21. diffusers/models/activations.py +11 -3
  22. diffusers/models/attention.py +10 -11
  23. diffusers/models/attention_processor.py +367 -148
  24. diffusers/models/autoencoders/autoencoder_asym_kl.py +14 -16
  25. diffusers/models/autoencoders/autoencoder_kl.py +18 -19
  26. diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py +11 -11
  27. diffusers/models/autoencoders/autoencoder_tiny.py +16 -16
  28. diffusers/models/autoencoders/consistency_decoder_vae.py +36 -11
  29. diffusers/models/autoencoders/vae.py +23 -24
  30. diffusers/models/controlnet.py +12 -9
  31. diffusers/models/controlnet_flax.py +4 -4
  32. diffusers/models/controlnet_xs.py +1915 -0
  33. diffusers/models/downsampling.py +17 -18
  34. diffusers/models/embeddings.py +147 -24
  35. diffusers/models/model_loading_utils.py +149 -0
  36. diffusers/models/modeling_flax_pytorch_utils.py +2 -1
  37. diffusers/models/modeling_flax_utils.py +4 -4
  38. diffusers/models/modeling_pytorch_flax_utils.py +1 -1
  39. diffusers/models/modeling_utils.py +118 -98
  40. diffusers/models/resnet.py +18 -23
  41. diffusers/models/transformer_temporal.py +3 -3
  42. diffusers/models/transformers/dual_transformer_2d.py +4 -4
  43. diffusers/models/transformers/prior_transformer.py +7 -7
  44. diffusers/models/transformers/t5_film_transformer.py +17 -19
  45. diffusers/models/transformers/transformer_2d.py +272 -156
  46. diffusers/models/transformers/transformer_temporal.py +10 -10
  47. diffusers/models/unets/unet_1d.py +5 -5
  48. diffusers/models/unets/unet_1d_blocks.py +29 -29
  49. diffusers/models/unets/unet_2d.py +6 -6
  50. diffusers/models/unets/unet_2d_blocks.py +137 -128
  51. diffusers/models/unets/unet_2d_condition.py +20 -15
  52. diffusers/models/unets/unet_2d_condition_flax.py +6 -5
  53. diffusers/models/unets/unet_3d_blocks.py +79 -77
  54. diffusers/models/unets/unet_3d_condition.py +13 -9
  55. diffusers/models/unets/unet_i2vgen_xl.py +14 -13
  56. diffusers/models/unets/unet_kandinsky3.py +1 -1
  57. diffusers/models/unets/unet_motion_model.py +114 -14
  58. diffusers/models/unets/unet_spatio_temporal_condition.py +15 -14
  59. diffusers/models/unets/unet_stable_cascade.py +16 -13
  60. diffusers/models/upsampling.py +17 -20
  61. diffusers/models/vq_model.py +16 -15
  62. diffusers/pipelines/__init__.py +25 -3
  63. diffusers/pipelines/amused/pipeline_amused.py +12 -12
  64. diffusers/pipelines/amused/pipeline_amused_img2img.py +14 -12
  65. diffusers/pipelines/amused/pipeline_amused_inpaint.py +13 -11
  66. diffusers/pipelines/animatediff/__init__.py +2 -0
  67. diffusers/pipelines/animatediff/pipeline_animatediff.py +24 -46
  68. diffusers/pipelines/animatediff/pipeline_animatediff_sdxl.py +1284 -0
  69. diffusers/pipelines/animatediff/pipeline_animatediff_video2video.py +82 -72
  70. diffusers/pipelines/animatediff/pipeline_output.py +3 -2
  71. diffusers/pipelines/audioldm/pipeline_audioldm.py +14 -14
  72. diffusers/pipelines/audioldm2/modeling_audioldm2.py +54 -35
  73. diffusers/pipelines/audioldm2/pipeline_audioldm2.py +120 -36
  74. diffusers/pipelines/auto_pipeline.py +21 -17
  75. diffusers/pipelines/blip_diffusion/blip_image_processing.py +1 -1
  76. diffusers/pipelines/blip_diffusion/modeling_blip2.py +5 -5
  77. diffusers/pipelines/blip_diffusion/modeling_ctx_clip.py +1 -1
  78. diffusers/pipelines/blip_diffusion/pipeline_blip_diffusion.py +2 -2
  79. diffusers/pipelines/consistency_models/pipeline_consistency_models.py +5 -5
  80. diffusers/pipelines/controlnet/multicontrolnet.py +4 -8
  81. diffusers/pipelines/controlnet/pipeline_controlnet.py +87 -52
  82. diffusers/pipelines/controlnet/pipeline_controlnet_blip_diffusion.py +2 -2
  83. diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py +50 -43
  84. diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py +52 -40
  85. diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py +80 -47
  86. diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py +147 -49
  87. diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py +89 -55
  88. diffusers/pipelines/controlnet_xs/__init__.py +68 -0
  89. diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs.py +911 -0
  90. diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs_sd_xl.py +1115 -0
  91. diffusers/pipelines/deepfloyd_if/pipeline_if.py +14 -28
  92. diffusers/pipelines/deepfloyd_if/pipeline_if_img2img.py +18 -33
  93. diffusers/pipelines/deepfloyd_if/pipeline_if_img2img_superresolution.py +21 -39
  94. diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting.py +20 -36
  95. diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting_superresolution.py +23 -39
  96. diffusers/pipelines/deepfloyd_if/pipeline_if_superresolution.py +17 -32
  97. diffusers/pipelines/deprecated/alt_diffusion/modeling_roberta_series.py +11 -11
  98. diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion.py +43 -20
  99. diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion_img2img.py +36 -18
  100. diffusers/pipelines/deprecated/repaint/pipeline_repaint.py +2 -2
  101. diffusers/pipelines/deprecated/spectrogram_diffusion/pipeline_spectrogram_diffusion.py +7 -7
  102. diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_cycle_diffusion.py +12 -12
  103. diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_inpaint_legacy.py +18 -21
  104. diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_model_editing.py +20 -15
  105. diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_paradigms.py +20 -15
  106. diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_pix2pix_zero.py +30 -25
  107. diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py +69 -59
  108. diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion.py +13 -13
  109. diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_dual_guided.py +10 -5
  110. diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_image_variation.py +11 -6
  111. diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_text_to_image.py +10 -5
  112. diffusers/pipelines/deprecated/vq_diffusion/pipeline_vq_diffusion.py +5 -5
  113. diffusers/pipelines/dit/pipeline_dit.py +3 -0
  114. diffusers/pipelines/free_init_utils.py +39 -38
  115. diffusers/pipelines/i2vgen_xl/pipeline_i2vgen_xl.py +33 -48
  116. diffusers/pipelines/kandinsky/pipeline_kandinsky.py +8 -8
  117. diffusers/pipelines/kandinsky/pipeline_kandinsky_combined.py +23 -20
  118. diffusers/pipelines/kandinsky/pipeline_kandinsky_img2img.py +11 -11
  119. diffusers/pipelines/kandinsky/pipeline_kandinsky_inpaint.py +12 -12
  120. diffusers/pipelines/kandinsky/pipeline_kandinsky_prior.py +10 -10
  121. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py +6 -6
  122. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py +32 -29
  123. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet.py +10 -10
  124. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet_img2img.py +10 -10
  125. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py +6 -6
  126. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpainting.py +8 -8
  127. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py +7 -7
  128. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior_emb2emb.py +6 -6
  129. diffusers/pipelines/kandinsky3/convert_kandinsky3_unet.py +3 -3
  130. diffusers/pipelines/kandinsky3/pipeline_kandinsky3.py +20 -33
  131. diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py +24 -35
  132. diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_img2img.py +48 -30
  133. diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_text2img.py +50 -28
  134. diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py +11 -11
  135. diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py +61 -67
  136. diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py +70 -69
  137. diffusers/pipelines/ledits_pp/pipeline_output.py +2 -2
  138. diffusers/pipelines/marigold/__init__.py +50 -0
  139. diffusers/pipelines/marigold/marigold_image_processing.py +561 -0
  140. diffusers/pipelines/marigold/pipeline_marigold_depth.py +813 -0
  141. diffusers/pipelines/marigold/pipeline_marigold_normals.py +690 -0
  142. diffusers/pipelines/musicldm/pipeline_musicldm.py +14 -14
  143. diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py +17 -12
  144. diffusers/pipelines/pia/pipeline_pia.py +39 -125
  145. diffusers/pipelines/pipeline_flax_utils.py +4 -4
  146. diffusers/pipelines/pipeline_loading_utils.py +268 -23
  147. diffusers/pipelines/pipeline_utils.py +266 -37
  148. diffusers/pipelines/pixart_alpha/__init__.py +8 -1
  149. diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py +65 -75
  150. diffusers/pipelines/pixart_alpha/pipeline_pixart_sigma.py +880 -0
  151. diffusers/pipelines/semantic_stable_diffusion/pipeline_semantic_stable_diffusion.py +10 -5
  152. diffusers/pipelines/shap_e/pipeline_shap_e.py +3 -3
  153. diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py +14 -14
  154. diffusers/pipelines/shap_e/renderer.py +1 -1
  155. diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py +36 -22
  156. diffusers/pipelines/stable_cascade/pipeline_stable_cascade_combined.py +23 -19
  157. diffusers/pipelines/stable_cascade/pipeline_stable_cascade_prior.py +33 -32
  158. diffusers/pipelines/stable_diffusion/__init__.py +0 -1
  159. diffusers/pipelines/stable_diffusion/convert_from_ckpt.py +18 -11
  160. diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py +2 -2
  161. diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_upscale.py +6 -6
  162. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +73 -39
  163. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_depth2img.py +24 -17
  164. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py +13 -8
  165. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py +66 -36
  166. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py +82 -46
  167. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py +123 -28
  168. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_latent_upscale.py +6 -6
  169. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py +16 -16
  170. diffusers/pipelines/stable_diffusion/pipeline_stable_unclip.py +24 -19
  171. diffusers/pipelines/stable_diffusion/pipeline_stable_unclip_img2img.py +37 -31
  172. diffusers/pipelines/stable_diffusion/safety_checker.py +2 -1
  173. diffusers/pipelines/stable_diffusion_attend_and_excite/pipeline_stable_diffusion_attend_and_excite.py +23 -15
  174. diffusers/pipelines/stable_diffusion_diffedit/pipeline_stable_diffusion_diffedit.py +44 -42
  175. diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen.py +23 -18
  176. diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen_text_image.py +19 -14
  177. diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_k_diffusion.py +20 -15
  178. diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_xl_k_diffusion.py +24 -19
  179. diffusers/pipelines/stable_diffusion_ldm3d/pipeline_stable_diffusion_ldm3d.py +65 -32
  180. diffusers/pipelines/stable_diffusion_panorama/pipeline_stable_diffusion_panorama.py +274 -38
  181. diffusers/pipelines/stable_diffusion_safe/pipeline_stable_diffusion_safe.py +10 -5
  182. diffusers/pipelines/stable_diffusion_safe/safety_checker.py +1 -1
  183. diffusers/pipelines/stable_diffusion_sag/pipeline_stable_diffusion_sag.py +92 -25
  184. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py +88 -44
  185. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py +108 -56
  186. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py +96 -51
  187. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_instruct_pix2pix.py +45 -25
  188. diffusers/pipelines/stable_diffusion_xl/watermark.py +9 -3
  189. diffusers/pipelines/stable_video_diffusion/pipeline_stable_video_diffusion.py +110 -57
  190. diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_adapter.py +59 -30
  191. diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_xl_adapter.py +71 -42
  192. diffusers/pipelines/text_to_video_synthesis/pipeline_output.py +3 -2
  193. diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth.py +18 -41
  194. diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth_img2img.py +21 -85
  195. diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero.py +28 -19
  196. diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero_sdxl.py +39 -33
  197. diffusers/pipelines/unclip/pipeline_unclip.py +6 -6
  198. diffusers/pipelines/unclip/pipeline_unclip_image_variation.py +6 -6
  199. diffusers/pipelines/unidiffuser/modeling_text_decoder.py +1 -1
  200. diffusers/pipelines/unidiffuser/modeling_uvit.py +9 -9
  201. diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py +23 -23
  202. diffusers/pipelines/wuerstchen/modeling_paella_vq_model.py +5 -5
  203. diffusers/pipelines/wuerstchen/modeling_wuerstchen_common.py +5 -10
  204. diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py +4 -6
  205. diffusers/pipelines/wuerstchen/pipeline_wuerstchen.py +4 -4
  206. diffusers/pipelines/wuerstchen/pipeline_wuerstchen_combined.py +12 -12
  207. diffusers/pipelines/wuerstchen/pipeline_wuerstchen_prior.py +10 -10
  208. diffusers/schedulers/__init__.py +2 -2
  209. diffusers/schedulers/deprecated/__init__.py +1 -1
  210. diffusers/schedulers/deprecated/scheduling_karras_ve.py +25 -25
  211. diffusers/schedulers/scheduling_amused.py +5 -5
  212. diffusers/schedulers/scheduling_consistency_decoder.py +11 -11
  213. diffusers/schedulers/scheduling_consistency_models.py +23 -25
  214. diffusers/schedulers/scheduling_ddim.py +22 -24
  215. diffusers/schedulers/scheduling_ddim_flax.py +2 -1
  216. diffusers/schedulers/scheduling_ddim_inverse.py +16 -16
  217. diffusers/schedulers/scheduling_ddim_parallel.py +28 -30
  218. diffusers/schedulers/scheduling_ddpm.py +20 -22
  219. diffusers/schedulers/scheduling_ddpm_flax.py +7 -3
  220. diffusers/schedulers/scheduling_ddpm_parallel.py +26 -28
  221. diffusers/schedulers/scheduling_ddpm_wuerstchen.py +14 -14
  222. diffusers/schedulers/scheduling_deis_multistep.py +46 -42
  223. diffusers/schedulers/scheduling_dpmsolver_multistep.py +107 -77
  224. diffusers/schedulers/scheduling_dpmsolver_multistep_flax.py +2 -2
  225. diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py +46 -46
  226. diffusers/schedulers/scheduling_dpmsolver_sde.py +26 -22
  227. diffusers/schedulers/scheduling_dpmsolver_singlestep.py +90 -65
  228. diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py +78 -53
  229. diffusers/schedulers/scheduling_edm_euler.py +53 -30
  230. diffusers/schedulers/scheduling_euler_ancestral_discrete.py +26 -28
  231. diffusers/schedulers/scheduling_euler_discrete.py +163 -67
  232. diffusers/schedulers/scheduling_heun_discrete.py +60 -38
  233. diffusers/schedulers/scheduling_ipndm.py +8 -8
  234. diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py +22 -18
  235. diffusers/schedulers/scheduling_k_dpm_2_discrete.py +22 -18
  236. diffusers/schedulers/scheduling_karras_ve_flax.py +6 -6
  237. diffusers/schedulers/scheduling_lcm.py +21 -23
  238. diffusers/schedulers/scheduling_lms_discrete.py +27 -25
  239. diffusers/schedulers/scheduling_pndm.py +20 -20
  240. diffusers/schedulers/scheduling_repaint.py +20 -20
  241. diffusers/schedulers/scheduling_sasolver.py +55 -54
  242. diffusers/schedulers/scheduling_sde_ve.py +19 -19
  243. diffusers/schedulers/scheduling_tcd.py +39 -30
  244. diffusers/schedulers/scheduling_unclip.py +15 -15
  245. diffusers/schedulers/scheduling_unipc_multistep.py +115 -41
  246. diffusers/schedulers/scheduling_utils.py +14 -5
  247. diffusers/schedulers/scheduling_utils_flax.py +3 -3
  248. diffusers/schedulers/scheduling_vq_diffusion.py +10 -10
  249. diffusers/training_utils.py +56 -1
  250. diffusers/utils/__init__.py +7 -0
  251. diffusers/utils/doc_utils.py +1 -0
  252. diffusers/utils/dummy_pt_objects.py +30 -0
  253. diffusers/utils/dummy_torch_and_transformers_objects.py +90 -0
  254. diffusers/utils/dynamic_modules_utils.py +24 -11
  255. diffusers/utils/hub_utils.py +3 -2
  256. diffusers/utils/import_utils.py +91 -0
  257. diffusers/utils/loading_utils.py +2 -2
  258. diffusers/utils/logging.py +1 -1
  259. diffusers/utils/peft_utils.py +32 -5
  260. diffusers/utils/state_dict_utils.py +11 -2
  261. diffusers/utils/testing_utils.py +71 -6
  262. diffusers/utils/torch_utils.py +1 -0
  263. diffusers/video_processor.py +113 -0
  264. {diffusers-0.27.1.dist-info → diffusers-0.28.0.dist-info}/METADATA +7 -7
  265. diffusers-0.28.0.dist-info/RECORD +414 -0
  266. diffusers-0.27.1.dist-info/RECORD +0 -399
  267. {diffusers-0.27.1.dist-info → diffusers-0.28.0.dist-info}/LICENSE +0 -0
  268. {diffusers-0.27.1.dist-info → diffusers-0.28.0.dist-info}/WHEEL +0 -0
  269. {diffusers-0.27.1.dist-info → diffusers-0.28.0.dist-info}/entry_points.txt +0 -0
  270. {diffusers-0.27.1.dist-info → diffusers-0.28.0.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,880 @@
1
+ # Copyright 2024 PixArt-Sigma Authors and 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
+ import html
16
+ import inspect
17
+ import re
18
+ import urllib.parse as ul
19
+ from typing import Callable, List, Optional, Tuple, Union
20
+
21
+ import torch
22
+ from transformers import T5EncoderModel, T5Tokenizer
23
+
24
+ from ...image_processor import PixArtImageProcessor
25
+ from ...models import AutoencoderKL, Transformer2DModel
26
+ from ...schedulers import KarrasDiffusionSchedulers
27
+ from ...utils import (
28
+ BACKENDS_MAPPING,
29
+ deprecate,
30
+ is_bs4_available,
31
+ is_ftfy_available,
32
+ logging,
33
+ replace_example_docstring,
34
+ )
35
+ from ...utils.torch_utils import randn_tensor
36
+ from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput
37
+ from .pipeline_pixart_alpha import (
38
+ ASPECT_RATIO_256_BIN,
39
+ ASPECT_RATIO_512_BIN,
40
+ ASPECT_RATIO_1024_BIN,
41
+ )
42
+
43
+
44
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
45
+
46
+ if is_bs4_available():
47
+ from bs4 import BeautifulSoup
48
+
49
+ if is_ftfy_available():
50
+ import ftfy
51
+
52
+
53
+ ASPECT_RATIO_2048_BIN = {
54
+ "0.25": [1024.0, 4096.0],
55
+ "0.26": [1024.0, 3968.0],
56
+ "0.27": [1024.0, 3840.0],
57
+ "0.28": [1024.0, 3712.0],
58
+ "0.32": [1152.0, 3584.0],
59
+ "0.33": [1152.0, 3456.0],
60
+ "0.35": [1152.0, 3328.0],
61
+ "0.4": [1280.0, 3200.0],
62
+ "0.42": [1280.0, 3072.0],
63
+ "0.48": [1408.0, 2944.0],
64
+ "0.5": [1408.0, 2816.0],
65
+ "0.52": [1408.0, 2688.0],
66
+ "0.57": [1536.0, 2688.0],
67
+ "0.6": [1536.0, 2560.0],
68
+ "0.68": [1664.0, 2432.0],
69
+ "0.72": [1664.0, 2304.0],
70
+ "0.78": [1792.0, 2304.0],
71
+ "0.82": [1792.0, 2176.0],
72
+ "0.88": [1920.0, 2176.0],
73
+ "0.94": [1920.0, 2048.0],
74
+ "1.0": [2048.0, 2048.0],
75
+ "1.07": [2048.0, 1920.0],
76
+ "1.13": [2176.0, 1920.0],
77
+ "1.21": [2176.0, 1792.0],
78
+ "1.29": [2304.0, 1792.0],
79
+ "1.38": [2304.0, 1664.0],
80
+ "1.46": [2432.0, 1664.0],
81
+ "1.67": [2560.0, 1536.0],
82
+ "1.75": [2688.0, 1536.0],
83
+ "2.0": [2816.0, 1408.0],
84
+ "2.09": [2944.0, 1408.0],
85
+ "2.4": [3072.0, 1280.0],
86
+ "2.5": [3200.0, 1280.0],
87
+ "2.89": [3328.0, 1152.0],
88
+ "3.0": [3456.0, 1152.0],
89
+ "3.11": [3584.0, 1152.0],
90
+ "3.62": [3712.0, 1024.0],
91
+ "3.75": [3840.0, 1024.0],
92
+ "3.88": [3968.0, 1024.0],
93
+ "4.0": [4096.0, 1024.0],
94
+ }
95
+
96
+
97
+ EXAMPLE_DOC_STRING = """
98
+ Examples:
99
+ ```py
100
+ >>> import torch
101
+ >>> from diffusers import PixArtSigmaPipeline
102
+
103
+ >>> # You can replace the checkpoint id with "PixArt-alpha/PixArt-Sigma-XL-2-512-MS" too.
104
+ >>> pipe = PixArtSigmaPipeline.from_pretrained(
105
+ ... "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", torch_dtype=torch.float16
106
+ ... )
107
+ >>> # Enable memory optimizations.
108
+ >>> # pipe.enable_model_cpu_offload()
109
+
110
+ >>> prompt = "A small cactus with a happy face in the Sahara desert."
111
+ >>> image = pipe(prompt).images[0]
112
+ ```
113
+ """
114
+
115
+
116
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
117
+ def retrieve_timesteps(
118
+ scheduler,
119
+ num_inference_steps: Optional[int] = None,
120
+ device: Optional[Union[str, torch.device]] = None,
121
+ timesteps: Optional[List[int]] = None,
122
+ sigmas: Optional[List[float]] = None,
123
+ **kwargs,
124
+ ):
125
+ """
126
+ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
127
+ custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
128
+
129
+ Args:
130
+ scheduler (`SchedulerMixin`):
131
+ The scheduler to get timesteps from.
132
+ num_inference_steps (`int`):
133
+ The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
134
+ must be `None`.
135
+ device (`str` or `torch.device`, *optional*):
136
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
137
+ timesteps (`List[int]`, *optional*):
138
+ Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
139
+ `num_inference_steps` and `sigmas` must be `None`.
140
+ sigmas (`List[float]`, *optional*):
141
+ Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
142
+ `num_inference_steps` and `timesteps` must be `None`.
143
+
144
+ Returns:
145
+ `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
146
+ second element is the number of inference steps.
147
+ """
148
+ if timesteps is not None and sigmas is not None:
149
+ raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
150
+ if timesteps is not None:
151
+ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
152
+ if not accepts_timesteps:
153
+ raise ValueError(
154
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
155
+ f" timestep schedules. Please check whether you are using the correct scheduler."
156
+ )
157
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
158
+ timesteps = scheduler.timesteps
159
+ num_inference_steps = len(timesteps)
160
+ elif sigmas is not None:
161
+ accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
162
+ if not accept_sigmas:
163
+ raise ValueError(
164
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
165
+ f" sigmas schedules. Please check whether you are using the correct scheduler."
166
+ )
167
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
168
+ timesteps = scheduler.timesteps
169
+ num_inference_steps = len(timesteps)
170
+ else:
171
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
172
+ timesteps = scheduler.timesteps
173
+ return timesteps, num_inference_steps
174
+
175
+
176
+ class PixArtSigmaPipeline(DiffusionPipeline):
177
+ r"""
178
+ Pipeline for text-to-image generation using PixArt-Sigma.
179
+ """
180
+
181
+ bad_punct_regex = re.compile(
182
+ r"["
183
+ + "#®•©™&@·º½¾¿¡§~"
184
+ + r"\)"
185
+ + r"\("
186
+ + r"\]"
187
+ + r"\["
188
+ + r"\}"
189
+ + r"\{"
190
+ + r"\|"
191
+ + "\\"
192
+ + r"\/"
193
+ + r"\*"
194
+ + r"]{1,}"
195
+ ) # noqa
196
+
197
+ _optional_components = ["tokenizer", "text_encoder"]
198
+ model_cpu_offload_seq = "text_encoder->transformer->vae"
199
+
200
+ def __init__(
201
+ self,
202
+ tokenizer: T5Tokenizer,
203
+ text_encoder: T5EncoderModel,
204
+ vae: AutoencoderKL,
205
+ transformer: Transformer2DModel,
206
+ scheduler: KarrasDiffusionSchedulers,
207
+ ):
208
+ super().__init__()
209
+
210
+ self.register_modules(
211
+ tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler
212
+ )
213
+
214
+ self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
215
+ self.image_processor = PixArtImageProcessor(vae_scale_factor=self.vae_scale_factor)
216
+
217
+ # Copied from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha.PixArtAlphaPipeline.encode_prompt with 120->300
218
+ def encode_prompt(
219
+ self,
220
+ prompt: Union[str, List[str]],
221
+ do_classifier_free_guidance: bool = True,
222
+ negative_prompt: str = "",
223
+ num_images_per_prompt: int = 1,
224
+ device: Optional[torch.device] = None,
225
+ prompt_embeds: Optional[torch.Tensor] = None,
226
+ negative_prompt_embeds: Optional[torch.Tensor] = None,
227
+ prompt_attention_mask: Optional[torch.Tensor] = None,
228
+ negative_prompt_attention_mask: Optional[torch.Tensor] = None,
229
+ clean_caption: bool = False,
230
+ max_sequence_length: int = 300,
231
+ **kwargs,
232
+ ):
233
+ r"""
234
+ Encodes the prompt into text encoder hidden states.
235
+
236
+ Args:
237
+ prompt (`str` or `List[str]`, *optional*):
238
+ prompt to be encoded
239
+ negative_prompt (`str` or `List[str]`, *optional*):
240
+ The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`
241
+ instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For
242
+ PixArt-Alpha, this should be "".
243
+ do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
244
+ whether to use classifier free guidance or not
245
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
246
+ number of images that should be generated per prompt
247
+ device: (`torch.device`, *optional*):
248
+ torch device to place the resulting embeddings on
249
+ prompt_embeds (`torch.Tensor`, *optional*):
250
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
251
+ provided, text embeddings will be generated from `prompt` input argument.
252
+ negative_prompt_embeds (`torch.Tensor`, *optional*):
253
+ Pre-generated negative text embeddings. For PixArt-Alpha, it's should be the embeddings of the ""
254
+ string.
255
+ clean_caption (`bool`, defaults to `False`):
256
+ If `True`, the function will preprocess and clean the provided caption before encoding.
257
+ max_sequence_length (`int`, defaults to 300): Maximum sequence length to use for the prompt.
258
+ """
259
+
260
+ if "mask_feature" in kwargs:
261
+ 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."
262
+ deprecate("mask_feature", "1.0.0", deprecation_message, standard_warn=False)
263
+
264
+ if device is None:
265
+ device = self._execution_device
266
+
267
+ if prompt is not None and isinstance(prompt, str):
268
+ batch_size = 1
269
+ elif prompt is not None and isinstance(prompt, list):
270
+ batch_size = len(prompt)
271
+ else:
272
+ batch_size = prompt_embeds.shape[0]
273
+
274
+ # See Section 3.1. of the paper.
275
+ max_length = max_sequence_length
276
+
277
+ if prompt_embeds is None:
278
+ prompt = self._text_preprocessing(prompt, clean_caption=clean_caption)
279
+ text_inputs = self.tokenizer(
280
+ prompt,
281
+ padding="max_length",
282
+ max_length=max_length,
283
+ truncation=True,
284
+ add_special_tokens=True,
285
+ return_tensors="pt",
286
+ )
287
+ text_input_ids = text_inputs.input_ids
288
+ untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
289
+
290
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
291
+ text_input_ids, untruncated_ids
292
+ ):
293
+ removed_text = self.tokenizer.batch_decode(untruncated_ids[:, max_length - 1 : -1])
294
+ logger.warning(
295
+ "The following part of your input was truncated because T5 can only handle sequences up to"
296
+ f" {max_length} tokens: {removed_text}"
297
+ )
298
+
299
+ prompt_attention_mask = text_inputs.attention_mask
300
+ prompt_attention_mask = prompt_attention_mask.to(device)
301
+
302
+ prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask)
303
+ prompt_embeds = prompt_embeds[0]
304
+
305
+ if self.text_encoder is not None:
306
+ dtype = self.text_encoder.dtype
307
+ elif self.transformer is not None:
308
+ dtype = self.transformer.dtype
309
+ else:
310
+ dtype = None
311
+
312
+ prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
313
+
314
+ bs_embed, seq_len, _ = prompt_embeds.shape
315
+ # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
316
+ prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
317
+ prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
318
+ prompt_attention_mask = prompt_attention_mask.view(bs_embed, -1)
319
+ prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1)
320
+
321
+ # get unconditional embeddings for classifier free guidance
322
+ if do_classifier_free_guidance and negative_prompt_embeds is None:
323
+ uncond_tokens = [negative_prompt] * batch_size
324
+ uncond_tokens = self._text_preprocessing(uncond_tokens, clean_caption=clean_caption)
325
+ max_length = prompt_embeds.shape[1]
326
+ uncond_input = self.tokenizer(
327
+ uncond_tokens,
328
+ padding="max_length",
329
+ max_length=max_length,
330
+ truncation=True,
331
+ return_attention_mask=True,
332
+ add_special_tokens=True,
333
+ return_tensors="pt",
334
+ )
335
+ negative_prompt_attention_mask = uncond_input.attention_mask
336
+ negative_prompt_attention_mask = negative_prompt_attention_mask.to(device)
337
+
338
+ negative_prompt_embeds = self.text_encoder(
339
+ uncond_input.input_ids.to(device), attention_mask=negative_prompt_attention_mask
340
+ )
341
+ negative_prompt_embeds = negative_prompt_embeds[0]
342
+
343
+ if do_classifier_free_guidance:
344
+ # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
345
+ seq_len = negative_prompt_embeds.shape[1]
346
+
347
+ negative_prompt_embeds = negative_prompt_embeds.to(dtype=dtype, device=device)
348
+
349
+ negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
350
+ negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
351
+
352
+ negative_prompt_attention_mask = negative_prompt_attention_mask.view(bs_embed, -1)
353
+ negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 1)
354
+ else:
355
+ negative_prompt_embeds = None
356
+ negative_prompt_attention_mask = None
357
+
358
+ return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask
359
+
360
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
361
+ def prepare_extra_step_kwargs(self, generator, eta):
362
+ # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
363
+ # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
364
+ # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
365
+ # and should be between [0, 1]
366
+
367
+ accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
368
+ extra_step_kwargs = {}
369
+ if accepts_eta:
370
+ extra_step_kwargs["eta"] = eta
371
+
372
+ # check if the scheduler accepts generator
373
+ accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
374
+ if accepts_generator:
375
+ extra_step_kwargs["generator"] = generator
376
+ return extra_step_kwargs
377
+
378
+ # Copied from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha.PixArtAlphaPipeline.check_inputs
379
+ def check_inputs(
380
+ self,
381
+ prompt,
382
+ height,
383
+ width,
384
+ negative_prompt,
385
+ callback_steps,
386
+ prompt_embeds=None,
387
+ negative_prompt_embeds=None,
388
+ prompt_attention_mask=None,
389
+ negative_prompt_attention_mask=None,
390
+ ):
391
+ if height % 8 != 0 or width % 8 != 0:
392
+ raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
393
+
394
+ if (callback_steps is None) or (
395
+ callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
396
+ ):
397
+ raise ValueError(
398
+ f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
399
+ f" {type(callback_steps)}."
400
+ )
401
+
402
+ if prompt is not None and prompt_embeds is not None:
403
+ raise ValueError(
404
+ f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
405
+ " only forward one of the two."
406
+ )
407
+ elif prompt is None and prompt_embeds is None:
408
+ raise ValueError(
409
+ "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
410
+ )
411
+ elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
412
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
413
+
414
+ if prompt is not None and negative_prompt_embeds is not None:
415
+ raise ValueError(
416
+ f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:"
417
+ f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
418
+ )
419
+
420
+ if negative_prompt is not None and negative_prompt_embeds is not None:
421
+ raise ValueError(
422
+ f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
423
+ f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
424
+ )
425
+
426
+ if prompt_embeds is not None and prompt_attention_mask is None:
427
+ raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.")
428
+
429
+ if negative_prompt_embeds is not None and negative_prompt_attention_mask is None:
430
+ raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.")
431
+
432
+ if prompt_embeds is not None and negative_prompt_embeds is not None:
433
+ if prompt_embeds.shape != negative_prompt_embeds.shape:
434
+ raise ValueError(
435
+ "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
436
+ f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
437
+ f" {negative_prompt_embeds.shape}."
438
+ )
439
+ if prompt_attention_mask.shape != negative_prompt_attention_mask.shape:
440
+ raise ValueError(
441
+ "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but"
442
+ f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`"
443
+ f" {negative_prompt_attention_mask.shape}."
444
+ )
445
+
446
+ # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing
447
+ def _text_preprocessing(self, text, clean_caption=False):
448
+ if clean_caption and not is_bs4_available():
449
+ logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`"))
450
+ logger.warning("Setting `clean_caption` to False...")
451
+ clean_caption = False
452
+
453
+ if clean_caption and not is_ftfy_available():
454
+ logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`"))
455
+ logger.warning("Setting `clean_caption` to False...")
456
+ clean_caption = False
457
+
458
+ if not isinstance(text, (tuple, list)):
459
+ text = [text]
460
+
461
+ def process(text: str):
462
+ if clean_caption:
463
+ text = self._clean_caption(text)
464
+ text = self._clean_caption(text)
465
+ else:
466
+ text = text.lower().strip()
467
+ return text
468
+
469
+ return [process(t) for t in text]
470
+
471
+ # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption
472
+ def _clean_caption(self, caption):
473
+ caption = str(caption)
474
+ caption = ul.unquote_plus(caption)
475
+ caption = caption.strip().lower()
476
+ caption = re.sub("<person>", "person", caption)
477
+ # urls:
478
+ caption = re.sub(
479
+ r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa
480
+ "",
481
+ caption,
482
+ ) # regex for urls
483
+ caption = re.sub(
484
+ r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa
485
+ "",
486
+ caption,
487
+ ) # regex for urls
488
+ # html:
489
+ caption = BeautifulSoup(caption, features="html.parser").text
490
+
491
+ # @<nickname>
492
+ caption = re.sub(r"@[\w\d]+\b", "", caption)
493
+
494
+ # 31C0—31EF CJK Strokes
495
+ # 31F0—31FF Katakana Phonetic Extensions
496
+ # 3200—32FF Enclosed CJK Letters and Months
497
+ # 3300—33FF CJK Compatibility
498
+ # 3400—4DBF CJK Unified Ideographs Extension A
499
+ # 4DC0—4DFF Yijing Hexagram Symbols
500
+ # 4E00—9FFF CJK Unified Ideographs
501
+ caption = re.sub(r"[\u31c0-\u31ef]+", "", caption)
502
+ caption = re.sub(r"[\u31f0-\u31ff]+", "", caption)
503
+ caption = re.sub(r"[\u3200-\u32ff]+", "", caption)
504
+ caption = re.sub(r"[\u3300-\u33ff]+", "", caption)
505
+ caption = re.sub(r"[\u3400-\u4dbf]+", "", caption)
506
+ caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption)
507
+ caption = re.sub(r"[\u4e00-\u9fff]+", "", caption)
508
+ #######################################################
509
+
510
+ # все виды тире / all types of dash --> "-"
511
+ caption = re.sub(
512
+ r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+", # noqa
513
+ "-",
514
+ caption,
515
+ )
516
+
517
+ # кавычки к одному стандарту
518
+ caption = re.sub(r"[`´«»“”¨]", '"', caption)
519
+ caption = re.sub(r"[‘’]", "'", caption)
520
+
521
+ # &quot;
522
+ caption = re.sub(r"&quot;?", "", caption)
523
+ # &amp
524
+ caption = re.sub(r"&amp", "", caption)
525
+
526
+ # ip adresses:
527
+ caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption)
528
+
529
+ # article ids:
530
+ caption = re.sub(r"\d:\d\d\s+$", "", caption)
531
+
532
+ # \n
533
+ caption = re.sub(r"\\n", " ", caption)
534
+
535
+ # "#123"
536
+ caption = re.sub(r"#\d{1,3}\b", "", caption)
537
+ # "#12345.."
538
+ caption = re.sub(r"#\d{5,}\b", "", caption)
539
+ # "123456.."
540
+ caption = re.sub(r"\b\d{6,}\b", "", caption)
541
+ # filenames:
542
+ caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption)
543
+
544
+ #
545
+ caption = re.sub(r"[\"\']{2,}", r'"', caption) # """AUSVERKAUFT"""
546
+ caption = re.sub(r"[\.]{2,}", r" ", caption) # """AUSVERKAUFT"""
547
+
548
+ caption = re.sub(self.bad_punct_regex, r" ", caption) # ***AUSVERKAUFT***, #AUSVERKAUFT
549
+ caption = re.sub(r"\s+\.\s+", r" ", caption) # " . "
550
+
551
+ # this-is-my-cute-cat / this_is_my_cute_cat
552
+ regex2 = re.compile(r"(?:\-|\_)")
553
+ if len(re.findall(regex2, caption)) > 3:
554
+ caption = re.sub(regex2, " ", caption)
555
+
556
+ caption = ftfy.fix_text(caption)
557
+ caption = html.unescape(html.unescape(caption))
558
+
559
+ caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption) # jc6640
560
+ caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption) # jc6640vc
561
+ caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption) # 6640vc231
562
+
563
+ caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption)
564
+ caption = re.sub(r"(free\s)?download(\sfree)?", "", caption)
565
+ caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption)
566
+ caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption)
567
+ caption = re.sub(r"\bpage\s+\d+\b", "", caption)
568
+
569
+ caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption) # j2d1a2a...
570
+
571
+ caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption)
572
+
573
+ caption = re.sub(r"\b\s+\:\s+", r": ", caption)
574
+ caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption)
575
+ caption = re.sub(r"\s+", " ", caption)
576
+
577
+ caption.strip()
578
+
579
+ caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption)
580
+ caption = re.sub(r"^[\'\_,\-\:;]", r"", caption)
581
+ caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption)
582
+ caption = re.sub(r"^\.\S+$", "", caption)
583
+
584
+ return caption.strip()
585
+
586
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents
587
+ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
588
+ shape = (
589
+ batch_size,
590
+ num_channels_latents,
591
+ int(height) // self.vae_scale_factor,
592
+ int(width) // self.vae_scale_factor,
593
+ )
594
+ if isinstance(generator, list) and len(generator) != batch_size:
595
+ raise ValueError(
596
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
597
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
598
+ )
599
+
600
+ if latents is None:
601
+ latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
602
+ else:
603
+ latents = latents.to(device)
604
+
605
+ # scale the initial noise by the standard deviation required by the scheduler
606
+ latents = latents * self.scheduler.init_noise_sigma
607
+ return latents
608
+
609
+ @torch.no_grad()
610
+ @replace_example_docstring(EXAMPLE_DOC_STRING)
611
+ def __call__(
612
+ self,
613
+ prompt: Union[str, List[str]] = None,
614
+ negative_prompt: str = "",
615
+ num_inference_steps: int = 20,
616
+ timesteps: List[int] = None,
617
+ sigmas: List[float] = None,
618
+ guidance_scale: float = 4.5,
619
+ num_images_per_prompt: Optional[int] = 1,
620
+ height: Optional[int] = None,
621
+ width: Optional[int] = None,
622
+ eta: float = 0.0,
623
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
624
+ latents: Optional[torch.Tensor] = None,
625
+ prompt_embeds: Optional[torch.Tensor] = None,
626
+ prompt_attention_mask: Optional[torch.Tensor] = None,
627
+ negative_prompt_embeds: Optional[torch.Tensor] = None,
628
+ negative_prompt_attention_mask: Optional[torch.Tensor] = None,
629
+ output_type: Optional[str] = "pil",
630
+ return_dict: bool = True,
631
+ callback: Optional[Callable[[int, int, torch.Tensor], None]] = None,
632
+ callback_steps: int = 1,
633
+ clean_caption: bool = True,
634
+ use_resolution_binning: bool = True,
635
+ max_sequence_length: int = 300,
636
+ **kwargs,
637
+ ) -> Union[ImagePipelineOutput, Tuple]:
638
+ """
639
+ Function invoked when calling the pipeline for generation.
640
+
641
+ Args:
642
+ prompt (`str` or `List[str]`, *optional*):
643
+ The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
644
+ instead.
645
+ negative_prompt (`str` or `List[str]`, *optional*):
646
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass
647
+ `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
648
+ less than `1`).
649
+ num_inference_steps (`int`, *optional*, defaults to 100):
650
+ The number of denoising steps. More denoising steps usually lead to a higher quality image at the
651
+ expense of slower inference.
652
+ timesteps (`List[int]`, *optional*):
653
+ Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
654
+ in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
655
+ passed will be used. Must be in descending order.
656
+ sigmas (`List[float]`, *optional*):
657
+ Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
658
+ their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
659
+ will be used.
660
+ guidance_scale (`float`, *optional*, defaults to 4.5):
661
+ Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
662
+ `guidance_scale` is defined as `w` of equation 2. of [Imagen
663
+ Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
664
+ 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
665
+ usually at the expense of lower image quality.
666
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
667
+ The number of images to generate per prompt.
668
+ height (`int`, *optional*, defaults to self.unet.config.sample_size):
669
+ The height in pixels of the generated image.
670
+ width (`int`, *optional*, defaults to self.unet.config.sample_size):
671
+ The width in pixels of the generated image.
672
+ eta (`float`, *optional*, defaults to 0.0):
673
+ Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
674
+ [`schedulers.DDIMScheduler`], will be ignored for others.
675
+ generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
676
+ One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
677
+ to make generation deterministic.
678
+ latents (`torch.Tensor`, *optional*):
679
+ Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
680
+ generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
681
+ tensor will ge generated by sampling using the supplied random `generator`.
682
+ prompt_embeds (`torch.Tensor`, *optional*):
683
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
684
+ provided, text embeddings will be generated from `prompt` input argument.
685
+ prompt_attention_mask (`torch.Tensor`, *optional*): Pre-generated attention mask for text embeddings.
686
+ negative_prompt_embeds (`torch.Tensor`, *optional*):
687
+ Pre-generated negative text embeddings. For PixArt-Sigma this negative prompt should be "". If not
688
+ provided, negative_prompt_embeds will be generated from `negative_prompt` input argument.
689
+ negative_prompt_attention_mask (`torch.Tensor`, *optional*):
690
+ Pre-generated attention mask for negative text embeddings.
691
+ output_type (`str`, *optional*, defaults to `"pil"`):
692
+ The output format of the generate image. Choose between
693
+ [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
694
+ return_dict (`bool`, *optional*, defaults to `True`):
695
+ Whether or not to return a [`~pipelines.stable_diffusion.IFPipelineOutput`] instead of a plain tuple.
696
+ callback (`Callable`, *optional*):
697
+ A function that will be called every `callback_steps` steps during inference. The function will be
698
+ called with the following arguments: `callback(step: int, timestep: int, latents: torch.Tensor)`.
699
+ callback_steps (`int`, *optional*, defaults to 1):
700
+ The frequency at which the `callback` function will be called. If not specified, the callback will be
701
+ called at every step.
702
+ clean_caption (`bool`, *optional*, defaults to `True`):
703
+ Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to
704
+ be installed. If the dependencies are not installed, the embeddings will be created from the raw
705
+ prompt.
706
+ use_resolution_binning (`bool` defaults to `True`):
707
+ If set to `True`, the requested height and width are first mapped to the closest resolutions using
708
+ `ASPECT_RATIO_1024_BIN`. After the produced latents are decoded into images, they are resized back to
709
+ the requested resolution. Useful for generating non-square images.
710
+ max_sequence_length (`int` defaults to 300): Maximum sequence length to use with the `prompt`.
711
+
712
+ Examples:
713
+
714
+ Returns:
715
+ [`~pipelines.ImagePipelineOutput`] or `tuple`:
716
+ If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is
717
+ returned where the first element is a list with the generated images
718
+ """
719
+ # 1. Check inputs. Raise error if not correct
720
+ height = height or self.transformer.config.sample_size * self.vae_scale_factor
721
+ width = width or self.transformer.config.sample_size * self.vae_scale_factor
722
+ if use_resolution_binning:
723
+ if self.transformer.config.sample_size == 256:
724
+ aspect_ratio_bin = ASPECT_RATIO_2048_BIN
725
+ elif self.transformer.config.sample_size == 128:
726
+ aspect_ratio_bin = ASPECT_RATIO_1024_BIN
727
+ elif self.transformer.config.sample_size == 64:
728
+ aspect_ratio_bin = ASPECT_RATIO_512_BIN
729
+ elif self.transformer.config.sample_size == 32:
730
+ aspect_ratio_bin = ASPECT_RATIO_256_BIN
731
+ else:
732
+ raise ValueError("Invalid sample size")
733
+ orig_height, orig_width = height, width
734
+ height, width = self.image_processor.classify_height_width_bin(height, width, ratios=aspect_ratio_bin)
735
+
736
+ self.check_inputs(
737
+ prompt,
738
+ height,
739
+ width,
740
+ negative_prompt,
741
+ callback_steps,
742
+ prompt_embeds,
743
+ negative_prompt_embeds,
744
+ prompt_attention_mask,
745
+ negative_prompt_attention_mask,
746
+ )
747
+
748
+ # 2. Default height and width to transformer
749
+ if prompt is not None and isinstance(prompt, str):
750
+ batch_size = 1
751
+ elif prompt is not None and isinstance(prompt, list):
752
+ batch_size = len(prompt)
753
+ else:
754
+ batch_size = prompt_embeds.shape[0]
755
+
756
+ device = self._execution_device
757
+
758
+ # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
759
+ # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
760
+ # corresponds to doing no classifier free guidance.
761
+ do_classifier_free_guidance = guidance_scale > 1.0
762
+
763
+ # 3. Encode input prompt
764
+ (
765
+ prompt_embeds,
766
+ prompt_attention_mask,
767
+ negative_prompt_embeds,
768
+ negative_prompt_attention_mask,
769
+ ) = self.encode_prompt(
770
+ prompt,
771
+ do_classifier_free_guidance,
772
+ negative_prompt=negative_prompt,
773
+ num_images_per_prompt=num_images_per_prompt,
774
+ device=device,
775
+ prompt_embeds=prompt_embeds,
776
+ negative_prompt_embeds=negative_prompt_embeds,
777
+ prompt_attention_mask=prompt_attention_mask,
778
+ negative_prompt_attention_mask=negative_prompt_attention_mask,
779
+ clean_caption=clean_caption,
780
+ max_sequence_length=max_sequence_length,
781
+ )
782
+ if do_classifier_free_guidance:
783
+ prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
784
+ prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
785
+
786
+ # 4. Prepare timesteps
787
+ timesteps, num_inference_steps = retrieve_timesteps(
788
+ self.scheduler, num_inference_steps, device, timesteps, sigmas
789
+ )
790
+
791
+ # 5. Prepare latents.
792
+ latent_channels = self.transformer.config.in_channels
793
+ latents = self.prepare_latents(
794
+ batch_size * num_images_per_prompt,
795
+ latent_channels,
796
+ height,
797
+ width,
798
+ prompt_embeds.dtype,
799
+ device,
800
+ generator,
801
+ latents,
802
+ )
803
+
804
+ # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
805
+ extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
806
+
807
+ # 6.1 Prepare micro-conditions.
808
+ added_cond_kwargs = {"resolution": None, "aspect_ratio": None}
809
+
810
+ # 7. Denoising loop
811
+ num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
812
+
813
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
814
+ for i, t in enumerate(timesteps):
815
+ latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
816
+ latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
817
+
818
+ current_timestep = t
819
+ if not torch.is_tensor(current_timestep):
820
+ # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
821
+ # This would be a good case for the `match` statement (Python 3.10+)
822
+ is_mps = latent_model_input.device.type == "mps"
823
+ if isinstance(current_timestep, float):
824
+ dtype = torch.float32 if is_mps else torch.float64
825
+ else:
826
+ dtype = torch.int32 if is_mps else torch.int64
827
+ current_timestep = torch.tensor([current_timestep], dtype=dtype, device=latent_model_input.device)
828
+ elif len(current_timestep.shape) == 0:
829
+ current_timestep = current_timestep[None].to(latent_model_input.device)
830
+ # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
831
+ current_timestep = current_timestep.expand(latent_model_input.shape[0])
832
+
833
+ # predict noise model_output
834
+ noise_pred = self.transformer(
835
+ latent_model_input,
836
+ encoder_hidden_states=prompt_embeds,
837
+ encoder_attention_mask=prompt_attention_mask,
838
+ timestep=current_timestep,
839
+ added_cond_kwargs=added_cond_kwargs,
840
+ return_dict=False,
841
+ )[0]
842
+
843
+ # perform guidance
844
+ if do_classifier_free_guidance:
845
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
846
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
847
+
848
+ # learned sigma
849
+ if self.transformer.config.out_channels // 2 == latent_channels:
850
+ noise_pred = noise_pred.chunk(2, dim=1)[0]
851
+ else:
852
+ noise_pred = noise_pred
853
+
854
+ # compute previous image: x_t -> x_t-1
855
+ latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
856
+
857
+ # call the callback, if provided
858
+ if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
859
+ progress_bar.update()
860
+ if callback is not None and i % callback_steps == 0:
861
+ step_idx = i // getattr(self.scheduler, "order", 1)
862
+ callback(step_idx, t, latents)
863
+
864
+ if not output_type == "latent":
865
+ image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
866
+ if use_resolution_binning:
867
+ image = self.image_processor.resize_and_crop_tensor(image, orig_width, orig_height)
868
+ else:
869
+ image = latents
870
+
871
+ if not output_type == "latent":
872
+ image = self.image_processor.postprocess(image, output_type=output_type)
873
+
874
+ # Offload all models
875
+ self.maybe_free_model_hooks()
876
+
877
+ if not return_dict:
878
+ return (image,)
879
+
880
+ return ImagePipelineOutput(images=image)