diffusers 0.24.0__py3-none-any.whl → 0.25.1__py3-none-any.whl

Sign up to get free protection for your applications and to get access to all the features.
Files changed (174) hide show
  1. diffusers/__init__.py +11 -1
  2. diffusers/commands/fp16_safetensors.py +10 -11
  3. diffusers/configuration_utils.py +12 -8
  4. diffusers/dependency_versions_table.py +3 -2
  5. diffusers/experimental/rl/value_guided_sampling.py +1 -1
  6. diffusers/image_processor.py +286 -46
  7. diffusers/loaders/ip_adapter.py +11 -9
  8. diffusers/loaders/lora.py +198 -60
  9. diffusers/loaders/single_file.py +24 -18
  10. diffusers/loaders/textual_inversion.py +10 -14
  11. diffusers/loaders/unet.py +130 -37
  12. diffusers/models/__init__.py +18 -12
  13. diffusers/models/activations.py +9 -6
  14. diffusers/models/attention.py +137 -16
  15. diffusers/models/attention_processor.py +133 -46
  16. diffusers/models/autoencoders/__init__.py +5 -0
  17. diffusers/models/{autoencoder_asym_kl.py → autoencoders/autoencoder_asym_kl.py} +4 -4
  18. diffusers/models/{autoencoder_kl.py → autoencoders/autoencoder_kl.py} +45 -6
  19. diffusers/models/{autoencoder_kl_temporal_decoder.py → autoencoders/autoencoder_kl_temporal_decoder.py} +8 -8
  20. diffusers/models/{autoencoder_tiny.py → autoencoders/autoencoder_tiny.py} +4 -4
  21. diffusers/models/{consistency_decoder_vae.py → autoencoders/consistency_decoder_vae.py} +14 -14
  22. diffusers/models/{vae.py → autoencoders/vae.py} +9 -5
  23. diffusers/models/downsampling.py +338 -0
  24. diffusers/models/embeddings.py +112 -29
  25. diffusers/models/modeling_flax_utils.py +12 -7
  26. diffusers/models/modeling_utils.py +10 -10
  27. diffusers/models/normalization.py +108 -2
  28. diffusers/models/resnet.py +15 -699
  29. diffusers/models/transformer_2d.py +2 -2
  30. diffusers/models/unet_2d_condition.py +37 -0
  31. diffusers/models/{unet_kandi3.py → unet_kandinsky3.py} +105 -159
  32. diffusers/models/upsampling.py +454 -0
  33. diffusers/models/uvit_2d.py +471 -0
  34. diffusers/models/vq_model.py +9 -2
  35. diffusers/pipelines/__init__.py +81 -73
  36. diffusers/pipelines/amused/__init__.py +62 -0
  37. diffusers/pipelines/amused/pipeline_amused.py +328 -0
  38. diffusers/pipelines/amused/pipeline_amused_img2img.py +347 -0
  39. diffusers/pipelines/amused/pipeline_amused_inpaint.py +378 -0
  40. diffusers/pipelines/animatediff/pipeline_animatediff.py +38 -10
  41. diffusers/pipelines/auto_pipeline.py +17 -13
  42. diffusers/pipelines/controlnet/pipeline_controlnet.py +27 -10
  43. diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py +47 -5
  44. diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py +25 -8
  45. diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py +4 -6
  46. diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py +26 -10
  47. diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py +4 -3
  48. diffusers/pipelines/deprecated/__init__.py +153 -0
  49. diffusers/pipelines/{alt_diffusion → deprecated/alt_diffusion}/__init__.py +3 -3
  50. diffusers/pipelines/{alt_diffusion → deprecated/alt_diffusion}/pipeline_alt_diffusion.py +91 -18
  51. diffusers/pipelines/{alt_diffusion → deprecated/alt_diffusion}/pipeline_alt_diffusion_img2img.py +91 -18
  52. diffusers/pipelines/{alt_diffusion → deprecated/alt_diffusion}/pipeline_output.py +1 -1
  53. diffusers/pipelines/{audio_diffusion → deprecated/audio_diffusion}/__init__.py +1 -1
  54. diffusers/pipelines/{audio_diffusion → deprecated/audio_diffusion}/mel.py +2 -2
  55. diffusers/pipelines/{audio_diffusion → deprecated/audio_diffusion}/pipeline_audio_diffusion.py +4 -4
  56. diffusers/pipelines/{latent_diffusion_uncond → deprecated/latent_diffusion_uncond}/__init__.py +1 -1
  57. diffusers/pipelines/{latent_diffusion_uncond → deprecated/latent_diffusion_uncond}/pipeline_latent_diffusion_uncond.py +4 -4
  58. diffusers/pipelines/{pndm → deprecated/pndm}/__init__.py +1 -1
  59. diffusers/pipelines/{pndm → deprecated/pndm}/pipeline_pndm.py +4 -4
  60. diffusers/pipelines/{repaint → deprecated/repaint}/__init__.py +1 -1
  61. diffusers/pipelines/{repaint → deprecated/repaint}/pipeline_repaint.py +5 -5
  62. diffusers/pipelines/{score_sde_ve → deprecated/score_sde_ve}/__init__.py +1 -1
  63. diffusers/pipelines/{score_sde_ve → deprecated/score_sde_ve}/pipeline_score_sde_ve.py +4 -4
  64. diffusers/pipelines/{spectrogram_diffusion → deprecated/spectrogram_diffusion}/__init__.py +6 -6
  65. diffusers/pipelines/{spectrogram_diffusion/continous_encoder.py → deprecated/spectrogram_diffusion/continuous_encoder.py} +2 -2
  66. diffusers/pipelines/{spectrogram_diffusion → deprecated/spectrogram_diffusion}/midi_utils.py +1 -1
  67. diffusers/pipelines/{spectrogram_diffusion → deprecated/spectrogram_diffusion}/notes_encoder.py +2 -2
  68. diffusers/pipelines/{spectrogram_diffusion → deprecated/spectrogram_diffusion}/pipeline_spectrogram_diffusion.py +7 -7
  69. diffusers/pipelines/deprecated/stable_diffusion_variants/__init__.py +55 -0
  70. diffusers/pipelines/{stable_diffusion → deprecated/stable_diffusion_variants}/pipeline_cycle_diffusion.py +16 -11
  71. diffusers/pipelines/{stable_diffusion → deprecated/stable_diffusion_variants}/pipeline_onnx_stable_diffusion_inpaint_legacy.py +6 -6
  72. diffusers/pipelines/{stable_diffusion → deprecated/stable_diffusion_variants}/pipeline_stable_diffusion_inpaint_legacy.py +11 -11
  73. diffusers/pipelines/{stable_diffusion → deprecated/stable_diffusion_variants}/pipeline_stable_diffusion_model_editing.py +16 -11
  74. diffusers/pipelines/{stable_diffusion → deprecated/stable_diffusion_variants}/pipeline_stable_diffusion_paradigms.py +10 -10
  75. diffusers/pipelines/{stable_diffusion → deprecated/stable_diffusion_variants}/pipeline_stable_diffusion_pix2pix_zero.py +13 -13
  76. diffusers/pipelines/{stochastic_karras_ve → deprecated/stochastic_karras_ve}/__init__.py +1 -1
  77. diffusers/pipelines/{stochastic_karras_ve → deprecated/stochastic_karras_ve}/pipeline_stochastic_karras_ve.py +4 -4
  78. diffusers/pipelines/{versatile_diffusion → deprecated/versatile_diffusion}/__init__.py +3 -3
  79. diffusers/pipelines/{versatile_diffusion → deprecated/versatile_diffusion}/modeling_text_unet.py +54 -11
  80. diffusers/pipelines/{versatile_diffusion → deprecated/versatile_diffusion}/pipeline_versatile_diffusion.py +4 -4
  81. diffusers/pipelines/{versatile_diffusion → deprecated/versatile_diffusion}/pipeline_versatile_diffusion_dual_guided.py +6 -6
  82. diffusers/pipelines/{versatile_diffusion → deprecated/versatile_diffusion}/pipeline_versatile_diffusion_image_variation.py +6 -6
  83. diffusers/pipelines/{versatile_diffusion → deprecated/versatile_diffusion}/pipeline_versatile_diffusion_text_to_image.py +6 -6
  84. diffusers/pipelines/{vq_diffusion → deprecated/vq_diffusion}/__init__.py +3 -3
  85. diffusers/pipelines/{vq_diffusion → deprecated/vq_diffusion}/pipeline_vq_diffusion.py +5 -5
  86. diffusers/pipelines/kandinsky3/__init__.py +4 -4
  87. diffusers/pipelines/kandinsky3/convert_kandinsky3_unet.py +98 -0
  88. diffusers/pipelines/kandinsky3/{kandinsky3_pipeline.py → pipeline_kandinsky3.py} +172 -35
  89. diffusers/pipelines/kandinsky3/{kandinsky3img2img_pipeline.py → pipeline_kandinsky3_img2img.py} +228 -34
  90. diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_img2img.py +46 -5
  91. diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_text2img.py +47 -6
  92. diffusers/pipelines/onnx_utils.py +8 -5
  93. diffusers/pipelines/pipeline_flax_utils.py +7 -6
  94. diffusers/pipelines/pipeline_utils.py +32 -31
  95. diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py +51 -2
  96. diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py +3 -3
  97. diffusers/pipelines/stable_diffusion/__init__.py +1 -72
  98. diffusers/pipelines/stable_diffusion/convert_from_ckpt.py +67 -75
  99. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +92 -8
  100. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py +92 -8
  101. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py +138 -10
  102. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py +57 -7
  103. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_latent_upscale.py +3 -0
  104. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py +6 -0
  105. diffusers/pipelines/stable_diffusion/pipeline_stable_unclip.py +5 -0
  106. diffusers/pipelines/stable_diffusion/pipeline_stable_unclip_img2img.py +5 -0
  107. diffusers/pipelines/stable_diffusion_attend_and_excite/__init__.py +48 -0
  108. diffusers/pipelines/{stable_diffusion → stable_diffusion_attend_and_excite}/pipeline_stable_diffusion_attend_and_excite.py +5 -2
  109. diffusers/pipelines/stable_diffusion_diffedit/__init__.py +48 -0
  110. diffusers/pipelines/{stable_diffusion → stable_diffusion_diffedit}/pipeline_stable_diffusion_diffedit.py +2 -3
  111. diffusers/pipelines/stable_diffusion_gligen/__init__.py +50 -0
  112. diffusers/pipelines/{stable_diffusion → stable_diffusion_gligen}/pipeline_stable_diffusion_gligen.py +2 -2
  113. diffusers/pipelines/{stable_diffusion → stable_diffusion_gligen}/pipeline_stable_diffusion_gligen_text_image.py +3 -3
  114. diffusers/pipelines/stable_diffusion_k_diffusion/__init__.py +60 -0
  115. diffusers/pipelines/{stable_diffusion → stable_diffusion_k_diffusion}/pipeline_stable_diffusion_k_diffusion.py +6 -1
  116. diffusers/pipelines/stable_diffusion_ldm3d/__init__.py +48 -0
  117. diffusers/pipelines/{stable_diffusion → stable_diffusion_ldm3d}/pipeline_stable_diffusion_ldm3d.py +50 -7
  118. diffusers/pipelines/stable_diffusion_panorama/__init__.py +48 -0
  119. diffusers/pipelines/{stable_diffusion → stable_diffusion_panorama}/pipeline_stable_diffusion_panorama.py +56 -8
  120. diffusers/pipelines/stable_diffusion_safe/pipeline_stable_diffusion_safe.py +58 -6
  121. diffusers/pipelines/stable_diffusion_sag/__init__.py +48 -0
  122. diffusers/pipelines/{stable_diffusion → stable_diffusion_sag}/pipeline_stable_diffusion_sag.py +67 -10
  123. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py +97 -15
  124. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py +98 -14
  125. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py +97 -14
  126. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_instruct_pix2pix.py +7 -5
  127. diffusers/pipelines/stable_video_diffusion/pipeline_stable_video_diffusion.py +12 -9
  128. diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_xl_adapter.py +6 -0
  129. diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth.py +5 -0
  130. diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth_img2img.py +5 -0
  131. diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero.py +331 -9
  132. diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero_sdxl.py +468 -9
  133. diffusers/pipelines/unclip/pipeline_unclip.py +2 -1
  134. diffusers/pipelines/unclip/pipeline_unclip_image_variation.py +1 -0
  135. diffusers/pipelines/wuerstchen/modeling_paella_vq_model.py +1 -1
  136. diffusers/pipelines/wuerstchen/pipeline_wuerstchen_prior.py +4 -0
  137. diffusers/schedulers/__init__.py +2 -0
  138. diffusers/schedulers/scheduling_amused.py +162 -0
  139. diffusers/schedulers/scheduling_consistency_models.py +2 -0
  140. diffusers/schedulers/scheduling_ddim_inverse.py +1 -4
  141. diffusers/schedulers/scheduling_ddpm.py +46 -0
  142. diffusers/schedulers/scheduling_ddpm_parallel.py +46 -0
  143. diffusers/schedulers/scheduling_deis_multistep.py +13 -1
  144. diffusers/schedulers/scheduling_dpmsolver_multistep.py +13 -1
  145. diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py +13 -1
  146. diffusers/schedulers/scheduling_dpmsolver_sde.py +2 -0
  147. diffusers/schedulers/scheduling_dpmsolver_singlestep.py +13 -1
  148. diffusers/schedulers/scheduling_euler_ancestral_discrete.py +58 -0
  149. diffusers/schedulers/scheduling_euler_discrete.py +62 -3
  150. diffusers/schedulers/scheduling_heun_discrete.py +2 -0
  151. diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py +2 -0
  152. diffusers/schedulers/scheduling_k_dpm_2_discrete.py +2 -0
  153. diffusers/schedulers/scheduling_lms_discrete.py +2 -0
  154. diffusers/schedulers/scheduling_unipc_multistep.py +13 -1
  155. diffusers/schedulers/scheduling_utils.py +3 -1
  156. diffusers/schedulers/scheduling_utils_flax.py +3 -1
  157. diffusers/training_utils.py +1 -1
  158. diffusers/utils/__init__.py +0 -2
  159. diffusers/utils/constants.py +2 -5
  160. diffusers/utils/dummy_pt_objects.py +30 -0
  161. diffusers/utils/dummy_torch_and_transformers_objects.py +45 -0
  162. diffusers/utils/dynamic_modules_utils.py +14 -18
  163. diffusers/utils/hub_utils.py +24 -36
  164. diffusers/utils/logging.py +1 -1
  165. diffusers/utils/state_dict_utils.py +8 -0
  166. diffusers/utils/testing_utils.py +199 -1
  167. diffusers/utils/torch_utils.py +3 -3
  168. {diffusers-0.24.0.dist-info → diffusers-0.25.1.dist-info}/METADATA +55 -53
  169. {diffusers-0.24.0.dist-info → diffusers-0.25.1.dist-info}/RECORD +174 -155
  170. {diffusers-0.24.0.dist-info → diffusers-0.25.1.dist-info}/WHEEL +1 -1
  171. {diffusers-0.24.0.dist-info → diffusers-0.25.1.dist-info}/entry_points.txt +0 -1
  172. /diffusers/pipelines/{alt_diffusion → deprecated/alt_diffusion}/modeling_roberta_series.py +0 -0
  173. {diffusers-0.24.0.dist-info → diffusers-0.25.1.dist-info}/LICENSE +0 -0
  174. {diffusers-0.24.0.dist-info → diffusers-0.25.1.dist-info}/top_level.txt +0 -0
