diffusers 0.27.0__py3-none-any.whl → 0.32.2__py3-none-any.whl

Sign up to get free protection for your applications and to get access to all the features.
Files changed (445) hide show
  1. diffusers/__init__.py +233 -6
  2. diffusers/callbacks.py +209 -0
  3. diffusers/commands/env.py +102 -6
  4. diffusers/configuration_utils.py +45 -16
  5. diffusers/dependency_versions_table.py +4 -3
  6. diffusers/image_processor.py +434 -110
  7. diffusers/loaders/__init__.py +42 -9
  8. diffusers/loaders/ip_adapter.py +626 -36
  9. diffusers/loaders/lora_base.py +900 -0
  10. diffusers/loaders/lora_conversion_utils.py +991 -125
  11. diffusers/loaders/lora_pipeline.py +3812 -0
  12. diffusers/loaders/peft.py +571 -7
  13. diffusers/loaders/single_file.py +405 -173
  14. diffusers/loaders/single_file_model.py +385 -0
  15. diffusers/loaders/single_file_utils.py +1783 -713
  16. diffusers/loaders/textual_inversion.py +41 -23
  17. diffusers/loaders/transformer_flux.py +181 -0
  18. diffusers/loaders/transformer_sd3.py +89 -0
  19. diffusers/loaders/unet.py +464 -540
  20. diffusers/loaders/unet_loader_utils.py +163 -0
  21. diffusers/models/__init__.py +76 -7
  22. diffusers/models/activations.py +65 -10
  23. diffusers/models/adapter.py +53 -53
  24. diffusers/models/attention.py +605 -18
  25. diffusers/models/attention_flax.py +1 -1
  26. diffusers/models/attention_processor.py +4304 -687
  27. diffusers/models/autoencoders/__init__.py +8 -0
  28. diffusers/models/autoencoders/autoencoder_asym_kl.py +15 -17
  29. diffusers/models/autoencoders/autoencoder_dc.py +620 -0
  30. diffusers/models/autoencoders/autoencoder_kl.py +110 -28
  31. diffusers/models/autoencoders/autoencoder_kl_allegro.py +1149 -0
  32. diffusers/models/autoencoders/autoencoder_kl_cogvideox.py +1482 -0
  33. diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py +1176 -0
  34. diffusers/models/autoencoders/autoencoder_kl_ltx.py +1338 -0
  35. diffusers/models/autoencoders/autoencoder_kl_mochi.py +1166 -0
  36. diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py +19 -24
  37. diffusers/models/autoencoders/autoencoder_oobleck.py +464 -0
  38. diffusers/models/autoencoders/autoencoder_tiny.py +21 -18
  39. diffusers/models/autoencoders/consistency_decoder_vae.py +45 -20
  40. diffusers/models/autoencoders/vae.py +41 -29
  41. diffusers/models/autoencoders/vq_model.py +182 -0
  42. diffusers/models/controlnet.py +47 -800
  43. diffusers/models/controlnet_flux.py +70 -0
  44. diffusers/models/controlnet_sd3.py +68 -0
  45. diffusers/models/controlnet_sparsectrl.py +116 -0
  46. diffusers/models/controlnets/__init__.py +23 -0
  47. diffusers/models/controlnets/controlnet.py +872 -0
  48. diffusers/models/{controlnet_flax.py → controlnets/controlnet_flax.py} +9 -9
  49. diffusers/models/controlnets/controlnet_flux.py +536 -0
  50. diffusers/models/controlnets/controlnet_hunyuan.py +401 -0
  51. diffusers/models/controlnets/controlnet_sd3.py +489 -0
  52. diffusers/models/controlnets/controlnet_sparsectrl.py +788 -0
  53. diffusers/models/controlnets/controlnet_union.py +832 -0
  54. diffusers/models/controlnets/controlnet_xs.py +1946 -0
  55. diffusers/models/controlnets/multicontrolnet.py +183 -0
  56. diffusers/models/downsampling.py +85 -18
  57. diffusers/models/embeddings.py +1856 -158
  58. diffusers/models/embeddings_flax.py +23 -9
  59. diffusers/models/model_loading_utils.py +480 -0
  60. diffusers/models/modeling_flax_pytorch_utils.py +2 -1
  61. diffusers/models/modeling_flax_utils.py +2 -7
  62. diffusers/models/modeling_outputs.py +14 -0
  63. diffusers/models/modeling_pytorch_flax_utils.py +1 -1
  64. diffusers/models/modeling_utils.py +611 -146
  65. diffusers/models/normalization.py +361 -20
  66. diffusers/models/resnet.py +18 -23
  67. diffusers/models/transformers/__init__.py +16 -0
  68. diffusers/models/transformers/auraflow_transformer_2d.py +544 -0
  69. diffusers/models/transformers/cogvideox_transformer_3d.py +542 -0
  70. diffusers/models/transformers/dit_transformer_2d.py +240 -0
  71. diffusers/models/transformers/dual_transformer_2d.py +9 -8
  72. diffusers/models/transformers/hunyuan_transformer_2d.py +578 -0
  73. diffusers/models/transformers/latte_transformer_3d.py +327 -0
  74. diffusers/models/transformers/lumina_nextdit2d.py +340 -0
  75. diffusers/models/transformers/pixart_transformer_2d.py +445 -0
  76. diffusers/models/transformers/prior_transformer.py +13 -13
  77. diffusers/models/transformers/sana_transformer.py +488 -0
  78. diffusers/models/transformers/stable_audio_transformer.py +458 -0
  79. diffusers/models/transformers/t5_film_transformer.py +17 -19
  80. diffusers/models/transformers/transformer_2d.py +297 -187
  81. diffusers/models/transformers/transformer_allegro.py +422 -0
  82. diffusers/models/transformers/transformer_cogview3plus.py +386 -0
  83. diffusers/models/transformers/transformer_flux.py +593 -0
  84. diffusers/models/transformers/transformer_hunyuan_video.py +791 -0
  85. diffusers/models/transformers/transformer_ltx.py +469 -0
  86. diffusers/models/transformers/transformer_mochi.py +499 -0
  87. diffusers/models/transformers/transformer_sd3.py +461 -0
  88. diffusers/models/transformers/transformer_temporal.py +21 -19
  89. diffusers/models/unets/unet_1d.py +8 -8
  90. diffusers/models/unets/unet_1d_blocks.py +31 -31
  91. diffusers/models/unets/unet_2d.py +17 -10
  92. diffusers/models/unets/unet_2d_blocks.py +225 -149
  93. diffusers/models/unets/unet_2d_condition.py +50 -53
  94. diffusers/models/unets/unet_2d_condition_flax.py +6 -5
  95. diffusers/models/unets/unet_3d_blocks.py +192 -1057
  96. diffusers/models/unets/unet_3d_condition.py +22 -27
  97. diffusers/models/unets/unet_i2vgen_xl.py +22 -18
  98. diffusers/models/unets/unet_kandinsky3.py +2 -2
  99. diffusers/models/unets/unet_motion_model.py +1413 -89
  100. diffusers/models/unets/unet_spatio_temporal_condition.py +40 -16
  101. diffusers/models/unets/unet_stable_cascade.py +19 -18
  102. diffusers/models/unets/uvit_2d.py +2 -2
  103. diffusers/models/upsampling.py +95 -26
  104. diffusers/models/vq_model.py +12 -164
  105. diffusers/optimization.py +1 -1
  106. diffusers/pipelines/__init__.py +202 -3
  107. diffusers/pipelines/allegro/__init__.py +48 -0
  108. diffusers/pipelines/allegro/pipeline_allegro.py +938 -0
  109. diffusers/pipelines/allegro/pipeline_output.py +23 -0
  110. diffusers/pipelines/amused/pipeline_amused.py +12 -12
  111. diffusers/pipelines/amused/pipeline_amused_img2img.py +14 -12
  112. diffusers/pipelines/amused/pipeline_amused_inpaint.py +13 -11
  113. diffusers/pipelines/animatediff/__init__.py +8 -0
  114. diffusers/pipelines/animatediff/pipeline_animatediff.py +122 -109
  115. diffusers/pipelines/animatediff/pipeline_animatediff_controlnet.py +1106 -0
  116. diffusers/pipelines/animatediff/pipeline_animatediff_sdxl.py +1288 -0
  117. diffusers/pipelines/animatediff/pipeline_animatediff_sparsectrl.py +1010 -0
  118. diffusers/pipelines/animatediff/pipeline_animatediff_video2video.py +236 -180
  119. diffusers/pipelines/animatediff/pipeline_animatediff_video2video_controlnet.py +1341 -0
  120. diffusers/pipelines/animatediff/pipeline_output.py +3 -2
  121. diffusers/pipelines/audioldm/pipeline_audioldm.py +14 -14
  122. diffusers/pipelines/audioldm2/modeling_audioldm2.py +58 -39
  123. diffusers/pipelines/audioldm2/pipeline_audioldm2.py +121 -36
  124. diffusers/pipelines/aura_flow/__init__.py +48 -0
  125. diffusers/pipelines/aura_flow/pipeline_aura_flow.py +584 -0
  126. diffusers/pipelines/auto_pipeline.py +196 -28
  127. diffusers/pipelines/blip_diffusion/blip_image_processing.py +1 -1
  128. diffusers/pipelines/blip_diffusion/modeling_blip2.py +6 -6
  129. diffusers/pipelines/blip_diffusion/modeling_ctx_clip.py +1 -1
  130. diffusers/pipelines/blip_diffusion/pipeline_blip_diffusion.py +2 -2
  131. diffusers/pipelines/cogvideo/__init__.py +54 -0
  132. diffusers/pipelines/cogvideo/pipeline_cogvideox.py +772 -0
  133. diffusers/pipelines/cogvideo/pipeline_cogvideox_fun_control.py +825 -0
  134. diffusers/pipelines/cogvideo/pipeline_cogvideox_image2video.py +885 -0
  135. diffusers/pipelines/cogvideo/pipeline_cogvideox_video2video.py +851 -0
  136. diffusers/pipelines/cogvideo/pipeline_output.py +20 -0
  137. diffusers/pipelines/cogview3/__init__.py +47 -0
  138. diffusers/pipelines/cogview3/pipeline_cogview3plus.py +674 -0
  139. diffusers/pipelines/cogview3/pipeline_output.py +21 -0
  140. diffusers/pipelines/consistency_models/pipeline_consistency_models.py +6 -6
  141. diffusers/pipelines/controlnet/__init__.py +86 -80
  142. diffusers/pipelines/controlnet/multicontrolnet.py +7 -182
  143. diffusers/pipelines/controlnet/pipeline_controlnet.py +134 -87
  144. diffusers/pipelines/controlnet/pipeline_controlnet_blip_diffusion.py +2 -2
  145. diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py +93 -77
  146. diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py +88 -197
  147. diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py +136 -90
  148. diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py +176 -80
  149. diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py +125 -89
  150. diffusers/pipelines/controlnet/pipeline_controlnet_union_inpaint_sd_xl.py +1790 -0
  151. diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl.py +1501 -0
  152. diffusers/pipelines/controlnet/pipeline_controlnet_union_sd_xl_img2img.py +1627 -0
  153. diffusers/pipelines/controlnet/pipeline_flax_controlnet.py +2 -2
  154. diffusers/pipelines/controlnet_hunyuandit/__init__.py +48 -0
  155. diffusers/pipelines/controlnet_hunyuandit/pipeline_hunyuandit_controlnet.py +1060 -0
  156. diffusers/pipelines/controlnet_sd3/__init__.py +57 -0
  157. diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py +1133 -0
  158. diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet_inpainting.py +1153 -0
  159. diffusers/pipelines/controlnet_xs/__init__.py +68 -0
  160. diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs.py +916 -0
  161. diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs_sd_xl.py +1111 -0
  162. diffusers/pipelines/ddpm/pipeline_ddpm.py +2 -2
  163. diffusers/pipelines/deepfloyd_if/pipeline_if.py +16 -30
  164. diffusers/pipelines/deepfloyd_if/pipeline_if_img2img.py +20 -35
  165. diffusers/pipelines/deepfloyd_if/pipeline_if_img2img_superresolution.py +23 -41
  166. diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting.py +22 -38
  167. diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting_superresolution.py +25 -41
  168. diffusers/pipelines/deepfloyd_if/pipeline_if_superresolution.py +19 -34
  169. diffusers/pipelines/deepfloyd_if/pipeline_output.py +6 -5
  170. diffusers/pipelines/deepfloyd_if/watermark.py +1 -1
  171. diffusers/pipelines/deprecated/alt_diffusion/modeling_roberta_series.py +11 -11
  172. diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion.py +70 -30
  173. diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion_img2img.py +48 -25
  174. diffusers/pipelines/deprecated/repaint/pipeline_repaint.py +2 -2
  175. diffusers/pipelines/deprecated/spectrogram_diffusion/pipeline_spectrogram_diffusion.py +7 -7
  176. diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_cycle_diffusion.py +21 -20
  177. diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_inpaint_legacy.py +27 -29
  178. diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_model_editing.py +33 -27
  179. diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_paradigms.py +33 -23
  180. diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_pix2pix_zero.py +36 -30
  181. diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py +102 -69
  182. diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion.py +13 -13
  183. diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_dual_guided.py +10 -5
  184. diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_image_variation.py +11 -6
  185. diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_text_to_image.py +10 -5
  186. diffusers/pipelines/deprecated/vq_diffusion/pipeline_vq_diffusion.py +5 -5
  187. diffusers/pipelines/dit/pipeline_dit.py +7 -4
  188. diffusers/pipelines/flux/__init__.py +69 -0
  189. diffusers/pipelines/flux/modeling_flux.py +47 -0
  190. diffusers/pipelines/flux/pipeline_flux.py +957 -0
  191. diffusers/pipelines/flux/pipeline_flux_control.py +889 -0
  192. diffusers/pipelines/flux/pipeline_flux_control_img2img.py +945 -0
  193. diffusers/pipelines/flux/pipeline_flux_control_inpaint.py +1141 -0
  194. diffusers/pipelines/flux/pipeline_flux_controlnet.py +1006 -0
  195. diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py +998 -0
  196. diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py +1204 -0
  197. diffusers/pipelines/flux/pipeline_flux_fill.py +969 -0
  198. diffusers/pipelines/flux/pipeline_flux_img2img.py +856 -0
  199. diffusers/pipelines/flux/pipeline_flux_inpaint.py +1022 -0
  200. diffusers/pipelines/flux/pipeline_flux_prior_redux.py +492 -0
  201. diffusers/pipelines/flux/pipeline_output.py +37 -0
  202. diffusers/pipelines/free_init_utils.py +41 -38
  203. diffusers/pipelines/free_noise_utils.py +596 -0
  204. diffusers/pipelines/hunyuan_video/__init__.py +48 -0
  205. diffusers/pipelines/hunyuan_video/pipeline_hunyuan_video.py +687 -0
  206. diffusers/pipelines/hunyuan_video/pipeline_output.py +20 -0
  207. diffusers/pipelines/hunyuandit/__init__.py +48 -0
  208. diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py +916 -0
  209. diffusers/pipelines/i2vgen_xl/pipeline_i2vgen_xl.py +33 -48
  210. diffusers/pipelines/kandinsky/pipeline_kandinsky.py +8 -8
  211. diffusers/pipelines/kandinsky/pipeline_kandinsky_combined.py +32 -29
  212. diffusers/pipelines/kandinsky/pipeline_kandinsky_img2img.py +11 -11
  213. diffusers/pipelines/kandinsky/pipeline_kandinsky_inpaint.py +12 -12
  214. diffusers/pipelines/kandinsky/pipeline_kandinsky_prior.py +10 -10
  215. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py +6 -6
  216. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py +34 -31
  217. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet.py +10 -10
  218. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet_img2img.py +10 -10
  219. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py +6 -6
  220. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpainting.py +8 -8
  221. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py +7 -7
  222. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior_emb2emb.py +6 -6
  223. diffusers/pipelines/kandinsky3/convert_kandinsky3_unet.py +3 -3
  224. diffusers/pipelines/kandinsky3/pipeline_kandinsky3.py +22 -35
  225. diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py +26 -37
  226. diffusers/pipelines/kolors/__init__.py +54 -0
  227. diffusers/pipelines/kolors/pipeline_kolors.py +1070 -0
  228. diffusers/pipelines/kolors/pipeline_kolors_img2img.py +1250 -0
  229. diffusers/pipelines/kolors/pipeline_output.py +21 -0
  230. diffusers/pipelines/kolors/text_encoder.py +889 -0
  231. diffusers/pipelines/kolors/tokenizer.py +338 -0
  232. diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_img2img.py +82 -62
  233. diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_text2img.py +77 -60
  234. diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py +12 -12
  235. diffusers/pipelines/latte/__init__.py +48 -0
  236. diffusers/pipelines/latte/pipeline_latte.py +881 -0
  237. diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py +80 -74
  238. diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py +85 -76
  239. diffusers/pipelines/ledits_pp/pipeline_output.py +2 -2
  240. diffusers/pipelines/ltx/__init__.py +50 -0
  241. diffusers/pipelines/ltx/pipeline_ltx.py +789 -0
  242. diffusers/pipelines/ltx/pipeline_ltx_image2video.py +885 -0
  243. diffusers/pipelines/ltx/pipeline_output.py +20 -0
  244. diffusers/pipelines/lumina/__init__.py +48 -0
  245. diffusers/pipelines/lumina/pipeline_lumina.py +890 -0
  246. diffusers/pipelines/marigold/__init__.py +50 -0
  247. diffusers/pipelines/marigold/marigold_image_processing.py +576 -0
  248. diffusers/pipelines/marigold/pipeline_marigold_depth.py +813 -0
  249. diffusers/pipelines/marigold/pipeline_marigold_normals.py +690 -0
  250. diffusers/pipelines/mochi/__init__.py +48 -0
  251. diffusers/pipelines/mochi/pipeline_mochi.py +748 -0
  252. diffusers/pipelines/mochi/pipeline_output.py +20 -0
  253. diffusers/pipelines/musicldm/pipeline_musicldm.py +14 -14
  254. diffusers/pipelines/pag/__init__.py +80 -0
  255. diffusers/pipelines/pag/pag_utils.py +243 -0
  256. diffusers/pipelines/pag/pipeline_pag_controlnet_sd.py +1328 -0
  257. diffusers/pipelines/pag/pipeline_pag_controlnet_sd_inpaint.py +1543 -0
  258. diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl.py +1610 -0
  259. diffusers/pipelines/pag/pipeline_pag_controlnet_sd_xl_img2img.py +1683 -0
  260. diffusers/pipelines/pag/pipeline_pag_hunyuandit.py +969 -0
  261. diffusers/pipelines/pag/pipeline_pag_kolors.py +1136 -0
  262. diffusers/pipelines/pag/pipeline_pag_pixart_sigma.py +865 -0
  263. diffusers/pipelines/pag/pipeline_pag_sana.py +886 -0
  264. diffusers/pipelines/pag/pipeline_pag_sd.py +1062 -0
  265. diffusers/pipelines/pag/pipeline_pag_sd_3.py +994 -0
  266. diffusers/pipelines/pag/pipeline_pag_sd_3_img2img.py +1058 -0
  267. diffusers/pipelines/pag/pipeline_pag_sd_animatediff.py +866 -0
  268. diffusers/pipelines/pag/pipeline_pag_sd_img2img.py +1094 -0
  269. diffusers/pipelines/pag/pipeline_pag_sd_inpaint.py +1356 -0
  270. diffusers/pipelines/pag/pipeline_pag_sd_xl.py +1345 -0
  271. diffusers/pipelines/pag/pipeline_pag_sd_xl_img2img.py +1544 -0
  272. diffusers/pipelines/pag/pipeline_pag_sd_xl_inpaint.py +1776 -0
  273. diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py +17 -12
  274. diffusers/pipelines/pia/pipeline_pia.py +74 -164
  275. diffusers/pipelines/pipeline_flax_utils.py +5 -10
  276. diffusers/pipelines/pipeline_loading_utils.py +515 -53
  277. diffusers/pipelines/pipeline_utils.py +411 -222
  278. diffusers/pipelines/pixart_alpha/__init__.py +8 -1
  279. diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py +76 -93
  280. diffusers/pipelines/pixart_alpha/pipeline_pixart_sigma.py +873 -0
  281. diffusers/pipelines/sana/__init__.py +47 -0
  282. diffusers/pipelines/sana/pipeline_output.py +21 -0
  283. diffusers/pipelines/sana/pipeline_sana.py +884 -0
  284. diffusers/pipelines/semantic_stable_diffusion/pipeline_semantic_stable_diffusion.py +27 -23
  285. diffusers/pipelines/shap_e/pipeline_shap_e.py +3 -3
  286. diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py +14 -14
  287. diffusers/pipelines/shap_e/renderer.py +1 -1
  288. diffusers/pipelines/stable_audio/__init__.py +50 -0
  289. diffusers/pipelines/stable_audio/modeling_stable_audio.py +158 -0
  290. diffusers/pipelines/stable_audio/pipeline_stable_audio.py +756 -0
  291. diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py +71 -25
  292. diffusers/pipelines/stable_cascade/pipeline_stable_cascade_combined.py +23 -19
  293. diffusers/pipelines/stable_cascade/pipeline_stable_cascade_prior.py +35 -34
  294. diffusers/pipelines/stable_diffusion/__init__.py +0 -1
  295. diffusers/pipelines/stable_diffusion/convert_from_ckpt.py +20 -11
  296. diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py +1 -1
  297. diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py +2 -2
  298. diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_upscale.py +6 -6
  299. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +145 -79
  300. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_depth2img.py +43 -28
  301. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py +13 -8
  302. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py +100 -68
  303. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py +109 -201
  304. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py +131 -32
  305. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_latent_upscale.py +247 -87
  306. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py +30 -29
  307. diffusers/pipelines/stable_diffusion/pipeline_stable_unclip.py +35 -27
  308. diffusers/pipelines/stable_diffusion/pipeline_stable_unclip_img2img.py +49 -42
  309. diffusers/pipelines/stable_diffusion/safety_checker.py +2 -1
  310. diffusers/pipelines/stable_diffusion_3/__init__.py +54 -0
  311. diffusers/pipelines/stable_diffusion_3/pipeline_output.py +21 -0
  312. diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py +1140 -0
  313. diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py +1036 -0
  314. diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py +1250 -0
  315. diffusers/pipelines/stable_diffusion_attend_and_excite/pipeline_stable_diffusion_attend_and_excite.py +29 -20
  316. diffusers/pipelines/stable_diffusion_diffedit/pipeline_stable_diffusion_diffedit.py +59 -58
  317. diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen.py +31 -25
  318. diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen_text_image.py +38 -22
  319. diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_k_diffusion.py +30 -24
  320. diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_xl_k_diffusion.py +24 -23
  321. diffusers/pipelines/stable_diffusion_ldm3d/pipeline_stable_diffusion_ldm3d.py +107 -67
  322. diffusers/pipelines/stable_diffusion_panorama/pipeline_stable_diffusion_panorama.py +316 -69
  323. diffusers/pipelines/stable_diffusion_safe/pipeline_stable_diffusion_safe.py +10 -5
  324. diffusers/pipelines/stable_diffusion_safe/safety_checker.py +1 -1
  325. diffusers/pipelines/stable_diffusion_sag/pipeline_stable_diffusion_sag.py +98 -30
  326. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py +121 -83
  327. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py +161 -105
  328. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py +142 -218
  329. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_instruct_pix2pix.py +45 -29
  330. diffusers/pipelines/stable_diffusion_xl/watermark.py +9 -3
  331. diffusers/pipelines/stable_video_diffusion/pipeline_stable_video_diffusion.py +110 -57
  332. diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_adapter.py +69 -39
  333. diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_xl_adapter.py +105 -74
  334. diffusers/pipelines/text_to_video_synthesis/pipeline_output.py +3 -2
  335. diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth.py +29 -49
  336. diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth_img2img.py +32 -93
  337. diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero.py +37 -25
  338. diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero_sdxl.py +54 -40
  339. diffusers/pipelines/unclip/pipeline_unclip.py +6 -6
  340. diffusers/pipelines/unclip/pipeline_unclip_image_variation.py +6 -6
  341. diffusers/pipelines/unidiffuser/modeling_text_decoder.py +1 -1
  342. diffusers/pipelines/unidiffuser/modeling_uvit.py +12 -12
  343. diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py +29 -28
  344. diffusers/pipelines/wuerstchen/modeling_paella_vq_model.py +5 -5
  345. diffusers/pipelines/wuerstchen/modeling_wuerstchen_common.py +5 -10
  346. diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py +6 -8
  347. diffusers/pipelines/wuerstchen/pipeline_wuerstchen.py +4 -4
  348. diffusers/pipelines/wuerstchen/pipeline_wuerstchen_combined.py +12 -12
  349. diffusers/pipelines/wuerstchen/pipeline_wuerstchen_prior.py +15 -14
  350. diffusers/{models/dual_transformer_2d.py → quantizers/__init__.py} +2 -6
  351. diffusers/quantizers/auto.py +139 -0
  352. diffusers/quantizers/base.py +233 -0
  353. diffusers/quantizers/bitsandbytes/__init__.py +2 -0
  354. diffusers/quantizers/bitsandbytes/bnb_quantizer.py +561 -0
  355. diffusers/quantizers/bitsandbytes/utils.py +306 -0
  356. diffusers/quantizers/gguf/__init__.py +1 -0
  357. diffusers/quantizers/gguf/gguf_quantizer.py +159 -0
  358. diffusers/quantizers/gguf/utils.py +456 -0
  359. diffusers/quantizers/quantization_config.py +669 -0
  360. diffusers/quantizers/torchao/__init__.py +15 -0
  361. diffusers/quantizers/torchao/torchao_quantizer.py +292 -0
  362. diffusers/schedulers/__init__.py +12 -2
  363. diffusers/schedulers/deprecated/__init__.py +1 -1
  364. diffusers/schedulers/deprecated/scheduling_karras_ve.py +25 -25
  365. diffusers/schedulers/scheduling_amused.py +5 -5
  366. diffusers/schedulers/scheduling_consistency_decoder.py +11 -11
  367. diffusers/schedulers/scheduling_consistency_models.py +23 -25
  368. diffusers/schedulers/scheduling_cosine_dpmsolver_multistep.py +572 -0
  369. diffusers/schedulers/scheduling_ddim.py +27 -26
  370. diffusers/schedulers/scheduling_ddim_cogvideox.py +452 -0
  371. diffusers/schedulers/scheduling_ddim_flax.py +2 -1
  372. diffusers/schedulers/scheduling_ddim_inverse.py +16 -16
  373. diffusers/schedulers/scheduling_ddim_parallel.py +32 -31
  374. diffusers/schedulers/scheduling_ddpm.py +27 -30
  375. diffusers/schedulers/scheduling_ddpm_flax.py +7 -3
  376. diffusers/schedulers/scheduling_ddpm_parallel.py +33 -36
  377. diffusers/schedulers/scheduling_ddpm_wuerstchen.py +14 -14
  378. diffusers/schedulers/scheduling_deis_multistep.py +150 -50
  379. diffusers/schedulers/scheduling_dpm_cogvideox.py +489 -0
  380. diffusers/schedulers/scheduling_dpmsolver_multistep.py +221 -84
  381. diffusers/schedulers/scheduling_dpmsolver_multistep_flax.py +2 -2
  382. diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py +158 -52
  383. diffusers/schedulers/scheduling_dpmsolver_sde.py +153 -34
  384. diffusers/schedulers/scheduling_dpmsolver_singlestep.py +275 -86
  385. diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py +81 -57
  386. diffusers/schedulers/scheduling_edm_euler.py +62 -39
  387. diffusers/schedulers/scheduling_euler_ancestral_discrete.py +30 -29
  388. diffusers/schedulers/scheduling_euler_discrete.py +255 -74
  389. diffusers/schedulers/scheduling_flow_match_euler_discrete.py +458 -0
  390. diffusers/schedulers/scheduling_flow_match_heun_discrete.py +320 -0
  391. diffusers/schedulers/scheduling_heun_discrete.py +174 -46
  392. diffusers/schedulers/scheduling_ipndm.py +9 -9
  393. diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py +138 -29
  394. diffusers/schedulers/scheduling_k_dpm_2_discrete.py +132 -26
  395. diffusers/schedulers/scheduling_karras_ve_flax.py +6 -6
  396. diffusers/schedulers/scheduling_lcm.py +23 -29
  397. diffusers/schedulers/scheduling_lms_discrete.py +105 -28
  398. diffusers/schedulers/scheduling_pndm.py +20 -20
  399. diffusers/schedulers/scheduling_repaint.py +21 -21
  400. diffusers/schedulers/scheduling_sasolver.py +157 -60
  401. diffusers/schedulers/scheduling_sde_ve.py +19 -19
  402. diffusers/schedulers/scheduling_tcd.py +41 -36
  403. diffusers/schedulers/scheduling_unclip.py +19 -16
  404. diffusers/schedulers/scheduling_unipc_multistep.py +243 -47
  405. diffusers/schedulers/scheduling_utils.py +12 -5
  406. diffusers/schedulers/scheduling_utils_flax.py +1 -3
  407. diffusers/schedulers/scheduling_vq_diffusion.py +10 -10
  408. diffusers/training_utils.py +214 -30
  409. diffusers/utils/__init__.py +17 -1
  410. diffusers/utils/constants.py +3 -0
  411. diffusers/utils/doc_utils.py +1 -0
  412. diffusers/utils/dummy_pt_objects.py +592 -7
  413. diffusers/utils/dummy_torch_and_torchsde_objects.py +15 -0
  414. diffusers/utils/dummy_torch_and_transformers_and_sentencepiece_objects.py +47 -0
  415. diffusers/utils/dummy_torch_and_transformers_objects.py +1001 -71
  416. diffusers/utils/dynamic_modules_utils.py +34 -29
  417. diffusers/utils/export_utils.py +50 -6
  418. diffusers/utils/hub_utils.py +131 -17
  419. diffusers/utils/import_utils.py +210 -8
  420. diffusers/utils/loading_utils.py +118 -5
  421. diffusers/utils/logging.py +4 -2
  422. diffusers/utils/peft_utils.py +37 -7
  423. diffusers/utils/state_dict_utils.py +13 -2
  424. diffusers/utils/testing_utils.py +193 -11
  425. diffusers/utils/torch_utils.py +4 -0
  426. diffusers/video_processor.py +113 -0
  427. {diffusers-0.27.0.dist-info → diffusers-0.32.2.dist-info}/METADATA +82 -91
  428. diffusers-0.32.2.dist-info/RECORD +550 -0
  429. {diffusers-0.27.0.dist-info → diffusers-0.32.2.dist-info}/WHEEL +1 -1
  430. diffusers/loaders/autoencoder.py +0 -146
  431. diffusers/loaders/controlnet.py +0 -136
  432. diffusers/loaders/lora.py +0 -1349
  433. diffusers/models/prior_transformer.py +0 -12
  434. diffusers/models/t5_film_transformer.py +0 -70
  435. diffusers/models/transformer_2d.py +0 -25
  436. diffusers/models/transformer_temporal.py +0 -34
  437. diffusers/models/unet_1d.py +0 -26
  438. diffusers/models/unet_1d_blocks.py +0 -203
  439. diffusers/models/unet_2d.py +0 -27
  440. diffusers/models/unet_2d_blocks.py +0 -375
  441. diffusers/models/unet_2d_condition.py +0 -25
  442. diffusers-0.27.0.dist-info/RECORD +0 -399
  443. {diffusers-0.27.0.dist-info → diffusers-0.32.2.dist-info}/LICENSE +0 -0
  444. {diffusers-0.27.0.dist-info → diffusers-0.32.2.dist-info}/entry_points.txt +0 -0
  445. {diffusers-0.27.0.dist-info → diffusers-0.32.2.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,690 @@
1
+ # Copyright 2024 Marigold authors, PRS ETH Zurich. All rights reserved.
2
+ # Copyright 2024 The HuggingFace Team. All rights reserved.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ # --------------------------------------------------------------------------
16
+ # More information and citation instructions are available on the
17
+ # Marigold project website: https://marigoldmonodepth.github.io
18
+ # --------------------------------------------------------------------------
19
+ from dataclasses import dataclass
20
+ from typing import Any, Dict, List, Optional, Tuple, Union
21
+
22
+ import numpy as np
23
+ import torch
24
+ from PIL import Image
25
+ from tqdm.auto import tqdm
26
+ from transformers import CLIPTextModel, CLIPTokenizer
27
+
28
+ from ...image_processor import PipelineImageInput
29
+ from ...models import (
30
+ AutoencoderKL,
31
+ UNet2DConditionModel,
32
+ )
33
+ from ...schedulers import (
34
+ DDIMScheduler,
35
+ LCMScheduler,
36
+ )
37
+ from ...utils import (
38
+ BaseOutput,
39
+ logging,
40
+ replace_example_docstring,
41
+ )
42
+ from ...utils.torch_utils import randn_tensor
43
+ from ..pipeline_utils import DiffusionPipeline
44
+ from .marigold_image_processing import MarigoldImageProcessor
45
+
46
+
47
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
48
+
49
+
50
+ EXAMPLE_DOC_STRING = """
51
+ Examples:
52
+ ```py
53
+ >>> import diffusers
54
+ >>> import torch
55
+
56
+ >>> pipe = diffusers.MarigoldNormalsPipeline.from_pretrained(
57
+ ... "prs-eth/marigold-normals-lcm-v0-1", variant="fp16", torch_dtype=torch.float16
58
+ ... ).to("cuda")
59
+
60
+ >>> image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg")
61
+ >>> normals = pipe(image)
62
+
63
+ >>> vis = pipe.image_processor.visualize_normals(normals.prediction)
64
+ >>> vis[0].save("einstein_normals.png")
65
+ ```
66
+ """
67
+
68
+
69
+ @dataclass
70
+ class MarigoldNormalsOutput(BaseOutput):
71
+ """
72
+ Output class for Marigold monocular normals prediction pipeline.
73
+
74
+ Args:
75
+ prediction (`np.ndarray`, `torch.Tensor`):
76
+ Predicted normals with values in the range [-1, 1]. The shape is always $numimages \times 3 \times height
77
+ \times width$, regardless of whether the images were passed as a 4D array or a list.
78
+ uncertainty (`None`, `np.ndarray`, `torch.Tensor`):
79
+ Uncertainty maps computed from the ensemble, with values in the range [0, 1]. The shape is $numimages
80
+ \times 1 \times height \times width$.
81
+ latent (`None`, `torch.Tensor`):
82
+ Latent features corresponding to the predictions, compatible with the `latents` argument of the pipeline.
83
+ The shape is $numimages * numensemble \times 4 \times latentheight \times latentwidth$.
84
+ """
85
+
86
+ prediction: Union[np.ndarray, torch.Tensor]
87
+ uncertainty: Union[None, np.ndarray, torch.Tensor]
88
+ latent: Union[None, torch.Tensor]
89
+
90
+
91
+ class MarigoldNormalsPipeline(DiffusionPipeline):
92
+ """
93
+ Pipeline for monocular normals estimation using the Marigold method: https://marigoldmonodepth.github.io.
94
+
95
+ This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
96
+ library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
97
+
98
+ Args:
99
+ unet (`UNet2DConditionModel`):
100
+ Conditional U-Net to denoise the normals latent, conditioned on image latent.
101
+ vae (`AutoencoderKL`):
102
+ Variational Auto-Encoder (VAE) Model to encode and decode images and predictions to and from latent
103
+ representations.
104
+ scheduler (`DDIMScheduler` or `LCMScheduler`):
105
+ A scheduler to be used in combination with `unet` to denoise the encoded image latents.
106
+ text_encoder (`CLIPTextModel`):
107
+ Text-encoder, for empty text embedding.
108
+ tokenizer (`CLIPTokenizer`):
109
+ CLIP tokenizer.
110
+ prediction_type (`str`, *optional*):
111
+ Type of predictions made by the model.
112
+ use_full_z_range (`bool`, *optional*):
113
+ Whether the normals predicted by this model utilize the full range of the Z dimension, or only its positive
114
+ half.
115
+ default_denoising_steps (`int`, *optional*):
116
+ The minimum number of denoising diffusion steps that are required to produce a prediction of reasonable
117
+ quality with the given model. This value must be set in the model config. When the pipeline is called
118
+ without explicitly setting `num_inference_steps`, the default value is used. This is required to ensure
119
+ reasonable results with various model flavors compatible with the pipeline, such as those relying on very
120
+ short denoising schedules (`LCMScheduler`) and those with full diffusion schedules (`DDIMScheduler`).
121
+ default_processing_resolution (`int`, *optional*):
122
+ The recommended value of the `processing_resolution` parameter of the pipeline. This value must be set in
123
+ the model config. When the pipeline is called without explicitly setting `processing_resolution`, the
124
+ default value is used. This is required to ensure reasonable results with various model flavors trained
125
+ with varying optimal processing resolution values.
126
+ """
127
+
128
+ model_cpu_offload_seq = "text_encoder->unet->vae"
129
+ supported_prediction_types = ("normals",)
130
+
131
+ def __init__(
132
+ self,
133
+ unet: UNet2DConditionModel,
134
+ vae: AutoencoderKL,
135
+ scheduler: Union[DDIMScheduler, LCMScheduler],
136
+ text_encoder: CLIPTextModel,
137
+ tokenizer: CLIPTokenizer,
138
+ prediction_type: Optional[str] = None,
139
+ use_full_z_range: Optional[bool] = True,
140
+ default_denoising_steps: Optional[int] = None,
141
+ default_processing_resolution: Optional[int] = None,
142
+ ):
143
+ super().__init__()
144
+
145
+ if prediction_type not in self.supported_prediction_types:
146
+ logger.warning(
147
+ f"Potentially unsupported `prediction_type='{prediction_type}'`; values supported by the pipeline: "
148
+ f"{self.supported_prediction_types}."
149
+ )
150
+
151
+ self.register_modules(
152
+ unet=unet,
153
+ vae=vae,
154
+ scheduler=scheduler,
155
+ text_encoder=text_encoder,
156
+ tokenizer=tokenizer,
157
+ )
158
+ self.register_to_config(
159
+ use_full_z_range=use_full_z_range,
160
+ default_denoising_steps=default_denoising_steps,
161
+ default_processing_resolution=default_processing_resolution,
162
+ )
163
+
164
+ self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
165
+
166
+ self.use_full_z_range = use_full_z_range
167
+ self.default_denoising_steps = default_denoising_steps
168
+ self.default_processing_resolution = default_processing_resolution
169
+
170
+ self.empty_text_embedding = None
171
+
172
+ self.image_processor = MarigoldImageProcessor(vae_scale_factor=self.vae_scale_factor)
173
+
174
+ def check_inputs(
175
+ self,
176
+ image: PipelineImageInput,
177
+ num_inference_steps: int,
178
+ ensemble_size: int,
179
+ processing_resolution: int,
180
+ resample_method_input: str,
181
+ resample_method_output: str,
182
+ batch_size: int,
183
+ ensembling_kwargs: Optional[Dict[str, Any]],
184
+ latents: Optional[torch.Tensor],
185
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]],
186
+ output_type: str,
187
+ output_uncertainty: bool,
188
+ ) -> int:
189
+ if num_inference_steps is None:
190
+ raise ValueError("`num_inference_steps` is not specified and could not be resolved from the model config.")
191
+ if num_inference_steps < 1:
192
+ raise ValueError("`num_inference_steps` must be positive.")
193
+ if ensemble_size < 1:
194
+ raise ValueError("`ensemble_size` must be positive.")
195
+ if ensemble_size == 2:
196
+ logger.warning(
197
+ "`ensemble_size` == 2 results are similar to no ensembling (1); "
198
+ "consider increasing the value to at least 3."
199
+ )
200
+ if ensemble_size == 1 and output_uncertainty:
201
+ raise ValueError(
202
+ "Computing uncertainty by setting `output_uncertainty=True` also requires setting `ensemble_size` "
203
+ "greater than 1."
204
+ )
205
+ if processing_resolution is None:
206
+ raise ValueError(
207
+ "`processing_resolution` is not specified and could not be resolved from the model config."
208
+ )
209
+ if processing_resolution < 0:
210
+ raise ValueError(
211
+ "`processing_resolution` must be non-negative: 0 for native resolution, or any positive value for "
212
+ "downsampled processing."
213
+ )
214
+ if processing_resolution % self.vae_scale_factor != 0:
215
+ raise ValueError(f"`processing_resolution` must be a multiple of {self.vae_scale_factor}.")
216
+ if resample_method_input not in ("nearest", "nearest-exact", "bilinear", "bicubic", "area"):
217
+ raise ValueError(
218
+ "`resample_method_input` takes string values compatible with PIL library: "
219
+ "nearest, nearest-exact, bilinear, bicubic, area."
220
+ )
221
+ if resample_method_output not in ("nearest", "nearest-exact", "bilinear", "bicubic", "area"):
222
+ raise ValueError(
223
+ "`resample_method_output` takes string values compatible with PIL library: "
224
+ "nearest, nearest-exact, bilinear, bicubic, area."
225
+ )
226
+ if batch_size < 1:
227
+ raise ValueError("`batch_size` must be positive.")
228
+ if output_type not in ["pt", "np"]:
229
+ raise ValueError("`output_type` must be one of `pt` or `np`.")
230
+ if latents is not None and generator is not None:
231
+ raise ValueError("`latents` and `generator` cannot be used together.")
232
+ if ensembling_kwargs is not None:
233
+ if not isinstance(ensembling_kwargs, dict):
234
+ raise ValueError("`ensembling_kwargs` must be a dictionary.")
235
+ if "reduction" in ensembling_kwargs and ensembling_kwargs["reduction"] not in ("closest", "mean"):
236
+ raise ValueError("`ensembling_kwargs['reduction']` can be either `'closest'` or `'mean'`.")
237
+
238
+ # image checks
239
+ num_images = 0
240
+ W, H = None, None
241
+ if not isinstance(image, list):
242
+ image = [image]
243
+ for i, img in enumerate(image):
244
+ if isinstance(img, np.ndarray) or torch.is_tensor(img):
245
+ if img.ndim not in (2, 3, 4):
246
+ raise ValueError(f"`image[{i}]` has unsupported dimensions or shape: {img.shape}.")
247
+ H_i, W_i = img.shape[-2:]
248
+ N_i = 1
249
+ if img.ndim == 4:
250
+ N_i = img.shape[0]
251
+ elif isinstance(img, Image.Image):
252
+ W_i, H_i = img.size
253
+ N_i = 1
254
+ else:
255
+ raise ValueError(f"Unsupported `image[{i}]` type: {type(img)}.")
256
+ if W is None:
257
+ W, H = W_i, H_i
258
+ elif (W, H) != (W_i, H_i):
259
+ raise ValueError(
260
+ f"Input `image[{i}]` has incompatible dimensions {(W_i, H_i)} with the previous images {(W, H)}"
261
+ )
262
+ num_images += N_i
263
+
264
+ # latents checks
265
+ if latents is not None:
266
+ if not torch.is_tensor(latents):
267
+ raise ValueError("`latents` must be a torch.Tensor.")
268
+ if latents.dim() != 4:
269
+ raise ValueError(f"`latents` has unsupported dimensions or shape: {latents.shape}.")
270
+
271
+ if processing_resolution > 0:
272
+ max_orig = max(H, W)
273
+ new_H = H * processing_resolution // max_orig
274
+ new_W = W * processing_resolution // max_orig
275
+ if new_H == 0 or new_W == 0:
276
+ raise ValueError(f"Extreme aspect ratio of the input image: [{W} x {H}]")
277
+ W, H = new_W, new_H
278
+ w = (W + self.vae_scale_factor - 1) // self.vae_scale_factor
279
+ h = (H + self.vae_scale_factor - 1) // self.vae_scale_factor
280
+ shape_expected = (num_images * ensemble_size, self.vae.config.latent_channels, h, w)
281
+
282
+ if latents.shape != shape_expected:
283
+ raise ValueError(f"`latents` has unexpected shape={latents.shape} expected={shape_expected}.")
284
+
285
+ # generator checks
286
+ if generator is not None:
287
+ if isinstance(generator, list):
288
+ if len(generator) != num_images * ensemble_size:
289
+ raise ValueError(
290
+ "The number of generators must match the total number of ensemble members for all input images."
291
+ )
292
+ if not all(g.device.type == generator[0].device.type for g in generator):
293
+ raise ValueError("`generator` device placement is not consistent in the list.")
294
+ elif not isinstance(generator, torch.Generator):
295
+ raise ValueError(f"Unsupported generator type: {type(generator)}.")
296
+
297
+ return num_images
298
+
299
+ def progress_bar(self, iterable=None, total=None, desc=None, leave=True):
300
+ if not hasattr(self, "_progress_bar_config"):
301
+ self._progress_bar_config = {}
302
+ elif not isinstance(self._progress_bar_config, dict):
303
+ raise ValueError(
304
+ f"`self._progress_bar_config` should be of type `dict`, but is {type(self._progress_bar_config)}."
305
+ )
306
+
307
+ progress_bar_config = dict(**self._progress_bar_config)
308
+ progress_bar_config["desc"] = progress_bar_config.get("desc", desc)
309
+ progress_bar_config["leave"] = progress_bar_config.get("leave", leave)
310
+ if iterable is not None:
311
+ return tqdm(iterable, **progress_bar_config)
312
+ elif total is not None:
313
+ return tqdm(total=total, **progress_bar_config)
314
+ else:
315
+ raise ValueError("Either `total` or `iterable` has to be defined.")
316
+
317
+ @torch.no_grad()
318
+ @replace_example_docstring(EXAMPLE_DOC_STRING)
319
+ def __call__(
320
+ self,
321
+ image: PipelineImageInput,
322
+ num_inference_steps: Optional[int] = None,
323
+ ensemble_size: int = 1,
324
+ processing_resolution: Optional[int] = None,
325
+ match_input_resolution: bool = True,
326
+ resample_method_input: str = "bilinear",
327
+ resample_method_output: str = "bilinear",
328
+ batch_size: int = 1,
329
+ ensembling_kwargs: Optional[Dict[str, Any]] = None,
330
+ latents: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None,
331
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
332
+ output_type: str = "np",
333
+ output_uncertainty: bool = False,
334
+ output_latent: bool = False,
335
+ return_dict: bool = True,
336
+ ):
337
+ """
338
+ Function invoked when calling the pipeline.
339
+
340
+ Args:
341
+ image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`),
342
+ `List[torch.Tensor]`: An input image or images used as an input for the normals estimation task. For
343
+ arrays and tensors, the expected value range is between `[0, 1]`. Passing a batch of images is possible
344
+ by providing a four-dimensional array or a tensor. Additionally, a list of images of two- or
345
+ three-dimensional arrays or tensors can be passed. In the latter case, all list elements must have the
346
+ same width and height.
347
+ num_inference_steps (`int`, *optional*, defaults to `None`):
348
+ Number of denoising diffusion steps during inference. The default value `None` results in automatic
349
+ selection. The number of steps should be at least 10 with the full Marigold models, and between 1 and 4
350
+ for Marigold-LCM models.
351
+ ensemble_size (`int`, defaults to `1`):
352
+ Number of ensemble predictions. Recommended values are 5 and higher for better precision, or 1 for
353
+ faster inference.
354
+ processing_resolution (`int`, *optional*, defaults to `None`):
355
+ Effective processing resolution. When set to `0`, matches the larger input image dimension. This
356
+ produces crisper predictions, but may also lead to the overall loss of global context. The default
357
+ value `None` resolves to the optimal value from the model config.
358
+ match_input_resolution (`bool`, *optional*, defaults to `True`):
359
+ When enabled, the output prediction is resized to match the input dimensions. When disabled, the longer
360
+ side of the output will equal to `processing_resolution`.
361
+ resample_method_input (`str`, *optional*, defaults to `"bilinear"`):
362
+ Resampling method used to resize input images to `processing_resolution`. The accepted values are:
363
+ `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`.
364
+ resample_method_output (`str`, *optional*, defaults to `"bilinear"`):
365
+ Resampling method used to resize output predictions to match the input resolution. The accepted values
366
+ are `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`.
367
+ batch_size (`int`, *optional*, defaults to `1`):
368
+ Batch size; only matters when setting `ensemble_size` or passing a tensor of images.
369
+ ensembling_kwargs (`dict`, *optional*, defaults to `None`)
370
+ Extra dictionary with arguments for precise ensembling control. The following options are available:
371
+ - reduction (`str`, *optional*, defaults to `"closest"`): Defines the ensembling function applied in
372
+ every pixel location, can be either `"closest"` or `"mean"`.
373
+ latents (`torch.Tensor`, *optional*, defaults to `None`):
374
+ Latent noise tensors to replace the random initialization. These can be taken from the previous
375
+ function call's output.
376
+ generator (`torch.Generator`, or `List[torch.Generator]`, *optional*, defaults to `None`):
377
+ Random number generator object to ensure reproducibility.
378
+ output_type (`str`, *optional*, defaults to `"np"`):
379
+ Preferred format of the output's `prediction` and the optional `uncertainty` fields. The accepted
380
+ values are: `"np"` (numpy array) or `"pt"` (torch tensor).
381
+ output_uncertainty (`bool`, *optional*, defaults to `False`):
382
+ When enabled, the output's `uncertainty` field contains the predictive uncertainty map, provided that
383
+ the `ensemble_size` argument is set to a value above 2.
384
+ output_latent (`bool`, *optional*, defaults to `False`):
385
+ When enabled, the output's `latent` field contains the latent codes corresponding to the predictions
386
+ within the ensemble. These codes can be saved, modified, and used for subsequent calls with the
387
+ `latents` argument.
388
+ return_dict (`bool`, *optional*, defaults to `True`):
389
+ Whether or not to return a [`~pipelines.marigold.MarigoldDepthOutput`] instead of a plain tuple.
390
+
391
+ Examples:
392
+
393
+ Returns:
394
+ [`~pipelines.marigold.MarigoldNormalsOutput`] or `tuple`:
395
+ If `return_dict` is `True`, [`~pipelines.marigold.MarigoldNormalsOutput`] is returned, otherwise a
396
+ `tuple` is returned where the first element is the prediction, the second element is the uncertainty
397
+ (or `None`), and the third is the latent (or `None`).
398
+ """
399
+
400
+ # 0. Resolving variables.
401
+ device = self._execution_device
402
+ dtype = self.dtype
403
+
404
+ # Model-specific optimal default values leading to fast and reasonable results.
405
+ if num_inference_steps is None:
406
+ num_inference_steps = self.default_denoising_steps
407
+ if processing_resolution is None:
408
+ processing_resolution = self.default_processing_resolution
409
+
410
+ # 1. Check inputs.
411
+ num_images = self.check_inputs(
412
+ image,
413
+ num_inference_steps,
414
+ ensemble_size,
415
+ processing_resolution,
416
+ resample_method_input,
417
+ resample_method_output,
418
+ batch_size,
419
+ ensembling_kwargs,
420
+ latents,
421
+ generator,
422
+ output_type,
423
+ output_uncertainty,
424
+ )
425
+
426
+ # 2. Prepare empty text conditioning.
427
+ # Model invocation: self.tokenizer, self.text_encoder.
428
+ if self.empty_text_embedding is None:
429
+ prompt = ""
430
+ text_inputs = self.tokenizer(
431
+ prompt,
432
+ padding="do_not_pad",
433
+ max_length=self.tokenizer.model_max_length,
434
+ truncation=True,
435
+ return_tensors="pt",
436
+ )
437
+ text_input_ids = text_inputs.input_ids.to(device)
438
+ self.empty_text_embedding = self.text_encoder(text_input_ids)[0] # [1,2,1024]
439
+
440
+ # 3. Preprocess input images. This function loads input image or images of compatible dimensions `(H, W)`,
441
+ # optionally downsamples them to the `processing_resolution` `(PH, PW)`, where
442
+ # `max(PH, PW) == processing_resolution`, and pads the dimensions to `(PPH, PPW)` such that these values are
443
+ # divisible by the latent space downscaling factor (typically 8 in Stable Diffusion). The default value `None`
444
+ # of `processing_resolution` resolves to the optimal value from the model config. It is a recommended mode of
445
+ # operation and leads to the most reasonable results. Using the native image resolution or any other processing
446
+ # resolution can lead to loss of either fine details or global context in the output predictions.
447
+ image, padding, original_resolution = self.image_processor.preprocess(
448
+ image, processing_resolution, resample_method_input, device, dtype
449
+ ) # [N,3,PPH,PPW]
450
+
451
+ # 4. Encode input image into latent space. At this step, each of the `N` input images is represented with `E`
452
+ # ensemble members. Each ensemble member is an independent diffused prediction, just initialized independently.
453
+ # Latents of each such predictions across all input images and all ensemble members are represented in the
454
+ # `pred_latent` variable. The variable `image_latent` is of the same shape: it contains each input image encoded
455
+ # into latent space and replicated `E` times. The latents can be either generated (see `generator` to ensure
456
+ # reproducibility), or passed explicitly via the `latents` argument. The latter can be set outside the pipeline
457
+ # code. For example, in the Marigold-LCM video processing demo, the latents initialization of a frame is taken
458
+ # as a convex combination of the latents output of the pipeline for the previous frame and a newly-sampled
459
+ # noise. This behavior can be achieved by setting the `output_latent` argument to `True`. The latent space
460
+ # dimensions are `(h, w)`. Encoding into latent space happens in batches of size `batch_size`.
461
+ # Model invocation: self.vae.encoder.
462
+ image_latent, pred_latent = self.prepare_latents(
463
+ image, latents, generator, ensemble_size, batch_size
464
+ ) # [N*E,4,h,w], [N*E,4,h,w]
465
+
466
+ del image
467
+
468
+ batch_empty_text_embedding = self.empty_text_embedding.to(device=device, dtype=dtype).repeat(
469
+ batch_size, 1, 1
470
+ ) # [B,1024,2]
471
+
472
+ # 5. Process the denoising loop. All `N * E` latents are processed sequentially in batches of size `batch_size`.
473
+ # The unet model takes concatenated latent spaces of the input image and the predicted modality as an input, and
474
+ # outputs noise for the predicted modality's latent space. The number of denoising diffusion steps is defined by
475
+ # `num_inference_steps`. It is either set directly, or resolves to the optimal value specific to the loaded
476
+ # model.
477
+ # Model invocation: self.unet.
478
+ pred_latents = []
479
+
480
+ for i in self.progress_bar(
481
+ range(0, num_images * ensemble_size, batch_size), leave=True, desc="Marigold predictions..."
482
+ ):
483
+ batch_image_latent = image_latent[i : i + batch_size] # [B,4,h,w]
484
+ batch_pred_latent = pred_latent[i : i + batch_size] # [B,4,h,w]
485
+ effective_batch_size = batch_image_latent.shape[0]
486
+ text = batch_empty_text_embedding[:effective_batch_size] # [B,2,1024]
487
+
488
+ self.scheduler.set_timesteps(num_inference_steps, device=device)
489
+ for t in self.progress_bar(self.scheduler.timesteps, leave=False, desc="Diffusion steps..."):
490
+ batch_latent = torch.cat([batch_image_latent, batch_pred_latent], dim=1) # [B,8,h,w]
491
+ noise = self.unet(batch_latent, t, encoder_hidden_states=text, return_dict=False)[0] # [B,4,h,w]
492
+ batch_pred_latent = self.scheduler.step(
493
+ noise, t, batch_pred_latent, generator=generator
494
+ ).prev_sample # [B,4,h,w]
495
+
496
+ pred_latents.append(batch_pred_latent)
497
+
498
+ pred_latent = torch.cat(pred_latents, dim=0) # [N*E,4,h,w]
499
+
500
+ del (
501
+ pred_latents,
502
+ image_latent,
503
+ batch_empty_text_embedding,
504
+ batch_image_latent,
505
+ batch_pred_latent,
506
+ text,
507
+ batch_latent,
508
+ noise,
509
+ )
510
+
511
+ # 6. Decode predictions from latent into pixel space. The resulting `N * E` predictions have shape `(PPH, PPW)`,
512
+ # which requires slight postprocessing. Decoding into pixel space happens in batches of size `batch_size`.
513
+ # Model invocation: self.vae.decoder.
514
+ prediction = torch.cat(
515
+ [
516
+ self.decode_prediction(pred_latent[i : i + batch_size])
517
+ for i in range(0, pred_latent.shape[0], batch_size)
518
+ ],
519
+ dim=0,
520
+ ) # [N*E,3,PPH,PPW]
521
+
522
+ if not output_latent:
523
+ pred_latent = None
524
+
525
+ # 7. Remove padding. The output shape is (PH, PW).
526
+ prediction = self.image_processor.unpad_image(prediction, padding) # [N*E,3,PH,PW]
527
+
528
+ # 8. Ensemble and compute uncertainty (when `output_uncertainty` is set). This code treats each of the `N`
529
+ # groups of `E` ensemble predictions independently. For each group it computes an ensembled prediction of shape
530
+ # `(PH, PW)` and an optional uncertainty map of the same dimensions. After computing this pair of outputs for
531
+ # each group independently, it stacks them respectively into batches of `N` almost final predictions and
532
+ # uncertainty maps.
533
+ uncertainty = None
534
+ if ensemble_size > 1:
535
+ prediction = prediction.reshape(num_images, ensemble_size, *prediction.shape[1:]) # [N,E,3,PH,PW]
536
+ prediction = [
537
+ self.ensemble_normals(prediction[i], output_uncertainty, **(ensembling_kwargs or {}))
538
+ for i in range(num_images)
539
+ ] # [ [[1,3,PH,PW], [1,1,PH,PW]], ... ]
540
+ prediction, uncertainty = zip(*prediction) # [[1,3,PH,PW], ... ], [[1,1,PH,PW], ... ]
541
+ prediction = torch.cat(prediction, dim=0) # [N,3,PH,PW]
542
+ if output_uncertainty:
543
+ uncertainty = torch.cat(uncertainty, dim=0) # [N,1,PH,PW]
544
+ else:
545
+ uncertainty = None
546
+
547
+ # 9. If `match_input_resolution` is set, the output prediction and the uncertainty are upsampled to match the
548
+ # input resolution `(H, W)`. This step may introduce upsampling artifacts, and therefore can be disabled.
549
+ # After upsampling, the native resolution normal maps are renormalized to unit length to reduce the artifacts.
550
+ # Depending on the downstream use-case, upsampling can be also chosen based on the tolerated artifacts by
551
+ # setting the `resample_method_output` parameter (e.g., to `"nearest"`).
552
+ if match_input_resolution:
553
+ prediction = self.image_processor.resize_antialias(
554
+ prediction, original_resolution, resample_method_output, is_aa=False
555
+ ) # [N,3,H,W]
556
+ prediction = self.normalize_normals(prediction) # [N,3,H,W]
557
+ if uncertainty is not None and output_uncertainty:
558
+ uncertainty = self.image_processor.resize_antialias(
559
+ uncertainty, original_resolution, resample_method_output, is_aa=False
560
+ ) # [N,1,H,W]
561
+
562
+ # 10. Prepare the final outputs.
563
+ if output_type == "np":
564
+ prediction = self.image_processor.pt_to_numpy(prediction) # [N,H,W,3]
565
+ if uncertainty is not None and output_uncertainty:
566
+ uncertainty = self.image_processor.pt_to_numpy(uncertainty) # [N,H,W,1]
567
+
568
+ # 11. Offload all models
569
+ self.maybe_free_model_hooks()
570
+
571
+ if not return_dict:
572
+ return (prediction, uncertainty, pred_latent)
573
+
574
+ return MarigoldNormalsOutput(
575
+ prediction=prediction,
576
+ uncertainty=uncertainty,
577
+ latent=pred_latent,
578
+ )
579
+
580
+ # Copied from diffusers.pipelines.marigold.pipeline_marigold_depth.MarigoldDepthPipeline.prepare_latents
581
+ def prepare_latents(
582
+ self,
583
+ image: torch.Tensor,
584
+ latents: Optional[torch.Tensor],
585
+ generator: Optional[torch.Generator],
586
+ ensemble_size: int,
587
+ batch_size: int,
588
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
589
+ def retrieve_latents(encoder_output):
590
+ if hasattr(encoder_output, "latent_dist"):
591
+ return encoder_output.latent_dist.mode()
592
+ elif hasattr(encoder_output, "latents"):
593
+ return encoder_output.latents
594
+ else:
595
+ raise AttributeError("Could not access latents of provided encoder_output")
596
+
597
+ image_latent = torch.cat(
598
+ [
599
+ retrieve_latents(self.vae.encode(image[i : i + batch_size]))
600
+ for i in range(0, image.shape[0], batch_size)
601
+ ],
602
+ dim=0,
603
+ ) # [N,4,h,w]
604
+ image_latent = image_latent * self.vae.config.scaling_factor
605
+ image_latent = image_latent.repeat_interleave(ensemble_size, dim=0) # [N*E,4,h,w]
606
+
607
+ pred_latent = latents
608
+ if pred_latent is None:
609
+ pred_latent = randn_tensor(
610
+ image_latent.shape,
611
+ generator=generator,
612
+ device=image_latent.device,
613
+ dtype=image_latent.dtype,
614
+ ) # [N*E,4,h,w]
615
+
616
+ return image_latent, pred_latent
617
+
618
+ def decode_prediction(self, pred_latent: torch.Tensor) -> torch.Tensor:
619
+ if pred_latent.dim() != 4 or pred_latent.shape[1] != self.vae.config.latent_channels:
620
+ raise ValueError(
621
+ f"Expecting 4D tensor of shape [B,{self.vae.config.latent_channels},H,W]; got {pred_latent.shape}."
622
+ )
623
+
624
+ prediction = self.vae.decode(pred_latent / self.vae.config.scaling_factor, return_dict=False)[0] # [B,3,H,W]
625
+
626
+ prediction = torch.clip(prediction, -1.0, 1.0)
627
+
628
+ if not self.use_full_z_range:
629
+ prediction[:, 2, :, :] *= 0.5
630
+ prediction[:, 2, :, :] += 0.5
631
+
632
+ prediction = self.normalize_normals(prediction) # [B,3,H,W]
633
+
634
+ return prediction # [B,3,H,W]
635
+
636
+ @staticmethod
637
+ def normalize_normals(normals: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
638
+ if normals.dim() != 4 or normals.shape[1] != 3:
639
+ raise ValueError(f"Expecting 4D tensor of shape [B,3,H,W]; got {normals.shape}.")
640
+
641
+ norm = torch.norm(normals, dim=1, keepdim=True)
642
+ normals /= norm.clamp(min=eps)
643
+
644
+ return normals
645
+
646
+ @staticmethod
647
+ def ensemble_normals(
648
+ normals: torch.Tensor, output_uncertainty: bool, reduction: str = "closest"
649
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
650
+ """
651
+ Ensembles the normals maps represented by the `normals` tensor with expected shape `(B, 3, H, W)`, where B is
652
+ the number of ensemble members for a given prediction of size `(H x W)`.
653
+
654
+ Args:
655
+ normals (`torch.Tensor`):
656
+ Input ensemble normals maps.
657
+ output_uncertainty (`bool`, *optional*, defaults to `False`):
658
+ Whether to output uncertainty map.
659
+ reduction (`str`, *optional*, defaults to `"closest"`):
660
+ Reduction method used to ensemble aligned predictions. The accepted values are: `"closest"` and
661
+ `"mean"`.
662
+
663
+ Returns:
664
+ A tensor of aligned and ensembled normals maps with shape `(1, 3, H, W)` and optionally a tensor of
665
+ uncertainties of shape `(1, 1, H, W)`.
666
+ """
667
+ if normals.dim() != 4 or normals.shape[1] != 3:
668
+ raise ValueError(f"Expecting 4D tensor of shape [B,3,H,W]; got {normals.shape}.")
669
+ if reduction not in ("closest", "mean"):
670
+ raise ValueError(f"Unrecognized reduction method: {reduction}.")
671
+
672
+ mean_normals = normals.mean(dim=0, keepdim=True) # [1,3,H,W]
673
+ mean_normals = MarigoldNormalsPipeline.normalize_normals(mean_normals) # [1,3,H,W]
674
+
675
+ sim_cos = (mean_normals * normals).sum(dim=1, keepdim=True) # [E,1,H,W]
676
+ sim_cos = sim_cos.clamp(-1, 1) # required to avoid NaN in uncertainty with fp16
677
+
678
+ uncertainty = None
679
+ if output_uncertainty:
680
+ uncertainty = sim_cos.arccos() # [E,1,H,W]
681
+ uncertainty = uncertainty.mean(dim=0, keepdim=True) / np.pi # [1,1,H,W]
682
+
683
+ if reduction == "mean":
684
+ return mean_normals, uncertainty # [1,3,H,W], [1,1,H,W]
685
+
686
+ closest_indices = sim_cos.argmax(dim=0, keepdim=True) # [1,1,H,W]
687
+ closest_indices = closest_indices.repeat(1, 3, 1, 1) # [1,3,H,W]
688
+ closest_normals = torch.gather(normals, 0, closest_indices) # [1,3,H,W]
689
+
690
+ return closest_normals, uncertainty # [1,3,H,W], [1,1,H,W]