diffusers 0.29.2__py3-none-any.whl → 0.30.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 (220) hide show
  1. diffusers/__init__.py +94 -3
  2. diffusers/commands/env.py +1 -5
  3. diffusers/configuration_utils.py +4 -9
  4. diffusers/dependency_versions_table.py +2 -2
  5. diffusers/image_processor.py +1 -2
  6. diffusers/loaders/__init__.py +17 -2
  7. diffusers/loaders/ip_adapter.py +10 -7
  8. diffusers/loaders/lora_base.py +752 -0
  9. diffusers/loaders/lora_pipeline.py +2222 -0
  10. diffusers/loaders/peft.py +213 -5
  11. diffusers/loaders/single_file.py +1 -12
  12. diffusers/loaders/single_file_model.py +31 -10
  13. diffusers/loaders/single_file_utils.py +262 -2
  14. diffusers/loaders/textual_inversion.py +1 -6
  15. diffusers/loaders/unet.py +23 -208
  16. diffusers/models/__init__.py +20 -0
  17. diffusers/models/activations.py +22 -0
  18. diffusers/models/attention.py +386 -7
  19. diffusers/models/attention_processor.py +1795 -629
  20. diffusers/models/autoencoders/__init__.py +2 -0
  21. diffusers/models/autoencoders/autoencoder_kl.py +14 -3
  22. diffusers/models/autoencoders/autoencoder_kl_cogvideox.py +1035 -0
  23. diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py +1 -1
  24. diffusers/models/autoencoders/autoencoder_oobleck.py +464 -0
  25. diffusers/models/autoencoders/autoencoder_tiny.py +1 -0
  26. diffusers/models/autoencoders/consistency_decoder_vae.py +1 -1
  27. diffusers/models/autoencoders/vq_model.py +4 -4
  28. diffusers/models/controlnet.py +2 -3
  29. diffusers/models/controlnet_hunyuan.py +401 -0
  30. diffusers/models/controlnet_sd3.py +11 -11
  31. diffusers/models/controlnet_sparsectrl.py +789 -0
  32. diffusers/models/controlnet_xs.py +40 -10
  33. diffusers/models/downsampling.py +68 -0
  34. diffusers/models/embeddings.py +319 -36
  35. diffusers/models/model_loading_utils.py +1 -3
  36. diffusers/models/modeling_flax_utils.py +1 -6
  37. diffusers/models/modeling_utils.py +4 -16
  38. diffusers/models/normalization.py +203 -12
  39. diffusers/models/transformers/__init__.py +6 -0
  40. diffusers/models/transformers/auraflow_transformer_2d.py +527 -0
  41. diffusers/models/transformers/cogvideox_transformer_3d.py +345 -0
  42. diffusers/models/transformers/hunyuan_transformer_2d.py +19 -15
  43. diffusers/models/transformers/latte_transformer_3d.py +327 -0
  44. diffusers/models/transformers/lumina_nextdit2d.py +340 -0
  45. diffusers/models/transformers/pixart_transformer_2d.py +102 -1
  46. diffusers/models/transformers/prior_transformer.py +1 -1
  47. diffusers/models/transformers/stable_audio_transformer.py +458 -0
  48. diffusers/models/transformers/transformer_flux.py +455 -0
  49. diffusers/models/transformers/transformer_sd3.py +18 -4
  50. diffusers/models/unets/unet_1d_blocks.py +1 -1
  51. diffusers/models/unets/unet_2d_condition.py +8 -1
  52. diffusers/models/unets/unet_3d_blocks.py +51 -920
  53. diffusers/models/unets/unet_3d_condition.py +4 -1
  54. diffusers/models/unets/unet_i2vgen_xl.py +4 -1
  55. diffusers/models/unets/unet_kandinsky3.py +1 -1
  56. diffusers/models/unets/unet_motion_model.py +1330 -84
  57. diffusers/models/unets/unet_spatio_temporal_condition.py +1 -1
  58. diffusers/models/unets/unet_stable_cascade.py +1 -3
  59. diffusers/models/unets/uvit_2d.py +1 -1
  60. diffusers/models/upsampling.py +64 -0
  61. diffusers/models/vq_model.py +8 -4
  62. diffusers/optimization.py +1 -1
  63. diffusers/pipelines/__init__.py +100 -3
  64. diffusers/pipelines/animatediff/__init__.py +4 -0
  65. diffusers/pipelines/animatediff/pipeline_animatediff.py +50 -40
  66. diffusers/pipelines/animatediff/pipeline_animatediff_controlnet.py +1076 -0
  67. diffusers/pipelines/animatediff/pipeline_animatediff_sdxl.py +17 -27
  68. diffusers/pipelines/animatediff/pipeline_animatediff_sparsectrl.py +1008 -0
  69. diffusers/pipelines/animatediff/pipeline_animatediff_video2video.py +51 -38
  70. diffusers/pipelines/audioldm2/modeling_audioldm2.py +1 -1
  71. diffusers/pipelines/audioldm2/pipeline_audioldm2.py +1 -0
  72. diffusers/pipelines/aura_flow/__init__.py +48 -0
  73. diffusers/pipelines/aura_flow/pipeline_aura_flow.py +591 -0
  74. diffusers/pipelines/auto_pipeline.py +97 -19
  75. diffusers/pipelines/cogvideo/__init__.py +48 -0
  76. diffusers/pipelines/cogvideo/pipeline_cogvideox.py +687 -0
  77. diffusers/pipelines/consistency_models/pipeline_consistency_models.py +1 -1
  78. diffusers/pipelines/controlnet/pipeline_controlnet.py +24 -30
  79. diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py +31 -30
  80. diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py +24 -153
  81. diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py +19 -28
  82. diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py +18 -28
  83. diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py +29 -32
  84. diffusers/pipelines/controlnet/pipeline_flax_controlnet.py +2 -2
  85. diffusers/pipelines/controlnet_hunyuandit/__init__.py +48 -0
  86. diffusers/pipelines/controlnet_hunyuandit/pipeline_hunyuandit_controlnet.py +1042 -0
  87. diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py +35 -0
  88. diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs.py +10 -6
  89. diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs_sd_xl.py +0 -4
  90. diffusers/pipelines/deepfloyd_if/pipeline_if.py +2 -2
  91. diffusers/pipelines/deepfloyd_if/pipeline_if_img2img.py +2 -2
  92. diffusers/pipelines/deepfloyd_if/pipeline_if_img2img_superresolution.py +2 -2
  93. diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting.py +2 -2
  94. diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting_superresolution.py +2 -2
  95. diffusers/pipelines/deepfloyd_if/pipeline_if_superresolution.py +2 -2
  96. diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion.py +11 -6
  97. diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion_img2img.py +11 -6
  98. diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_cycle_diffusion.py +6 -6
  99. diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_inpaint_legacy.py +6 -6
  100. diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_model_editing.py +10 -10
  101. diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_paradigms.py +10 -6
  102. diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_pix2pix_zero.py +3 -3
  103. diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py +1 -1
  104. diffusers/pipelines/flux/__init__.py +47 -0
  105. diffusers/pipelines/flux/pipeline_flux.py +749 -0
  106. diffusers/pipelines/flux/pipeline_output.py +21 -0
  107. diffusers/pipelines/free_init_utils.py +2 -0
  108. diffusers/pipelines/free_noise_utils.py +236 -0
  109. diffusers/pipelines/kandinsky3/pipeline_kandinsky3.py +2 -2
  110. diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py +2 -2
  111. diffusers/pipelines/kolors/__init__.py +54 -0
  112. diffusers/pipelines/kolors/pipeline_kolors.py +1070 -0
  113. diffusers/pipelines/kolors/pipeline_kolors_img2img.py +1247 -0
  114. diffusers/pipelines/kolors/pipeline_output.py +21 -0
  115. diffusers/pipelines/kolors/text_encoder.py +889 -0
  116. diffusers/pipelines/kolors/tokenizer.py +334 -0
  117. diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_img2img.py +30 -29
  118. diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_text2img.py +23 -29
  119. diffusers/pipelines/latte/__init__.py +48 -0
  120. diffusers/pipelines/latte/pipeline_latte.py +881 -0
  121. diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py +4 -4
  122. diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py +0 -4
  123. diffusers/pipelines/lumina/__init__.py +48 -0
  124. diffusers/pipelines/lumina/pipeline_lumina.py +897 -0
  125. diffusers/pipelines/pag/__init__.py +67 -0
  126. diffusers/pipelines/pag/pag_utils.py +237 -0
  127. diffusers/pipelines/pag/pipeline_pag_controlnet_sd.py +1329 -0
  128. diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl.py +1612 -0
  129. diffusers/pipelines/pag/pipeline_pag_hunyuandit.py +953 -0
  130. diffusers/pipelines/pag/pipeline_pag_kolors.py +1136 -0
  131. diffusers/pipelines/pag/pipeline_pag_pixart_sigma.py +872 -0
  132. diffusers/pipelines/pag/pipeline_pag_sd.py +1050 -0
  133. diffusers/pipelines/pag/pipeline_pag_sd_3.py +985 -0
  134. diffusers/pipelines/pag/pipeline_pag_sd_animatediff.py +862 -0
  135. diffusers/pipelines/pag/pipeline_pag_sd_xl.py +1333 -0
  136. diffusers/pipelines/pag/pipeline_pag_sd_xl_img2img.py +1529 -0
  137. diffusers/pipelines/pag/pipeline_pag_sd_xl_inpaint.py +1753 -0
  138. diffusers/pipelines/pia/pipeline_pia.py +30 -37
  139. diffusers/pipelines/pipeline_flax_utils.py +4 -9
  140. diffusers/pipelines/pipeline_loading_utils.py +0 -3
  141. diffusers/pipelines/pipeline_utils.py +2 -14
  142. diffusers/pipelines/semantic_stable_diffusion/pipeline_semantic_stable_diffusion.py +0 -1
  143. diffusers/pipelines/stable_audio/__init__.py +50 -0
  144. diffusers/pipelines/stable_audio/modeling_stable_audio.py +158 -0
  145. diffusers/pipelines/stable_audio/pipeline_stable_audio.py +745 -0
  146. diffusers/pipelines/stable_diffusion/convert_from_ckpt.py +2 -0
  147. diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py +1 -1
  148. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +23 -29
  149. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_depth2img.py +15 -8
  150. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py +30 -29
  151. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py +23 -152
  152. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py +8 -4
  153. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py +11 -11
  154. diffusers/pipelines/stable_diffusion/pipeline_stable_unclip.py +8 -6
  155. diffusers/pipelines/stable_diffusion/pipeline_stable_unclip_img2img.py +6 -6
  156. diffusers/pipelines/stable_diffusion_3/__init__.py +2 -0
  157. diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py +34 -3
  158. diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py +33 -7
  159. diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py +1201 -0
  160. diffusers/pipelines/stable_diffusion_attend_and_excite/pipeline_stable_diffusion_attend_and_excite.py +3 -3
  161. diffusers/pipelines/stable_diffusion_diffedit/pipeline_stable_diffusion_diffedit.py +6 -6
  162. diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen.py +5 -5
  163. diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen_text_image.py +5 -5
  164. diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_k_diffusion.py +6 -6
  165. diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_xl_k_diffusion.py +0 -4
  166. diffusers/pipelines/stable_diffusion_ldm3d/pipeline_stable_diffusion_ldm3d.py +23 -29
  167. diffusers/pipelines/stable_diffusion_panorama/pipeline_stable_diffusion_panorama.py +27 -29
  168. diffusers/pipelines/stable_diffusion_sag/pipeline_stable_diffusion_sag.py +3 -3
  169. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py +17 -27
  170. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py +26 -29
  171. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py +17 -145
  172. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_instruct_pix2pix.py +0 -4
  173. diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_adapter.py +6 -6
  174. diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_xl_adapter.py +18 -28
  175. diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth.py +8 -6
  176. diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth_img2img.py +8 -6
  177. diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero.py +6 -4
  178. diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero_sdxl.py +0 -4
  179. diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py +3 -3
  180. diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py +1 -1
  181. diffusers/pipelines/wuerstchen/pipeline_wuerstchen_prior.py +5 -4
  182. diffusers/schedulers/__init__.py +8 -0
  183. diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py +572 -0
  184. diffusers/schedulers/scheduling_ddim.py +1 -1
  185. diffusers/schedulers/scheduling_ddim_cogvideox.py +449 -0
  186. diffusers/schedulers/scheduling_ddpm.py +1 -1
  187. diffusers/schedulers/scheduling_ddpm_parallel.py +1 -1
  188. diffusers/schedulers/scheduling_deis_multistep.py +2 -2
  189. diffusers/schedulers/scheduling_dpm_cogvideox.py +489 -0
  190. diffusers/schedulers/scheduling_dpmsolver_multistep.py +1 -1
  191. diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py +1 -1
  192. diffusers/schedulers/scheduling_dpmsolver_singlestep.py +64 -19
  193. diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py +2 -2
  194. diffusers/schedulers/scheduling_flow_match_euler_discrete.py +63 -39
  195. diffusers/schedulers/scheduling_flow_match_heun_discrete.py +321 -0
  196. diffusers/schedulers/scheduling_ipndm.py +1 -1
  197. diffusers/schedulers/scheduling_unipc_multistep.py +1 -1
  198. diffusers/schedulers/scheduling_utils.py +1 -3
  199. diffusers/schedulers/scheduling_utils_flax.py +1 -3
  200. diffusers/training_utils.py +99 -14
  201. diffusers/utils/__init__.py +2 -2
  202. diffusers/utils/dummy_pt_objects.py +210 -0
  203. diffusers/utils/dummy_torch_and_torchsde_objects.py +15 -0
  204. diffusers/utils/dummy_torch_and_transformers_and_sentencepiece_objects.py +47 -0
  205. diffusers/utils/dummy_torch_and_transformers_objects.py +315 -0
  206. diffusers/utils/dynamic_modules_utils.py +1 -11
  207. diffusers/utils/export_utils.py +1 -4
  208. diffusers/utils/hub_utils.py +45 -42
  209. diffusers/utils/import_utils.py +19 -16
  210. diffusers/utils/loading_utils.py +76 -3
  211. diffusers/utils/testing_utils.py +11 -8
  212. {diffusers-0.29.2.dist-info → diffusers-0.30.0.dist-info}/METADATA +73 -83
  213. {diffusers-0.29.2.dist-info → diffusers-0.30.0.dist-info}/RECORD +217 -164
  214. {diffusers-0.29.2.dist-info → diffusers-0.30.0.dist-info}/WHEEL +1 -1
  215. diffusers/loaders/autoencoder.py +0 -146
  216. diffusers/loaders/controlnet.py +0 -136
  217. diffusers/loaders/lora.py +0 -1728
  218. {diffusers-0.29.2.dist-info → diffusers-0.30.0.dist-info}/LICENSE +0 -0
  219. {diffusers-0.29.2.dist-info → diffusers-0.30.0.dist-info}/entry_points.txt +0 -0
  220. {diffusers-0.29.2.dist-info → diffusers-0.30.0.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,897 @@
1
+ # Copyright 2024 Alpha-VLLM 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 math
18
+ import re
19
+ import urllib.parse as ul
20
+ from typing import List, Optional, Tuple, Union
21
+
22
+ import torch
23
+ from transformers import AutoModel, AutoTokenizer
24
+
25
+ from ...image_processor import VaeImageProcessor
26
+ from ...models import AutoencoderKL
27
+ from ...models.embeddings import get_2d_rotary_pos_embed_lumina
28
+ from ...models.transformers.lumina_nextdit2d import LuminaNextDiT2DModel
29
+ from ...schedulers import FlowMatchEulerDiscreteScheduler
30
+ from ...utils import (
31
+ BACKENDS_MAPPING,
32
+ is_bs4_available,
33
+ is_ftfy_available,
34
+ logging,
35
+ replace_example_docstring,
36
+ )
37
+ from ...utils.torch_utils import randn_tensor
38
+ from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput
39
+
40
+
41
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
42
+
43
+ if is_bs4_available():
44
+ from bs4 import BeautifulSoup
45
+
46
+ if is_ftfy_available():
47
+ import ftfy
48
+
49
+ EXAMPLE_DOC_STRING = """
50
+ Examples:
51
+ ```py
52
+ >>> import torch
53
+ >>> from diffusers import LuminaText2ImgPipeline
54
+
55
+ >>> pipe = LuminaText2ImgPipeline.from_pretrained(
56
+ ... "Alpha-VLLM/Lumina-Next-SFT-diffusers", torch_dtype=torch.bfloat16
57
+ ... ).cuda()
58
+ >>> # Enable memory optimizations.
59
+ >>> pipe.enable_model_cpu_offload()
60
+
61
+ >>> prompt = "Upper body of a young woman in a Victorian-era outfit with brass goggles and leather straps. Background shows an industrial revolution cityscape with smoky skies and tall, metal structures"
62
+ >>> image = pipe(prompt).images[0]
63
+ ```
64
+ """
65
+
66
+
67
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
68
+ def retrieve_timesteps(
69
+ scheduler,
70
+ num_inference_steps: Optional[int] = None,
71
+ device: Optional[Union[str, torch.device]] = None,
72
+ timesteps: Optional[List[int]] = None,
73
+ sigmas: Optional[List[float]] = None,
74
+ **kwargs,
75
+ ):
76
+ """
77
+ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
78
+ custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
79
+
80
+ Args:
81
+ scheduler (`SchedulerMixin`):
82
+ The scheduler to get timesteps from.
83
+ num_inference_steps (`int`):
84
+ The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
85
+ must be `None`.
86
+ device (`str` or `torch.device`, *optional*):
87
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
88
+ timesteps (`List[int]`, *optional*):
89
+ Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
90
+ `num_inference_steps` and `sigmas` must be `None`.
91
+ sigmas (`List[float]`, *optional*):
92
+ Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
93
+ `num_inference_steps` and `timesteps` must be `None`.
94
+
95
+ Returns:
96
+ `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
97
+ second element is the number of inference steps.
98
+ """
99
+ if timesteps is not None and sigmas is not None:
100
+ raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
101
+ if timesteps is not None:
102
+ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
103
+ if not accepts_timesteps:
104
+ raise ValueError(
105
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
106
+ f" timestep schedules. Please check whether you are using the correct scheduler."
107
+ )
108
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
109
+ timesteps = scheduler.timesteps
110
+ num_inference_steps = len(timesteps)
111
+ elif sigmas is not None:
112
+ accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
113
+ if not accept_sigmas:
114
+ raise ValueError(
115
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
116
+ f" sigmas schedules. Please check whether you are using the correct scheduler."
117
+ )
118
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
119
+ timesteps = scheduler.timesteps
120
+ num_inference_steps = len(timesteps)
121
+ else:
122
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
123
+ timesteps = scheduler.timesteps
124
+ return timesteps, num_inference_steps
125
+
126
+
127
+ class LuminaText2ImgPipeline(DiffusionPipeline):
128
+ r"""
129
+ Pipeline for text-to-image generation using Lumina-T2I.
130
+
131
+ This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
132
+ library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
133
+
134
+ Args:
135
+ vae ([`AutoencoderKL`]):
136
+ Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
137
+ text_encoder ([`AutoModel`]):
138
+ Frozen text-encoder. Lumina-T2I uses
139
+ [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.AutoModel), specifically the
140
+ [t5-v1_1-xxl](https://huggingface.co/Alpha-VLLM/tree/main/t5-v1_1-xxl) variant.
141
+ tokenizer (`AutoModel`):
142
+ Tokenizer of class
143
+ [AutoModel](https://huggingface.co/docs/transformers/model_doc/t5#transformers.AutoModel).
144
+ transformer ([`Transformer2DModel`]):
145
+ A text conditioned `Transformer2DModel` to denoise the encoded image latents.
146
+ scheduler ([`SchedulerMixin`]):
147
+ A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
148
+ """
149
+
150
+ bad_punct_regex = re.compile(
151
+ r"["
152
+ + "#®•©™&@·º½¾¿¡§~"
153
+ + r"\)"
154
+ + r"\("
155
+ + r"\]"
156
+ + r"\["
157
+ + r"\}"
158
+ + r"\{"
159
+ + r"\|"
160
+ + "\\"
161
+ + r"\/"
162
+ + r"\*"
163
+ + r"]{1,}"
164
+ ) # noqa
165
+
166
+ _optional_components = []
167
+ model_cpu_offload_seq = "text_encoder->transformer->vae"
168
+
169
+ def __init__(
170
+ self,
171
+ transformer: LuminaNextDiT2DModel,
172
+ scheduler: FlowMatchEulerDiscreteScheduler,
173
+ vae: AutoencoderKL,
174
+ text_encoder: AutoModel,
175
+ tokenizer: AutoTokenizer,
176
+ ):
177
+ super().__init__()
178
+
179
+ self.register_modules(
180
+ vae=vae,
181
+ text_encoder=text_encoder,
182
+ tokenizer=tokenizer,
183
+ transformer=transformer,
184
+ scheduler=scheduler,
185
+ )
186
+ self.vae_scale_factor = 8
187
+ self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
188
+ self.max_sequence_length = 256
189
+ self.default_sample_size = (
190
+ self.transformer.config.sample_size
191
+ if hasattr(self, "transformer") and self.transformer is not None
192
+ else 128
193
+ )
194
+ self.default_image_size = self.default_sample_size * self.vae_scale_factor
195
+
196
+ def _get_gemma_prompt_embeds(
197
+ self,
198
+ prompt: Union[str, List[str]],
199
+ num_images_per_prompt: int = 1,
200
+ device: Optional[torch.device] = None,
201
+ clean_caption: Optional[bool] = False,
202
+ max_length: Optional[int] = None,
203
+ ):
204
+ device = device or self._execution_device
205
+ prompt = [prompt] if isinstance(prompt, str) else prompt
206
+ batch_size = len(prompt)
207
+
208
+ prompt = self._text_preprocessing(prompt, clean_caption=clean_caption)
209
+ text_inputs = self.tokenizer(
210
+ prompt,
211
+ pad_to_multiple_of=8,
212
+ max_length=self.max_sequence_length,
213
+ truncation=True,
214
+ padding=True,
215
+ return_tensors="pt",
216
+ )
217
+ text_input_ids = text_inputs.input_ids.to(device)
218
+ untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids.to(device)
219
+
220
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
221
+ removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.max_sequence_length - 1 : -1])
222
+ logger.warning(
223
+ "The following part of your input was truncated because Gemma can only handle sequences up to"
224
+ f" {self.max_sequence_length} tokens: {removed_text}"
225
+ )
226
+
227
+ prompt_attention_mask = text_inputs.attention_mask.to(device)
228
+ prompt_embeds = self.text_encoder(
229
+ text_input_ids, attention_mask=prompt_attention_mask, output_hidden_states=True
230
+ )
231
+ prompt_embeds = prompt_embeds.hidden_states[-2]
232
+
233
+ if self.text_encoder is not None:
234
+ dtype = self.text_encoder.dtype
235
+ elif self.transformer is not None:
236
+ dtype = self.transformer.dtype
237
+ else:
238
+ dtype = None
239
+
240
+ prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
241
+
242
+ _, seq_len, _ = prompt_embeds.shape
243
+ # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
244
+ prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
245
+ prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
246
+ prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1)
247
+ prompt_attention_mask = prompt_attention_mask.view(batch_size * num_images_per_prompt, -1)
248
+
249
+ return prompt_embeds, prompt_attention_mask
250
+
251
+ # Adapted from diffusers.pipelines.deepfloyd_if.pipeline_if.encode_prompt
252
+ def encode_prompt(
253
+ self,
254
+ prompt: Union[str, List[str]],
255
+ do_classifier_free_guidance: bool = True,
256
+ negative_prompt: Union[str, List[str]] = None,
257
+ num_images_per_prompt: int = 1,
258
+ device: Optional[torch.device] = None,
259
+ prompt_embeds: Optional[torch.Tensor] = None,
260
+ negative_prompt_embeds: Optional[torch.Tensor] = None,
261
+ prompt_attention_mask: Optional[torch.Tensor] = None,
262
+ negative_prompt_attention_mask: Optional[torch.Tensor] = None,
263
+ clean_caption: bool = False,
264
+ **kwargs,
265
+ ):
266
+ r"""
267
+ Encodes the prompt into text encoder hidden states.
268
+
269
+ Args:
270
+ prompt (`str` or `List[str]`, *optional*):
271
+ prompt to be encoded
272
+ negative_prompt (`str` or `List[str]`, *optional*):
273
+ The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`
274
+ instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For
275
+ Lumina-T2I, this should be "".
276
+ do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
277
+ whether to use classifier free guidance or not
278
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
279
+ number of images that should be generated per prompt
280
+ device: (`torch.device`, *optional*):
281
+ torch device to place the resulting embeddings on
282
+ prompt_embeds (`torch.Tensor`, *optional*):
283
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
284
+ provided, text embeddings will be generated from `prompt` input argument.
285
+ negative_prompt_embeds (`torch.Tensor`, *optional*):
286
+ Pre-generated negative text embeddings. For Lumina-T2I, it's should be the embeddings of the "" string.
287
+ clean_caption (`bool`, defaults to `False`):
288
+ If `True`, the function will preprocess and clean the provided caption before encoding.
289
+ max_sequence_length (`int`, defaults to 256): Maximum sequence length to use for the prompt.
290
+ """
291
+ if device is None:
292
+ device = self._execution_device
293
+
294
+ prompt = [prompt] if isinstance(prompt, str) else prompt
295
+ if prompt is not None:
296
+ batch_size = len(prompt)
297
+ else:
298
+ batch_size = prompt_embeds.shape[0]
299
+
300
+ if prompt_embeds is None:
301
+ prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds(
302
+ prompt=prompt,
303
+ num_images_per_prompt=num_images_per_prompt,
304
+ device=device,
305
+ clean_caption=clean_caption,
306
+ )
307
+
308
+ # Get negative embeddings for classifier free guidance
309
+ if do_classifier_free_guidance and negative_prompt_embeds is None:
310
+ negative_prompt = negative_prompt if negative_prompt is not None else ""
311
+
312
+ # Normalize str to list
313
+ negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
314
+
315
+ if prompt is not None and type(prompt) is not type(negative_prompt):
316
+ raise TypeError(
317
+ f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
318
+ f" {type(prompt)}."
319
+ )
320
+ elif isinstance(negative_prompt, str):
321
+ negative_prompt = [negative_prompt]
322
+ elif batch_size != len(negative_prompt):
323
+ raise ValueError(
324
+ f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
325
+ f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
326
+ " the batch size of `prompt`."
327
+ )
328
+ # Padding negative prompt to the same length with prompt
329
+ prompt_max_length = prompt_embeds.shape[1]
330
+ negative_text_inputs = self.tokenizer(
331
+ negative_prompt,
332
+ padding="max_length",
333
+ max_length=prompt_max_length,
334
+ truncation=True,
335
+ return_tensors="pt",
336
+ )
337
+ negative_text_input_ids = negative_text_inputs.input_ids.to(device)
338
+ negative_prompt_attention_mask = negative_text_inputs.attention_mask.to(device)
339
+ # Get the negative prompt embeddings
340
+ negative_prompt_embeds = self.text_encoder(
341
+ negative_text_input_ids,
342
+ attention_mask=negative_prompt_attention_mask,
343
+ output_hidden_states=True,
344
+ )
345
+
346
+ negative_dtype = self.text_encoder.dtype
347
+ negative_prompt_embeds = negative_prompt_embeds.hidden_states[-2]
348
+ _, seq_len, _ = negative_prompt_embeds.shape
349
+
350
+ negative_prompt_embeds = negative_prompt_embeds.to(dtype=negative_dtype, device=device)
351
+ # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
352
+ negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
353
+ negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
354
+ negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 1)
355
+ negative_prompt_attention_mask = negative_prompt_attention_mask.view(
356
+ batch_size * num_images_per_prompt, -1
357
+ )
358
+
359
+ return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask
360
+
361
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
362
+ def prepare_extra_step_kwargs(self, generator, eta):
363
+ # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
364
+ # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
365
+ # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
366
+ # and should be between [0, 1]
367
+
368
+ accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
369
+ extra_step_kwargs = {}
370
+ if accepts_eta:
371
+ extra_step_kwargs["eta"] = eta
372
+
373
+ # check if the scheduler accepts generator
374
+ accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
375
+ if accepts_generator:
376
+ extra_step_kwargs["generator"] = generator
377
+ return extra_step_kwargs
378
+
379
+ def check_inputs(
380
+ self,
381
+ prompt,
382
+ height,
383
+ width,
384
+ negative_prompt,
385
+ prompt_embeds=None,
386
+ negative_prompt_embeds=None,
387
+ prompt_attention_mask=None,
388
+ negative_prompt_attention_mask=None,
389
+ ):
390
+ if height % 8 != 0 or width % 8 != 0:
391
+ raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
392
+
393
+ if prompt is not None and prompt_embeds is not None:
394
+ raise ValueError(
395
+ f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
396
+ " only forward one of the two."
397
+ )
398
+ elif prompt is None and prompt_embeds is None:
399
+ raise ValueError(
400
+ "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
401
+ )
402
+ elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
403
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
404
+
405
+ if prompt is not None and negative_prompt_embeds is not None:
406
+ raise ValueError(
407
+ f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:"
408
+ f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
409
+ )
410
+
411
+ if negative_prompt is not None and negative_prompt_embeds is not None:
412
+ raise ValueError(
413
+ f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
414
+ f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
415
+ )
416
+
417
+ if prompt_embeds is not None and prompt_attention_mask is None:
418
+ raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.")
419
+
420
+ if negative_prompt_embeds is not None and negative_prompt_attention_mask is None:
421
+ raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.")
422
+
423
+ if prompt_embeds is not None and negative_prompt_embeds is not None:
424
+ if prompt_embeds.shape != negative_prompt_embeds.shape:
425
+ raise ValueError(
426
+ "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
427
+ f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
428
+ f" {negative_prompt_embeds.shape}."
429
+ )
430
+ if prompt_attention_mask.shape != negative_prompt_attention_mask.shape:
431
+ raise ValueError(
432
+ "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but"
433
+ f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`"
434
+ f" {negative_prompt_attention_mask.shape}."
435
+ )
436
+
437
+ # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing
438
+ def _text_preprocessing(self, text, clean_caption=False):
439
+ if clean_caption and not is_bs4_available():
440
+ logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`"))
441
+ logger.warning("Setting `clean_caption` to False...")
442
+ clean_caption = False
443
+
444
+ if clean_caption and not is_ftfy_available():
445
+ logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`"))
446
+ logger.warning("Setting `clean_caption` to False...")
447
+ clean_caption = False
448
+
449
+ if not isinstance(text, (tuple, list)):
450
+ text = [text]
451
+
452
+ def process(text: str):
453
+ if clean_caption:
454
+ text = self._clean_caption(text)
455
+ text = self._clean_caption(text)
456
+ else:
457
+ text = text.lower().strip()
458
+ return text
459
+
460
+ return [process(t) for t in text]
461
+
462
+ # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption
463
+ def _clean_caption(self, caption):
464
+ caption = str(caption)
465
+ caption = ul.unquote_plus(caption)
466
+ caption = caption.strip().lower()
467
+ caption = re.sub("<person>", "person", caption)
468
+ # urls:
469
+ caption = re.sub(
470
+ r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa
471
+ "",
472
+ caption,
473
+ ) # regex for urls
474
+ caption = re.sub(
475
+ r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa
476
+ "",
477
+ caption,
478
+ ) # regex for urls
479
+ # html:
480
+ caption = BeautifulSoup(caption, features="html.parser").text
481
+
482
+ # @<nickname>
483
+ caption = re.sub(r"@[\w\d]+\b", "", caption)
484
+
485
+ # 31C0—31EF CJK Strokes
486
+ # 31F0—31FF Katakana Phonetic Extensions
487
+ # 3200—32FF Enclosed CJK Letters and Months
488
+ # 3300—33FF CJK Compatibility
489
+ # 3400—4DBF CJK Unified Ideographs Extension A
490
+ # 4DC0—4DFF Yijing Hexagram Symbols
491
+ # 4E00—9FFF CJK Unified Ideographs
492
+ caption = re.sub(r"[\u31c0-\u31ef]+", "", caption)
493
+ caption = re.sub(r"[\u31f0-\u31ff]+", "", caption)
494
+ caption = re.sub(r"[\u3200-\u32ff]+", "", caption)
495
+ caption = re.sub(r"[\u3300-\u33ff]+", "", caption)
496
+ caption = re.sub(r"[\u3400-\u4dbf]+", "", caption)
497
+ caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption)
498
+ caption = re.sub(r"[\u4e00-\u9fff]+", "", caption)
499
+ #######################################################
500
+
501
+ # все виды тире / all types of dash --> "-"
502
+ caption = re.sub(
503
+ r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+", # noqa
504
+ "-",
505
+ caption,
506
+ )
507
+
508
+ # кавычки к одному стандарту
509
+ caption = re.sub(r"[`´«»“”¨]", '"', caption)
510
+ caption = re.sub(r"[‘’]", "'", caption)
511
+
512
+ # &quot;
513
+ caption = re.sub(r"&quot;?", "", caption)
514
+ # &amp
515
+ caption = re.sub(r"&amp", "", caption)
516
+
517
+ # ip adresses:
518
+ caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption)
519
+
520
+ # article ids:
521
+ caption = re.sub(r"\d:\d\d\s+$", "", caption)
522
+
523
+ # \n
524
+ caption = re.sub(r"\\n", " ", caption)
525
+
526
+ # "#123"
527
+ caption = re.sub(r"#\d{1,3}\b", "", caption)
528
+ # "#12345.."
529
+ caption = re.sub(r"#\d{5,}\b", "", caption)
530
+ # "123456.."
531
+ caption = re.sub(r"\b\d{6,}\b", "", caption)
532
+ # filenames:
533
+ caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption)
534
+
535
+ #
536
+ caption = re.sub(r"[\"\']{2,}", r'"', caption) # """AUSVERKAUFT"""
537
+ caption = re.sub(r"[\.]{2,}", r" ", caption) # """AUSVERKAUFT"""
538
+
539
+ caption = re.sub(self.bad_punct_regex, r" ", caption) # ***AUSVERKAUFT***, #AUSVERKAUFT
540
+ caption = re.sub(r"\s+\.\s+", r" ", caption) # " . "
541
+
542
+ # this-is-my-cute-cat / this_is_my_cute_cat
543
+ regex2 = re.compile(r"(?:\-|\_)")
544
+ if len(re.findall(regex2, caption)) > 3:
545
+ caption = re.sub(regex2, " ", caption)
546
+
547
+ caption = ftfy.fix_text(caption)
548
+ caption = html.unescape(html.unescape(caption))
549
+
550
+ caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption) # jc6640
551
+ caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption) # jc6640vc
552
+ caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption) # 6640vc231
553
+
554
+ caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption)
555
+ caption = re.sub(r"(free\s)?download(\sfree)?", "", caption)
556
+ caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption)
557
+ caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption)
558
+ caption = re.sub(r"\bpage\s+\d+\b", "", caption)
559
+
560
+ caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption) # j2d1a2a...
561
+
562
+ caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption)
563
+
564
+ caption = re.sub(r"\b\s+\:\s+", r": ", caption)
565
+ caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption)
566
+ caption = re.sub(r"\s+", " ", caption)
567
+
568
+ caption.strip()
569
+
570
+ caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption)
571
+ caption = re.sub(r"^[\'\_,\-\:;]", r"", caption)
572
+ caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption)
573
+ caption = re.sub(r"^\.\S+$", "", caption)
574
+
575
+ return caption.strip()
576
+
577
+ def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
578
+ shape = (
579
+ batch_size,
580
+ num_channels_latents,
581
+ int(height) // self.vae_scale_factor,
582
+ int(width) // self.vae_scale_factor,
583
+ )
584
+ if isinstance(generator, list) and len(generator) != batch_size:
585
+ raise ValueError(
586
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
587
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
588
+ )
589
+
590
+ if latents is None:
591
+ latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
592
+ else:
593
+ latents = latents.to(device)
594
+
595
+ return latents
596
+
597
+ @property
598
+ def guidance_scale(self):
599
+ return self._guidance_scale
600
+
601
+ # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
602
+ # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
603
+ # corresponds to doing no classifier free guidance.
604
+ @property
605
+ def do_classifier_free_guidance(self):
606
+ return self._guidance_scale > 1
607
+
608
+ @property
609
+ def num_timesteps(self):
610
+ return self._num_timesteps
611
+
612
+ @torch.no_grad()
613
+ @replace_example_docstring(EXAMPLE_DOC_STRING)
614
+ def __call__(
615
+ self,
616
+ prompt: Union[str, List[str]] = None,
617
+ width: Optional[int] = None,
618
+ height: Optional[int] = None,
619
+ num_inference_steps: int = 30,
620
+ timesteps: List[int] = None,
621
+ guidance_scale: float = 4.0,
622
+ negative_prompt: Union[str, List[str]] = None,
623
+ sigmas: List[float] = None,
624
+ num_images_per_prompt: Optional[int] = 1,
625
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
626
+ latents: Optional[torch.Tensor] = None,
627
+ prompt_embeds: Optional[torch.Tensor] = None,
628
+ negative_prompt_embeds: Optional[torch.Tensor] = None,
629
+ prompt_attention_mask: Optional[torch.Tensor] = None,
630
+ negative_prompt_attention_mask: Optional[torch.Tensor] = None,
631
+ output_type: Optional[str] = "pil",
632
+ return_dict: bool = True,
633
+ clean_caption: bool = True,
634
+ max_sequence_length: int = 256,
635
+ scaling_watershed: Optional[float] = 1.0,
636
+ proportional_attn: Optional[bool] = True,
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 30):
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.0):
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 Lumina-T2I 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
+ clean_caption (`bool`, *optional*, defaults to `True`):
697
+ Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to
698
+ be installed. If the dependencies are not installed, the embeddings will be created from the raw
699
+ prompt.
700
+ max_sequence_length (`int` defaults to 120):
701
+ Maximum sequence length to use with the `prompt`.
702
+ callback_on_step_end (`Callable`, *optional*):
703
+ A function that calls at the end of each denoising steps during the inference. The function is called
704
+ with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
705
+ callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
706
+ `callback_on_step_end_tensor_inputs`.
707
+ callback_on_step_end_tensor_inputs (`List`, *optional*):
708
+ The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
709
+ will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
710
+ `._callback_tensor_inputs` attribute of your pipeline class.
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
+ height = height or self.default_sample_size * self.vae_scale_factor
720
+ width = width or self.default_sample_size * self.vae_scale_factor
721
+
722
+ # 1. Check inputs. Raise error if not correct
723
+ self.check_inputs(
724
+ prompt,
725
+ height,
726
+ width,
727
+ negative_prompt,
728
+ prompt_embeds=prompt_embeds,
729
+ negative_prompt_embeds=negative_prompt_embeds,
730
+ prompt_attention_mask=prompt_attention_mask,
731
+ negative_prompt_attention_mask=negative_prompt_attention_mask,
732
+ )
733
+ cross_attention_kwargs = {}
734
+
735
+ # 2. Define call parameters
736
+ if prompt is not None and isinstance(prompt, str):
737
+ batch_size = 1
738
+ elif prompt is not None and isinstance(prompt, list):
739
+ batch_size = len(prompt)
740
+ else:
741
+ batch_size = prompt_embeds.shape[0]
742
+
743
+ if proportional_attn:
744
+ cross_attention_kwargs["base_sequence_length"] = (self.default_image_size // 16) ** 2
745
+
746
+ scaling_factor = math.sqrt(width * height / self.default_image_size**2)
747
+
748
+ device = self._execution_device
749
+
750
+ # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
751
+ # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
752
+ # corresponds to doing no classifier free guidance.
753
+ do_classifier_free_guidance = guidance_scale > 1.0
754
+
755
+ # 3. Encode input prompt
756
+ (
757
+ prompt_embeds,
758
+ prompt_attention_mask,
759
+ negative_prompt_embeds,
760
+ negative_prompt_attention_mask,
761
+ ) = self.encode_prompt(
762
+ prompt,
763
+ do_classifier_free_guidance,
764
+ negative_prompt=negative_prompt,
765
+ num_images_per_prompt=num_images_per_prompt,
766
+ device=device,
767
+ prompt_embeds=prompt_embeds,
768
+ negative_prompt_embeds=negative_prompt_embeds,
769
+ prompt_attention_mask=prompt_attention_mask,
770
+ negative_prompt_attention_mask=negative_prompt_attention_mask,
771
+ clean_caption=clean_caption,
772
+ max_sequence_length=max_sequence_length,
773
+ )
774
+ if do_classifier_free_guidance:
775
+ prompt_embeds = torch.cat([prompt_embeds, negative_prompt_embeds], dim=0)
776
+ prompt_attention_mask = torch.cat([prompt_attention_mask, negative_prompt_attention_mask], dim=0)
777
+
778
+ # 4. Prepare timesteps
779
+ timesteps, num_inference_steps = retrieve_timesteps(
780
+ self.scheduler, num_inference_steps, device, timesteps, sigmas
781
+ )
782
+
783
+ # 5. Prepare latents.
784
+ latent_channels = self.transformer.config.in_channels
785
+ latents = self.prepare_latents(
786
+ batch_size * num_images_per_prompt,
787
+ latent_channels,
788
+ height,
789
+ width,
790
+ prompt_embeds.dtype,
791
+ device,
792
+ generator,
793
+ latents,
794
+ )
795
+
796
+ # 6. Denoising loop
797
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
798
+ for i, t in enumerate(timesteps):
799
+ # expand the latents if we are doing classifier free guidance
800
+ latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
801
+
802
+ current_timestep = t
803
+ if not torch.is_tensor(current_timestep):
804
+ # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
805
+ # This would be a good case for the `match` statement (Python 3.10+)
806
+ is_mps = latent_model_input.device.type == "mps"
807
+ if isinstance(current_timestep, float):
808
+ dtype = torch.float32 if is_mps else torch.float64
809
+ else:
810
+ dtype = torch.int32 if is_mps else torch.int64
811
+ current_timestep = torch.tensor(
812
+ [current_timestep],
813
+ dtype=dtype,
814
+ device=latent_model_input.device,
815
+ )
816
+ elif len(current_timestep.shape) == 0:
817
+ current_timestep = current_timestep[None].to(latent_model_input.device)
818
+ # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
819
+ current_timestep = current_timestep.expand(latent_model_input.shape[0])
820
+
821
+ # reverse the timestep since Lumina uses t=0 as the noise and t=1 as the image
822
+ current_timestep = 1 - current_timestep / self.scheduler.config.num_train_timesteps
823
+
824
+ # prepare image_rotary_emb for positional encoding
825
+ # dynamic scaling_factor for different resolution.
826
+ # NOTE: For `Time-aware` denosing mechanism from Lumina-Next
827
+ # https://arxiv.org/abs/2406.18583, Sec 2.3
828
+ # NOTE: We should compute different image_rotary_emb with different timestep.
829
+ if current_timestep[0] < scaling_watershed:
830
+ linear_factor = scaling_factor
831
+ ntk_factor = 1.0
832
+ else:
833
+ linear_factor = 1.0
834
+ ntk_factor = scaling_factor
835
+ image_rotary_emb = get_2d_rotary_pos_embed_lumina(
836
+ self.transformer.head_dim,
837
+ 384,
838
+ 384,
839
+ linear_factor=linear_factor,
840
+ ntk_factor=ntk_factor,
841
+ )
842
+
843
+ noise_pred = self.transformer(
844
+ hidden_states=latent_model_input,
845
+ timestep=current_timestep,
846
+ encoder_hidden_states=prompt_embeds,
847
+ encoder_mask=prompt_attention_mask,
848
+ image_rotary_emb=image_rotary_emb,
849
+ cross_attention_kwargs=cross_attention_kwargs,
850
+ return_dict=False,
851
+ )[0]
852
+ noise_pred = noise_pred.chunk(2, dim=1)[0]
853
+
854
+ # perform guidance scale
855
+ # NOTE: For exact reproducibility reasons, we apply classifier-free guidance on only
856
+ # three channels by default. The standard approach to cfg applies it to all channels.
857
+ # This can be done by uncommenting the following line and commenting-out the line following that.
858
+ # eps, rest = model_out[:, :self.in_channels], model_out[:, self.in_channels:]
859
+ if do_classifier_free_guidance:
860
+ noise_pred_eps, noise_pred_rest = noise_pred[:, :3], noise_pred[:, 3:]
861
+ noise_pred_cond_eps, noise_pred_uncond_eps = torch.split(
862
+ noise_pred_eps, len(noise_pred_eps) // 2, dim=0
863
+ )
864
+ noise_pred_half = noise_pred_uncond_eps + guidance_scale * (
865
+ noise_pred_cond_eps - noise_pred_uncond_eps
866
+ )
867
+ noise_pred_eps = torch.cat([noise_pred_half, noise_pred_half], dim=0)
868
+
869
+ noise_pred = torch.cat([noise_pred_eps, noise_pred_rest], dim=1)
870
+ noise_pred, _ = noise_pred.chunk(2, dim=0)
871
+
872
+ # compute the previous noisy sample x_t -> x_t-1
873
+ latents_dtype = latents.dtype
874
+ noise_pred = -noise_pred
875
+ latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
876
+
877
+ if latents.dtype != latents_dtype:
878
+ if torch.backends.mps.is_available():
879
+ # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
880
+ latents = latents.to(latents_dtype)
881
+
882
+ progress_bar.update()
883
+
884
+ if not output_type == "latent":
885
+ latents = latents / self.vae.config.scaling_factor
886
+ image = self.vae.decode(latents, return_dict=False)[0]
887
+ image = self.image_processor.postprocess(image, output_type=output_type)
888
+ else:
889
+ image = latents
890
+
891
+ # Offload all models
892
+ self.maybe_free_model_hooks()
893
+
894
+ if not return_dict:
895
+ return (image,)
896
+
897
+ return ImagePipelineOutput(images=image)