@@ -1,5 +1,5 @@
1
1
  import inspect
2
- from typing import Callable, List, Optional, Union
2
+ from typing import Callable, Dict, List, Optional, Union
3
3
 
4
4
  import numpy as np
5
5
  import PIL
@@ -11,8 +11,10 @@ from ...loaders import LoraLoaderMixin
11
11
  from ...models import Kandinsky3UNet, VQModel
12
12
  from ...schedulers import DDPMScheduler
13
13
  from ...utils import (
14
+ deprecate,
14
15
  is_accelerate_available,
15
16
  logging,
17
+ replace_example_docstring,
16
18
  )
17
19
  from ...utils.torch_utils import randn_tensor
18
20
  from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput
@@ -20,6 +22,24 @@ from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput
20
22
 
21
23
  logger = logging.get_logger(__name__) # pylint: disable=invalid-name
22
24
 
25
+ EXAMPLE_DOC_STRING = """
26
+ Examples:
27
+ ```py
28
+ >>> from diffusers import AutoPipelineForImage2Image
29
+ >>> from diffusers.utils import load_image
30
+ >>> import torch
31
+
32
+ >>> pipe = AutoPipelineForImage2Image.from_pretrained("kandinsky-community/kandinsky-3", variant="fp16", torch_dtype=torch.float16)
33
+ >>> pipe.enable_model_cpu_offload()
34
+
35
+ >>> prompt = "A painting of the inside of a subway train with tiny raccoons."
36
+ >>> image = load_image("https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/kandinsky3/t2i.png")
37
+
38
+ >>> generator = torch.Generator(device="cpu").manual_seed(0)
39
+ >>> image = pipe(prompt, image=image, strength=0.75, num_inference_steps=25, generator=generator).images[0]
40
+ ```
41
+ """
42
+
23
43
 
24
44
  def downscale_height_and_width(height, width, scale_factor=8):
25
45
  new_height = height // scale_factor**2
@@ -40,7 +60,14 @@ def prepare_image(pil_image):
40
60
 
41
61
 
42
62
  class Kandinsky3Img2ImgPipeline(DiffusionPipeline, LoraLoaderMixin):
43
- model_cpu_offload_seq = "text_encoder->unet->movq"
63
+ model_cpu_offload_seq = "text_encoder->movq->unet->movq"
64
+ _callback_tensor_inputs = [
65
+ "latents",
66
+ "prompt_embeds",
67
+ "negative_prompt_embeds",
68
+ "negative_attention_mask",
69
+ "attention_mask",
70
+ ]
44
71
 
45
72
  def __init__(
46
73
  self,
@@ -99,6 +126,8 @@ class Kandinsky3Img2ImgPipeline(DiffusionPipeline, LoraLoaderMixin):
99
126
  prompt_embeds: Optional[torch.FloatTensor] = None,
100
127
  negative_prompt_embeds: Optional[torch.FloatTensor] = None,
101
128
  _cut_context=False,
129
+ attention_mask: Optional[torch.FloatTensor] = None,
130
+ negative_attention_mask: Optional[torch.FloatTensor] = None,
102
131
  ):
103
132
  r"""
104
133
  Encodes the prompt into text encoder hidden states.
@@ -123,6 +152,10 @@ class Kandinsky3Img2ImgPipeline(DiffusionPipeline, LoraLoaderMixin):
123
152
  Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
124
153
  weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
125
154
  argument.
155
+ attention_mask (`torch.FloatTensor`, *optional*):
156
+ Pre-generated attention mask. Must provide if passing `prompt_embeds` directly.
157
+ negative_attention_mask (`torch.FloatTensor`, *optional*):
158
+ Pre-generated negative attention mask. Must provide if passing `negative_prompt_embeds` directly.
126
159
  """
127
160
  if prompt is not None and negative_prompt is not None:
128
161
  if type(prompt) is not type(negative_prompt):
@@ -299,15 +332,23 @@ class Kandinsky3Img2ImgPipeline(DiffusionPipeline, LoraLoaderMixin):
299
332
  negative_prompt=None,
300
333
  prompt_embeds=None,
301
334
  negative_prompt_embeds=None,
335
+ callback_on_step_end_tensor_inputs=None,
336
+ attention_mask=None,
337
+ negative_attention_mask=None,
302
338
  ):
303
- if (callback_steps is None) or (
304
- callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
305
- ):
339
+ if callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0):
306
340
  raise ValueError(
307
341
  f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
308
342
  f" {type(callback_steps)}."
309
343
  )
310
344
 
345
+ if callback_on_step_end_tensor_inputs is not None and not all(
346
+ k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
347
+ ):
348
+ raise ValueError(
349
+ f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
350
+ )
351
+
311
352
  if prompt is not None and prompt_embeds is not None:
312
353
  raise ValueError(
313
354
  f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
@@ -334,7 +375,42 @@ class Kandinsky3Img2ImgPipeline(DiffusionPipeline, LoraLoaderMixin):
334
375
  f" {negative_prompt_embeds.shape}."
335
376
  )
336
377
 
378
+ if negative_prompt_embeds is not None and negative_attention_mask is None:
379
+ raise ValueError("Please provide `negative_attention_mask` along with `negative_prompt_embeds`")
380
+
381
+ if negative_prompt_embeds is not None and negative_attention_mask is not None:
382
+ if negative_prompt_embeds.shape[:2] != negative_attention_mask.shape:
383
+ raise ValueError(
384
+ "`negative_prompt_embeds` and `negative_attention_mask` must have the same batch_size and token length when passed directly, but"
385
+ f" got: `negative_prompt_embeds` {negative_prompt_embeds.shape[:2]} != `negative_attention_mask`"
386
+ f" {negative_attention_mask.shape}."
387
+ )
388
+
389
+ if prompt_embeds is not None and attention_mask is None:
390
+ raise ValueError("Please provide `attention_mask` along with `prompt_embeds`")
391
+
392
+ if prompt_embeds is not None and attention_mask is not None:
393
+ if prompt_embeds.shape[:2] != attention_mask.shape:
394
+ raise ValueError(
395
+ "`prompt_embeds` and `attention_mask` must have the same batch_size and token length when passed directly, but"
396
+ f" got: `prompt_embeds` {prompt_embeds.shape[:2]} != `attention_mask`"
397
+ f" {attention_mask.shape}."
398
+ )
399
+
400
+ @property
401
+ def guidance_scale(self):
402
+ return self._guidance_scale
403
+
404
+ @property
405
+ def do_classifier_free_guidance(self):
406
+ return self._guidance_scale > 1
407
+
408
+ @property
409
+ def num_timesteps(self):
410
+ return self._num_timesteps
411
+
337
412
  @torch.no_grad()
413
+ @replace_example_docstring(EXAMPLE_DOC_STRING)
338
414
  def __call__(
339
415
  self,
340
416
  prompt: Union[str, List[str]] = None,
@@ -347,15 +423,117 @@ class Kandinsky3Img2ImgPipeline(DiffusionPipeline, LoraLoaderMixin):
347
423
  generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
348
424
  prompt_embeds: Optional[torch.FloatTensor] = None,
349
425
  negative_prompt_embeds: Optional[torch.FloatTensor] = None,
426
+ attention_mask: Optional[torch.FloatTensor] = None,
427
+ negative_attention_mask: Optional[torch.FloatTensor] = None,
350
428
  output_type: Optional[str] = "pil",
351
429
  return_dict: bool = True,
352
- callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
353
- callback_steps: int = 1,
354
- latents=None,
430
+ callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
431
+ callback_on_step_end_tensor_inputs: List[str] = ["latents"],
432
+ **kwargs,
355
433
  ):
434
+ """
435
+ Function invoked when calling the pipeline for generation.
436
+
437
+ Args:
438
+ prompt (`str` or `List[str]`, *optional*):
439
+ The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
440
+ instead.
441
+ image (`torch.FloatTensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.FloatTensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
442
+ `Image`, or tensor representing an image batch, that will be used as the starting point for the
443
+ process.
444
+ strength (`float`, *optional*, defaults to 0.8):
445
+ Indicates extent to transform the reference `image`. Must be between 0 and 1. `image` is used as a
446
+ starting point and more noise is added the higher the `strength`. The number of denoising steps depends
447
+ on the amount of noise initially added. When `strength` is 1, added noise is maximum and the denoising
448
+ process runs for the full number of iterations specified in `num_inference_steps`. A value of 1
449
+ essentially ignores `image`.
450
+ num_inference_steps (`int`, *optional*, defaults to 50):
451
+ The number of denoising steps. More denoising steps usually lead to a higher quality image at the
452
+ expense of slower inference.
453
+ guidance_scale (`float`, *optional*, defaults to 3.0):
454
+ Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
455
+ `guidance_scale` is defined as `w` of equation 2. of [Imagen
456
+ Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
457
+ 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
458
+ usually at the expense of lower image quality.
459
+ negative_prompt (`str` or `List[str]`, *optional*):
460
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass
461
+ `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
462
+ less than `1`).
463
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
464
+ The number of images to generate per prompt.
465
+ generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
466
+ One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
467
+ to make generation deterministic.
468
+ prompt_embeds (`torch.FloatTensor`, *optional*):
469
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
470
+ provided, text embeddings will be generated from `prompt` input argument.
471
+ negative_prompt_embeds (`torch.FloatTensor`, *optional*):
472
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
473
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
474
+ argument.
475
+ attention_mask (`torch.FloatTensor`, *optional*):
476
+ Pre-generated attention mask. Must provide if passing `prompt_embeds` directly.
477
+ negative_attention_mask (`torch.FloatTensor`, *optional*):
478
+ Pre-generated negative attention mask. Must provide if passing `negative_prompt_embeds` directly.
479
+ output_type (`str`, *optional*, defaults to `"pil"`):
480
+ The output format of the generate image. Choose between
481
+ [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
482
+ return_dict (`bool`, *optional*, defaults to `True`):
483
+ Whether or not to return a [`~pipelines.stable_diffusion.IFPipelineOutput`] instead of a plain tuple.
484
+ callback_on_step_end (`Callable`, *optional*):
485
+ A function that calls at the end of each denoising steps during the inference. The function is called
486
+ with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
487
+ callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
488
+ `callback_on_step_end_tensor_inputs`.
489
+ callback_on_step_end_tensor_inputs (`List`, *optional*):
490
+ The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
491
+ will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
492
+ `._callback_tensor_inputs` attribute of your pipeline class.
493
+
494
+ Examples:
495
+
496
+ Returns:
497
+ [`~pipelines.ImagePipelineOutput`] or `tuple`
498
+
499
+ """
500
+ callback = kwargs.pop("callback", None)
501
+ callback_steps = kwargs.pop("callback_steps", None)
502
+
503
+ if callback is not None:
504
+ deprecate(
505
+ "callback",
506
+ "1.0.0",
507
+ "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
508
+ )
509
+ if callback_steps is not None:
510
+ deprecate(
511
+ "callback_steps",
512
+ "1.0.0",
513
+ "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
514
+ )
515
+
516
+ if callback_on_step_end_tensor_inputs is not None and not all(
517
+ k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
518
+ ):
519
+ raise ValueError(
520
+ f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
521
+ )
522
+
356
523
  cut_context = True
357
524
  # 1. Check inputs. Raise error if not correct
358
- self.check_inputs(prompt, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds)
525
+ self.check_inputs(
526
+ prompt,
527
+ callback_steps,
528
+ negative_prompt,
529
+ prompt_embeds,
530
+ negative_prompt_embeds,
531
+ callback_on_step_end_tensor_inputs,
532
+ attention_mask,
533
+ negative_attention_mask,
534
+ )
535
+
536
+ self._guidance_scale = guidance_scale
359
537
 
360
538
  if prompt is not None and isinstance(prompt, str):
361
539
  batch_size = 1
@@ -366,24 +544,21 @@ class Kandinsky3Img2ImgPipeline(DiffusionPipeline, LoraLoaderMixin):
366
544
 
367
545
  device = self._execution_device
368
546
 
369
- # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
370
- # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
371
- # corresponds to doing no classifier free guidance.
372
- do_classifier_free_guidance = guidance_scale > 1.0
373
-
374
547
  # 3. Encode input prompt
375
548
  prompt_embeds, negative_prompt_embeds, attention_mask, negative_attention_mask = self.encode_prompt(
376
549
  prompt,
377
- do_classifier_free_guidance,
550
+ self.do_classifier_free_guidance,
378
551
  num_images_per_prompt=num_images_per_prompt,
379
552
  device=device,
380
553
  negative_prompt=negative_prompt,
381
554
  prompt_embeds=prompt_embeds,
382
555
  negative_prompt_embeds=negative_prompt_embeds,
383
556
  _cut_context=cut_context,
557
+ attention_mask=attention_mask,
558
+ negative_attention_mask=negative_attention_mask,
384
559
  )
385
560
 
386
- if do_classifier_free_guidance:
561
+ if self.do_classifier_free_guidance:
387
562
  prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
388
563
  attention_mask = torch.cat([negative_attention_mask, attention_mask]).bool()
389
564
  if not isinstance(image, list):
@@ -409,11 +584,11 @@ class Kandinsky3Img2ImgPipeline(DiffusionPipeline, LoraLoaderMixin):
409
584
  self.text_encoder_offload_hook.offload()
410
585
 
411
586
  # 7. Denoising loop
412
- # TODO(Yiyi): Correct the following line and use correctly
413
- # num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
587
+ num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
588
+ self._num_timesteps = len(timesteps)
414
589
  with self.progress_bar(total=num_inference_steps) as progress_bar:
415
590
  for i, t in enumerate(timesteps):
416
- latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
591
+ latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
417
592
 
418
593
  # predict the noise residual
419
594
  noise_pred = self.unet(
@@ -422,7 +597,7 @@ class Kandinsky3Img2ImgPipeline(DiffusionPipeline, LoraLoaderMixin):
422
597
  encoder_hidden_states=prompt_embeds,
423
598
  encoder_attention_mask=attention_mask,
424
599
  )[0]
425
- if do_classifier_free_guidance:
600
+ if self.do_classifier_free_guidance:
426
601
  noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
427
602
 
428
603
  noise_pred = (guidance_scale + 1.0) * noise_pred_text - guidance_scale * noise_pred_uncond
@@ -434,25 +609,44 @@ class Kandinsky3Img2ImgPipeline(DiffusionPipeline, LoraLoaderMixin):
434
609
  latents,
435
610
  generator=generator,
436
611
  ).prev_sample
437
- progress_bar.update()
438
- if callback is not None and i % callback_steps == 0:
439
- step_idx = i // getattr(self.scheduler, "order", 1)
440
- callback(step_idx, t, latents)
441
- # post-processing
442
- image = self.movq.decode(latents, force_not_quantize=True)["sample"]
443
612
 
444
- if output_type not in ["pt", "np", "pil"]:
613
+ if callback_on_step_end is not None:
614
+ callback_kwargs = {}
615
+ for k in callback_on_step_end_tensor_inputs:
616
+ callback_kwargs[k] = locals()[k]
617
+ callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
618
+
619
+ latents = callback_outputs.pop("latents", latents)
620
+ prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
621
+ negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
622
+ attention_mask = callback_outputs.pop("attention_mask", attention_mask)
623
+ negative_attention_mask = callback_outputs.pop("negative_attention_mask", negative_attention_mask)
624
+
625
+ if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
626
+ progress_bar.update()
627
+ if callback is not None and i % callback_steps == 0:
628
+ step_idx = i // getattr(self.scheduler, "order", 1)
629
+ callback(step_idx, t, latents)
630
+
631
+ # post-processing
632
+ if output_type not in ["pt", "np", "pil", "latent"]:
445
633
  raise ValueError(
446
- f"Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}"
634
+ f"Only the output types `pt`, `pil`, `np` and `latent` are supported not output_type={output_type}"
447
635
  )
636
+ if not output_type == "latent":
637
+ image = self.movq.decode(latents, force_not_quantize=True)["sample"]
638
+
639
+ if output_type in ["np", "pil"]:
640
+ image = image * 0.5 + 0.5
641
+ image = image.clamp(0, 1)
642
+ image = image.cpu().permute(0, 2, 3, 1).float().numpy()
448
643
 
449
- if output_type in ["np", "pil"]:
450
- image = image * 0.5 + 0.5
451
- image = image.clamp(0, 1)
452
- image = image.cpu().permute(0, 2, 3, 1).float().numpy()
644
+ if output_type == "pil":
645
+ image = self.numpy_to_pil(image)
646
+ else:
647
+ image = latents
453
648
 
454
- if output_type == "pil":
455
- image = self.numpy_to_pil(image)
649
+ self.maybe_free_model_hooks()
456
650
 
457
651
  if not return_dict:
458
652
  return (image,)
@@ -20,11 +20,11 @@ from typing import Any, Callable, Dict, List, Optional, Union
20
20
 
21
21
  import PIL.Image
22
22
  import torch
23
- from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
23
+ from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection
24
24
 
25
25
  from ...image_processor import PipelineImageInput, VaeImageProcessor
26
- from ...loaders import FromSingleFileMixin, LoraLoaderMixin, TextualInversionLoaderMixin
27
- from ...models import AutoencoderKL, UNet2DConditionModel
26
+ from ...loaders import FromSingleFileMixin, IPAdapterMixin, LoraLoaderMixin, TextualInversionLoaderMixin
27
+ from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel
28
28
  from ...models.lora import adjust_lora_scale_text_encoder
29
29
  from ...schedulers import LCMScheduler
30
30
  from ...utils import (
@@ -129,7 +129,7 @@ EXAMPLE_DOC_STRING = """
129
129
 
130
130
 
131
131
  class LatentConsistencyModelImg2ImgPipeline(
132
- DiffusionPipeline, TextualInversionLoaderMixin, LoraLoaderMixin, FromSingleFileMixin
132
+ DiffusionPipeline, TextualInversionLoaderMixin, IPAdapterMixin, LoraLoaderMixin, FromSingleFileMixin
133
133
  ):
134
134
  r"""
135
135
  Pipeline for image-to-image generation using a latent consistency model.
@@ -142,6 +142,7 @@ class LatentConsistencyModelImg2ImgPipeline(
142
142
  - [`~loaders.LoraLoaderMixin.load_lora_weights`] for loading LoRA weights
143
143
  - [`~loaders.LoraLoaderMixin.save_lora_weights`] for saving LoRA weights
144
144
  - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files
145
+ - [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters
145
146
 
146
147
  Args:
147
148
  vae ([`AutoencoderKL`]):
@@ -166,7 +167,7 @@ class LatentConsistencyModelImg2ImgPipeline(
166
167
  """
167
168
 
168
169
  model_cpu_offload_seq = "text_encoder->unet->vae"
169
- _optional_components = ["safety_checker", "feature_extractor"]
170
+ _optional_components = ["safety_checker", "feature_extractor", "image_encoder"]
170
171
  _exclude_from_cpu_offload = ["safety_checker"]
171
172
  _callback_tensor_inputs = ["latents", "denoised", "prompt_embeds", "w_embedding"]
172
173
 
@@ -179,6 +180,7 @@ class LatentConsistencyModelImg2ImgPipeline(
179
180
  scheduler: LCMScheduler,
180
181
  safety_checker: StableDiffusionSafetyChecker,
181
182
  feature_extractor: CLIPImageProcessor,
183
+ image_encoder: Optional[CLIPVisionModelWithProjection] = None,
182
184
  requires_safety_checker: bool = True,
183
185
  ):
184
186
  super().__init__()
@@ -191,6 +193,7 @@ class LatentConsistencyModelImg2ImgPipeline(
191
193
  scheduler=scheduler,
192
194
  safety_checker=safety_checker,
193
195
  feature_extractor=feature_extractor,
196
+ image_encoder=image_encoder,
194
197
  )
195
198
 
196
199
  if safety_checker is None and requires_safety_checker:
@@ -449,6 +452,31 @@ class LatentConsistencyModelImg2ImgPipeline(
449
452
 
450
453
  return prompt_embeds, negative_prompt_embeds
451
454
 
455
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image
456
+ def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None):
457
+ dtype = next(self.image_encoder.parameters()).dtype
458
+
459
+ if not isinstance(image, torch.Tensor):
460
+ image = self.feature_extractor(image, return_tensors="pt").pixel_values
461
+
462
+ image = image.to(device=device, dtype=dtype)
463
+ if output_hidden_states:
464
+ image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
465
+ image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
466
+ uncond_image_enc_hidden_states = self.image_encoder(
467
+ torch.zeros_like(image), output_hidden_states=True
468
+ ).hidden_states[-2]
469
+ uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave(
470
+ num_images_per_prompt, dim=0
471
+ )
472
+ return image_enc_hidden_states, uncond_image_enc_hidden_states
473
+ else:
474
+ image_embeds = self.image_encoder(image).image_embeds
475
+ image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
476
+ uncond_image_embeds = torch.zeros_like(image_embeds)
477
+
478
+ return image_embeds, uncond_image_embeds
479
+
452
480
  # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker
453
481
  def run_safety_checker(self, image, device, dtype):
454
482
  if self.safety_checker is None:
@@ -647,6 +675,7 @@ class LatentConsistencyModelImg2ImgPipeline(
647
675
  generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
648
676
  latents: Optional[torch.FloatTensor] = None,
649
677
  prompt_embeds: Optional[torch.FloatTensor] = None,
678
+ ip_adapter_image: Optional[PipelineImageInput] = None,
650
679
  output_type: Optional[str] = "pil",
651
680
  return_dict: bool = True,
652
681
  cross_attention_kwargs: Optional[Dict[str, Any]] = None,
@@ -695,6 +724,8 @@ class LatentConsistencyModelImg2ImgPipeline(
695
724
  prompt_embeds (`torch.FloatTensor`, *optional*):
696
725
  Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
697
726
  provided, text embeddings are generated from the `prompt` input argument.
727
+ ip_adapter_image: (`PipelineImageInput`, *optional*):
728
+ Optional image input to work with IP Adapters.
698
729
  output_type (`str`, *optional*, defaults to `"pil"`):
699
730
  The output format of the generated image. Choose between `PIL.Image` or `np.array`.
700
731
  return_dict (`bool`, *optional*, defaults to `True`):
@@ -758,6 +789,12 @@ class LatentConsistencyModelImg2ImgPipeline(
758
789
  device = self._execution_device
759
790
  # do_classifier_free_guidance = guidance_scale > 1.0
760
791
 
792
+ if ip_adapter_image is not None:
793
+ output_hidden_state = False if isinstance(self.unet.encoder_hid_proj, ImageProjection) else True
794
+ image_embeds, negative_image_embeds = self.encode_image(
795
+ ip_adapter_image, device, num_images_per_prompt, output_hidden_state
796
+ )
797
+
761
798
  # 3. Encode input prompt
762
799
  lora_scale = (
763
800
  self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None
@@ -815,6 +852,9 @@ class LatentConsistencyModelImg2ImgPipeline(
815
852
  # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
816
853
  extra_step_kwargs = self.prepare_extra_step_kwargs(generator, None)
817
854
 
855
+ # 7.1 Add image embeds for IP-Adapter
856
+ added_cond_kwargs = {"image_embeds": image_embeds} if ip_adapter_image is not None else None
857
+
818
858
  # 8. LCM Multistep Sampling Loop
819
859
  num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
820
860
  self._num_timesteps = len(timesteps)
@@ -829,6 +869,7 @@ class LatentConsistencyModelImg2ImgPipeline(
829
869
  timestep_cond=w_embedding,
830
870
  encoder_hidden_states=prompt_embeds,
831
871
  cross_attention_kwargs=self.cross_attention_kwargs,
872
+ added_cond_kwargs=added_cond_kwargs,
832
873
  return_dict=False,
833
874
  )[0]
834
875
 
@@ -19,11 +19,11 @@ import inspect
19
19
  from typing import Any, Callable, Dict, List, Optional, Union
20
20
 
21
21
  import torch
22
- from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
22
+ from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection
23
23
 
24
- from ...image_processor import VaeImageProcessor
25
- from ...loaders import FromSingleFileMixin, LoraLoaderMixin, TextualInversionLoaderMixin
26
- from ...models import AutoencoderKL, UNet2DConditionModel
24
+ from ...image_processor import PipelineImageInput, VaeImageProcessor
25
+ from ...loaders import FromSingleFileMixin, IPAdapterMixin, LoraLoaderMixin, TextualInversionLoaderMixin
26
+ from ...models import AutoencoderKL, ImageProjection, UNet2DConditionModel
27
27
  from ...models.lora import adjust_lora_scale_text_encoder
28
28
  from ...schedulers import LCMScheduler
29
29
  from ...utils import (
@@ -107,7 +107,7 @@ def retrieve_timesteps(
107
107
 
108
108
 
109
109
  class LatentConsistencyModelPipeline(
110
- DiffusionPipeline, TextualInversionLoaderMixin, LoraLoaderMixin, FromSingleFileMixin
110
+ DiffusionPipeline, TextualInversionLoaderMixin, IPAdapterMixin, LoraLoaderMixin, FromSingleFileMixin
111
111
  ):
112
112
  r"""
113
113
  Pipeline for text-to-image generation using a latent consistency model.
@@ -120,6 +120,7 @@ class LatentConsistencyModelPipeline(
120
120
  - [`~loaders.LoraLoaderMixin.load_lora_weights`] for loading LoRA weights
121
121
  - [`~loaders.LoraLoaderMixin.save_lora_weights`] for saving LoRA weights
122
122
  - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files
123
+ - [`~loaders.IPAdapterMixin.load_ip_adapter`] for loading IP Adapters
123
124
 
124
125
  Args:
125
126
  vae ([`AutoencoderKL`]):
@@ -144,7 +145,7 @@ class LatentConsistencyModelPipeline(
144
145
  """
145
146
 
146
147
  model_cpu_offload_seq = "text_encoder->unet->vae"
147
- _optional_components = ["safety_checker", "feature_extractor"]
148
+ _optional_components = ["safety_checker", "feature_extractor", "image_encoder"]
148
149
  _exclude_from_cpu_offload = ["safety_checker"]
149
150
  _callback_tensor_inputs = ["latents", "denoised", "prompt_embeds", "w_embedding"]
150
151
 
@@ -157,6 +158,7 @@ class LatentConsistencyModelPipeline(
157
158
  scheduler: LCMScheduler,
158
159
  safety_checker: StableDiffusionSafetyChecker,
159
160
  feature_extractor: CLIPImageProcessor,
161
+ image_encoder: Optional[CLIPVisionModelWithProjection] = None,
160
162
  requires_safety_checker: bool = True,
161
163
  ):
162
164
  super().__init__()
@@ -185,6 +187,7 @@ class LatentConsistencyModelPipeline(
185
187
  scheduler=scheduler,
186
188
  safety_checker=safety_checker,
187
189
  feature_extractor=feature_extractor,
190
+ image_encoder=image_encoder,
188
191
  )
189
192
  self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
190
193
  self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
@@ -433,6 +436,31 @@ class LatentConsistencyModelPipeline(
433
436
 
434
437
  return prompt_embeds, negative_prompt_embeds
435
438
 
439
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image
440
+ def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None):
441
+ dtype = next(self.image_encoder.parameters()).dtype
442
+
443
+ if not isinstance(image, torch.Tensor):
444
+ image = self.feature_extractor(image, return_tensors="pt").pixel_values
445
+
446
+ image = image.to(device=device, dtype=dtype)
447
+ if output_hidden_states:
448
+ image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
449
+ image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
450
+ uncond_image_enc_hidden_states = self.image_encoder(
451
+ torch.zeros_like(image), output_hidden_states=True
452
+ ).hidden_states[-2]
453
+ uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave(
454
+ num_images_per_prompt, dim=0
455
+ )
456
+ return image_enc_hidden_states, uncond_image_enc_hidden_states
457
+ else:
458
+ image_embeds = self.image_encoder(image).image_embeds
459
+ image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
460
+ uncond_image_embeds = torch.zeros_like(image_embeds)
461
+
462
+ return image_embeds, uncond_image_embeds
463
+
436
464
  # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.run_safety_checker
437
465
  def run_safety_checker(self, image, device, dtype):
438
466
  if self.safety_checker is None:
@@ -581,6 +609,7 @@ class LatentConsistencyModelPipeline(
581
609
  generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
582
610
  latents: Optional[torch.FloatTensor] = None,
583
611
  prompt_embeds: Optional[torch.FloatTensor] = None,
612
+ ip_adapter_image: Optional[PipelineImageInput] = None,
584
613
  output_type: Optional[str] = "pil",
585
614
  return_dict: bool = True,
586
615
  cross_attention_kwargs: Optional[Dict[str, Any]] = None,
@@ -629,6 +658,8 @@ class LatentConsistencyModelPipeline(
629
658
  prompt_embeds (`torch.FloatTensor`, *optional*):
630
659
  Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
631
660
  provided, text embeddings are generated from the `prompt` input argument.
661
+ ip_adapter_image: (`PipelineImageInput`, *optional*):
662
+ Optional image input to work with IP Adapters.
632
663
  output_type (`str`, *optional*, defaults to `"pil"`):
633
664
  The output format of the generated image. Choose between `PIL.Image` or `np.array`.
634
665
  return_dict (`bool`, *optional*, defaults to `True`):
@@ -697,6 +728,12 @@ class LatentConsistencyModelPipeline(
697
728
  device = self._execution_device
698
729
  # do_classifier_free_guidance = guidance_scale > 1.0
699
730
 
731
+ if ip_adapter_image is not None:
732
+ output_hidden_state = False if isinstance(self.unet.encoder_hid_proj, ImageProjection) else True
733
+ image_embeds, negative_image_embeds = self.encode_image(
734
+ ip_adapter_image, device, num_images_per_prompt, output_hidden_state
735
+ )
736
+
700
737
  # 3. Encode input prompt
701
738
  lora_scale = (
702
739
  self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None
@@ -748,6 +785,9 @@ class LatentConsistencyModelPipeline(
748
785
  # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
749
786
  extra_step_kwargs = self.prepare_extra_step_kwargs(generator, None)
750
787
 
788
+ # 7.1 Add image embeds for IP-Adapter
789
+ added_cond_kwargs = {"image_embeds": image_embeds} if ip_adapter_image is not None else None
790
+
751
791
  # 8. LCM MultiStep Sampling Loop:
752
792
  num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
753
793
  self._num_timesteps = len(timesteps)
@@ -762,6 +802,7 @@ class LatentConsistencyModelPipeline(
762
802
  timestep_cond=w_embedding,
763
803
  encoder_hidden_states=prompt_embeds,
764
804
  cross_attention_kwargs=self.cross_attention_kwargs,
805
+ added_cond_kwargs=added_cond_kwargs,
765
806
  return_dict=False,
766
807
  )[0]
767
808