diffusers 0.27.1__py3-none-any.whl → 0.28.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (270) hide show
  1. diffusers/__init__.py +18 -1
  2. diffusers/callbacks.py +156 -0
  3. diffusers/commands/env.py +110 -6
  4. diffusers/configuration_utils.py +16 -11
  5. diffusers/dependency_versions_table.py +2 -1
  6. diffusers/image_processor.py +158 -45
  7. diffusers/loaders/__init__.py +2 -5
  8. diffusers/loaders/autoencoder.py +4 -4
  9. diffusers/loaders/controlnet.py +4 -4
  10. diffusers/loaders/ip_adapter.py +80 -22
  11. diffusers/loaders/lora.py +134 -20
  12. diffusers/loaders/lora_conversion_utils.py +46 -43
  13. diffusers/loaders/peft.py +4 -3
  14. diffusers/loaders/single_file.py +401 -170
  15. diffusers/loaders/single_file_model.py +290 -0
  16. diffusers/loaders/single_file_utils.py +616 -672
  17. diffusers/loaders/textual_inversion.py +41 -20
  18. diffusers/loaders/unet.py +168 -115
  19. diffusers/loaders/unet_loader_utils.py +163 -0
  20. diffusers/models/__init__.py +2 -0
  21. diffusers/models/activations.py +11 -3
  22. diffusers/models/attention.py +10 -11
  23. diffusers/models/attention_processor.py +367 -148
  24. diffusers/models/autoencoders/autoencoder_asym_kl.py +14 -16
  25. diffusers/models/autoencoders/autoencoder_kl.py +18 -19
  26. diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py +11 -11
  27. diffusers/models/autoencoders/autoencoder_tiny.py +16 -16
  28. diffusers/models/autoencoders/consistency_decoder_vae.py +36 -11
  29. diffusers/models/autoencoders/vae.py +23 -24
  30. diffusers/models/controlnet.py +12 -9
  31. diffusers/models/controlnet_flax.py +4 -4
  32. diffusers/models/controlnet_xs.py +1915 -0
  33. diffusers/models/downsampling.py +17 -18
  34. diffusers/models/embeddings.py +147 -24
  35. diffusers/models/model_loading_utils.py +149 -0
  36. diffusers/models/modeling_flax_pytorch_utils.py +2 -1
  37. diffusers/models/modeling_flax_utils.py +4 -4
  38. diffusers/models/modeling_pytorch_flax_utils.py +1 -1
  39. diffusers/models/modeling_utils.py +118 -98
  40. diffusers/models/resnet.py +18 -23
  41. diffusers/models/transformer_temporal.py +3 -3
  42. diffusers/models/transformers/dual_transformer_2d.py +4 -4
  43. diffusers/models/transformers/prior_transformer.py +7 -7
  44. diffusers/models/transformers/t5_film_transformer.py +17 -19
  45. diffusers/models/transformers/transformer_2d.py +272 -156
  46. diffusers/models/transformers/transformer_temporal.py +10 -10
  47. diffusers/models/unets/unet_1d.py +5 -5
  48. diffusers/models/unets/unet_1d_blocks.py +29 -29
  49. diffusers/models/unets/unet_2d.py +6 -6
  50. diffusers/models/unets/unet_2d_blocks.py +137 -128
  51. diffusers/models/unets/unet_2d_condition.py +20 -15
  52. diffusers/models/unets/unet_2d_condition_flax.py +6 -5
  53. diffusers/models/unets/unet_3d_blocks.py +79 -77
  54. diffusers/models/unets/unet_3d_condition.py +13 -9
  55. diffusers/models/unets/unet_i2vgen_xl.py +14 -13
  56. diffusers/models/unets/unet_kandinsky3.py +1 -1
  57. diffusers/models/unets/unet_motion_model.py +114 -14
  58. diffusers/models/unets/unet_spatio_temporal_condition.py +15 -14
  59. diffusers/models/unets/unet_stable_cascade.py +16 -13
  60. diffusers/models/upsampling.py +17 -20
  61. diffusers/models/vq_model.py +16 -15
  62. diffusers/pipelines/__init__.py +25 -3
  63. diffusers/pipelines/amused/pipeline_amused.py +12 -12
  64. diffusers/pipelines/amused/pipeline_amused_img2img.py +14 -12
  65. diffusers/pipelines/amused/pipeline_amused_inpaint.py +13 -11
  66. diffusers/pipelines/animatediff/__init__.py +2 -0
  67. diffusers/pipelines/animatediff/pipeline_animatediff.py +24 -46
  68. diffusers/pipelines/animatediff/pipeline_animatediff_sdxl.py +1284 -0
  69. diffusers/pipelines/animatediff/pipeline_animatediff_video2video.py +82 -72
  70. diffusers/pipelines/animatediff/pipeline_output.py +3 -2
  71. diffusers/pipelines/audioldm/pipeline_audioldm.py +14 -14
  72. diffusers/pipelines/audioldm2/modeling_audioldm2.py +54 -35
  73. diffusers/pipelines/audioldm2/pipeline_audioldm2.py +120 -36
  74. diffusers/pipelines/auto_pipeline.py +21 -17
  75. diffusers/pipelines/blip_diffusion/blip_image_processing.py +1 -1
  76. diffusers/pipelines/blip_diffusion/modeling_blip2.py +5 -5
  77. diffusers/pipelines/blip_diffusion/modeling_ctx_clip.py +1 -1
  78. diffusers/pipelines/blip_diffusion/pipeline_blip_diffusion.py +2 -2
  79. diffusers/pipelines/consistency_models/pipeline_consistency_models.py +5 -5
  80. diffusers/pipelines/controlnet/multicontrolnet.py +4 -8
  81. diffusers/pipelines/controlnet/pipeline_controlnet.py +87 -52
  82. diffusers/pipelines/controlnet/pipeline_controlnet_blip_diffusion.py +2 -2
  83. diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py +50 -43
  84. diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py +52 -40
  85. diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py +80 -47
  86. diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py +147 -49
  87. diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py +89 -55
  88. diffusers/pipelines/controlnet_xs/__init__.py +68 -0
  89. diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs.py +911 -0
  90. diffusers/pipelines/controlnet_xs/pipeline_controlnet_xs_sd_xl.py +1115 -0
  91. diffusers/pipelines/deepfloyd_if/pipeline_if.py +14 -28
  92. diffusers/pipelines/deepfloyd_if/pipeline_if_img2img.py +18 -33
  93. diffusers/pipelines/deepfloyd_if/pipeline_if_img2img_superresolution.py +21 -39
  94. diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting.py +20 -36
  95. diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting_superresolution.py +23 -39
  96. diffusers/pipelines/deepfloyd_if/pipeline_if_superresolution.py +17 -32
  97. diffusers/pipelines/deprecated/alt_diffusion/modeling_roberta_series.py +11 -11
  98. diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion.py +43 -20
  99. diffusers/pipelines/deprecated/alt_diffusion/pipeline_alt_diffusion_img2img.py +36 -18
  100. diffusers/pipelines/deprecated/repaint/pipeline_repaint.py +2 -2
  101. diffusers/pipelines/deprecated/spectrogram_diffusion/pipeline_spectrogram_diffusion.py +7 -7
  102. diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_cycle_diffusion.py +12 -12
  103. diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_inpaint_legacy.py +18 -21
  104. diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_model_editing.py +20 -15
  105. diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_paradigms.py +20 -15
  106. diffusers/pipelines/deprecated/stable_diffusion_variants/pipeline_stable_diffusion_pix2pix_zero.py +30 -25
  107. diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py +69 -59
  108. diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion.py +13 -13
  109. diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_dual_guided.py +10 -5
  110. diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_image_variation.py +11 -6
  111. diffusers/pipelines/deprecated/versatile_diffusion/pipeline_versatile_diffusion_text_to_image.py +10 -5
  112. diffusers/pipelines/deprecated/vq_diffusion/pipeline_vq_diffusion.py +5 -5
  113. diffusers/pipelines/dit/pipeline_dit.py +3 -0
  114. diffusers/pipelines/free_init_utils.py +39 -38
  115. diffusers/pipelines/i2vgen_xl/pipeline_i2vgen_xl.py +33 -48
  116. diffusers/pipelines/kandinsky/pipeline_kandinsky.py +8 -8
  117. diffusers/pipelines/kandinsky/pipeline_kandinsky_combined.py +23 -20
  118. diffusers/pipelines/kandinsky/pipeline_kandinsky_img2img.py +11 -11
  119. diffusers/pipelines/kandinsky/pipeline_kandinsky_inpaint.py +12 -12
  120. diffusers/pipelines/kandinsky/pipeline_kandinsky_prior.py +10 -10
  121. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py +6 -6
  122. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py +32 -29
  123. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet.py +10 -10
  124. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_controlnet_img2img.py +10 -10
  125. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_img2img.py +6 -6
  126. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_inpainting.py +8 -8
  127. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py +7 -7
  128. diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior_emb2emb.py +6 -6
  129. diffusers/pipelines/kandinsky3/convert_kandinsky3_unet.py +3 -3
  130. diffusers/pipelines/kandinsky3/pipeline_kandinsky3.py +20 -33
  131. diffusers/pipelines/kandinsky3/pipeline_kandinsky3_img2img.py +24 -35
  132. diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_img2img.py +48 -30
  133. diffusers/pipelines/latent_consistency_models/pipeline_latent_consistency_text2img.py +50 -28
  134. diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py +11 -11
  135. diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion.py +61 -67
  136. diffusers/pipelines/ledits_pp/pipeline_leditspp_stable_diffusion_xl.py +70 -69
  137. diffusers/pipelines/ledits_pp/pipeline_output.py +2 -2
  138. diffusers/pipelines/marigold/__init__.py +50 -0
  139. diffusers/pipelines/marigold/marigold_image_processing.py +561 -0
  140. diffusers/pipelines/marigold/pipeline_marigold_depth.py +813 -0
  141. diffusers/pipelines/marigold/pipeline_marigold_normals.py +690 -0
  142. diffusers/pipelines/musicldm/pipeline_musicldm.py +14 -14
  143. diffusers/pipelines/paint_by_example/pipeline_paint_by_example.py +17 -12
  144. diffusers/pipelines/pia/pipeline_pia.py +39 -125
  145. diffusers/pipelines/pipeline_flax_utils.py +4 -4
  146. diffusers/pipelines/pipeline_loading_utils.py +268 -23
  147. diffusers/pipelines/pipeline_utils.py +266 -37
  148. diffusers/pipelines/pixart_alpha/__init__.py +8 -1
  149. diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py +65 -75
  150. diffusers/pipelines/pixart_alpha/pipeline_pixart_sigma.py +880 -0
  151. diffusers/pipelines/semantic_stable_diffusion/pipeline_semantic_stable_diffusion.py +10 -5
  152. diffusers/pipelines/shap_e/pipeline_shap_e.py +3 -3
  153. diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py +14 -14
  154. diffusers/pipelines/shap_e/renderer.py +1 -1
  155. diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py +36 -22
  156. diffusers/pipelines/stable_cascade/pipeline_stable_cascade_combined.py +23 -19
  157. diffusers/pipelines/stable_cascade/pipeline_stable_cascade_prior.py +33 -32
  158. diffusers/pipelines/stable_diffusion/__init__.py +0 -1
  159. diffusers/pipelines/stable_diffusion/convert_from_ckpt.py +18 -11
  160. diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py +2 -2
  161. diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_upscale.py +6 -6
  162. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +73 -39
  163. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_depth2img.py +24 -17
  164. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py +13 -8
  165. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py +66 -36
  166. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py +82 -46
  167. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_instruct_pix2pix.py +123 -28
  168. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_latent_upscale.py +6 -6
  169. diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py +16 -16
  170. diffusers/pipelines/stable_diffusion/pipeline_stable_unclip.py +24 -19
  171. diffusers/pipelines/stable_diffusion/pipeline_stable_unclip_img2img.py +37 -31
  172. diffusers/pipelines/stable_diffusion/safety_checker.py +2 -1
  173. diffusers/pipelines/stable_diffusion_attend_and_excite/pipeline_stable_diffusion_attend_and_excite.py +23 -15
  174. diffusers/pipelines/stable_diffusion_diffedit/pipeline_stable_diffusion_diffedit.py +44 -42
  175. diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen.py +23 -18
  176. diffusers/pipelines/stable_diffusion_gligen/pipeline_stable_diffusion_gligen_text_image.py +19 -14
  177. diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_k_diffusion.py +20 -15
  178. diffusers/pipelines/stable_diffusion_k_diffusion/pipeline_stable_diffusion_xl_k_diffusion.py +24 -19
  179. diffusers/pipelines/stable_diffusion_ldm3d/pipeline_stable_diffusion_ldm3d.py +65 -32
  180. diffusers/pipelines/stable_diffusion_panorama/pipeline_stable_diffusion_panorama.py +274 -38
  181. diffusers/pipelines/stable_diffusion_safe/pipeline_stable_diffusion_safe.py +10 -5
  182. diffusers/pipelines/stable_diffusion_safe/safety_checker.py +1 -1
  183. diffusers/pipelines/stable_diffusion_sag/pipeline_stable_diffusion_sag.py +92 -25
  184. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py +88 -44
  185. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py +108 -56
  186. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_inpaint.py +96 -51
  187. diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_instruct_pix2pix.py +45 -25
  188. diffusers/pipelines/stable_diffusion_xl/watermark.py +9 -3
  189. diffusers/pipelines/stable_video_diffusion/pipeline_stable_video_diffusion.py +110 -57
  190. diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_adapter.py +59 -30
  191. diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_xl_adapter.py +71 -42
  192. diffusers/pipelines/text_to_video_synthesis/pipeline_output.py +3 -2
  193. diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth.py +18 -41
  194. diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth_img2img.py +21 -85
  195. diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero.py +28 -19
  196. diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_zero_sdxl.py +39 -33
  197. diffusers/pipelines/unclip/pipeline_unclip.py +6 -6
  198. diffusers/pipelines/unclip/pipeline_unclip_image_variation.py +6 -6
  199. diffusers/pipelines/unidiffuser/modeling_text_decoder.py +1 -1
  200. diffusers/pipelines/unidiffuser/modeling_uvit.py +9 -9
  201. diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py +23 -23
  202. diffusers/pipelines/wuerstchen/modeling_paella_vq_model.py +5 -5
  203. diffusers/pipelines/wuerstchen/modeling_wuerstchen_common.py +5 -10
  204. diffusers/pipelines/wuerstchen/modeling_wuerstchen_prior.py +4 -6
  205. diffusers/pipelines/wuerstchen/pipeline_wuerstchen.py +4 -4
  206. diffusers/pipelines/wuerstchen/pipeline_wuerstchen_combined.py +12 -12
  207. diffusers/pipelines/wuerstchen/pipeline_wuerstchen_prior.py +10 -10
  208. diffusers/schedulers/__init__.py +2 -2
  209. diffusers/schedulers/deprecated/__init__.py +1 -1
  210. diffusers/schedulers/deprecated/scheduling_karras_ve.py +25 -25
  211. diffusers/schedulers/scheduling_amused.py +5 -5
  212. diffusers/schedulers/scheduling_consistency_decoder.py +11 -11
  213. diffusers/schedulers/scheduling_consistency_models.py +23 -25
  214. diffusers/schedulers/scheduling_ddim.py +22 -24
  215. diffusers/schedulers/scheduling_ddim_flax.py +2 -1
  216. diffusers/schedulers/scheduling_ddim_inverse.py +16 -16
  217. diffusers/schedulers/scheduling_ddim_parallel.py +28 -30
  218. diffusers/schedulers/scheduling_ddpm.py +20 -22
  219. diffusers/schedulers/scheduling_ddpm_flax.py +7 -3
  220. diffusers/schedulers/scheduling_ddpm_parallel.py +26 -28
  221. diffusers/schedulers/scheduling_ddpm_wuerstchen.py +14 -14
  222. diffusers/schedulers/scheduling_deis_multistep.py +46 -42
  223. diffusers/schedulers/scheduling_dpmsolver_multistep.py +107 -77
  224. diffusers/schedulers/scheduling_dpmsolver_multistep_flax.py +2 -2
  225. diffusers/schedulers/scheduling_dpmsolver_multistep_inverse.py +46 -46
  226. diffusers/schedulers/scheduling_dpmsolver_sde.py +26 -22
  227. diffusers/schedulers/scheduling_dpmsolver_singlestep.py +90 -65
  228. diffusers/schedulers/scheduling_edm_dpmsolver_multistep.py +78 -53
  229. diffusers/schedulers/scheduling_edm_euler.py +53 -30
  230. diffusers/schedulers/scheduling_euler_ancestral_discrete.py +26 -28
  231. diffusers/schedulers/scheduling_euler_discrete.py +163 -67
  232. diffusers/schedulers/scheduling_heun_discrete.py +60 -38
  233. diffusers/schedulers/scheduling_ipndm.py +8 -8
  234. diffusers/schedulers/scheduling_k_dpm_2_ancestral_discrete.py +22 -18
  235. diffusers/schedulers/scheduling_k_dpm_2_discrete.py +22 -18
  236. diffusers/schedulers/scheduling_karras_ve_flax.py +6 -6
  237. diffusers/schedulers/scheduling_lcm.py +21 -23
  238. diffusers/schedulers/scheduling_lms_discrete.py +27 -25
  239. diffusers/schedulers/scheduling_pndm.py +20 -20
  240. diffusers/schedulers/scheduling_repaint.py +20 -20
  241. diffusers/schedulers/scheduling_sasolver.py +55 -54
  242. diffusers/schedulers/scheduling_sde_ve.py +19 -19
  243. diffusers/schedulers/scheduling_tcd.py +39 -30
  244. diffusers/schedulers/scheduling_unclip.py +15 -15
  245. diffusers/schedulers/scheduling_unipc_multistep.py +115 -41
  246. diffusers/schedulers/scheduling_utils.py +14 -5
  247. diffusers/schedulers/scheduling_utils_flax.py +3 -3
  248. diffusers/schedulers/scheduling_vq_diffusion.py +10 -10
  249. diffusers/training_utils.py +56 -1
  250. diffusers/utils/__init__.py +7 -0
  251. diffusers/utils/doc_utils.py +1 -0
  252. diffusers/utils/dummy_pt_objects.py +30 -0
  253. diffusers/utils/dummy_torch_and_transformers_objects.py +90 -0
  254. diffusers/utils/dynamic_modules_utils.py +24 -11
  255. diffusers/utils/hub_utils.py +3 -2
  256. diffusers/utils/import_utils.py +91 -0
  257. diffusers/utils/loading_utils.py +2 -2
  258. diffusers/utils/logging.py +1 -1
  259. diffusers/utils/peft_utils.py +32 -5
  260. diffusers/utils/state_dict_utils.py +11 -2
  261. diffusers/utils/testing_utils.py +71 -6
  262. diffusers/utils/torch_utils.py +1 -0
  263. diffusers/video_processor.py +113 -0
  264. {diffusers-0.27.1.dist-info → diffusers-0.28.0.dist-info}/METADATA +7 -7
  265. diffusers-0.28.0.dist-info/RECORD +414 -0
  266. diffusers-0.27.1.dist-info/RECORD +0 -399
  267. {diffusers-0.27.1.dist-info → diffusers-0.28.0.dist-info}/LICENSE +0 -0
  268. {diffusers-0.27.1.dist-info → diffusers-0.28.0.dist-info}/WHEEL +0 -0
  269. {diffusers-0.27.1.dist-info → diffusers-0.28.0.dist-info}/entry_points.txt +0 -0
  270. {diffusers-0.27.1.dist-info → diffusers-0.28.0.dist-info}/top_level.txt +0 -0
@@ -35,8 +35,8 @@ class LEditsPPInversionPipelineOutput(BaseOutput):
35
35
  List of the cropped and resized input images as PIL images of length `batch_size` or NumPy array of shape `
36
36
  (batch_size, height, width, num_channels)`.
37
37
  vae_reconstruction_images (`List[PIL.Image.Image]` or `np.ndarray`)
38
- List of VAE reconstruction of all input images as PIL images of length `batch_size` or NumPy array of shape `
39
- (batch_size, height, width, num_channels)`.
38
+ List of VAE reconstruction of all input images as PIL images of length `batch_size` or NumPy array of shape
39
+ ` (batch_size, height, width, num_channels)`.
40
40
  """
41
41
 
42
42
  images: Union[List[PIL.Image.Image], np.ndarray]
@@ -0,0 +1,50 @@
1
+ from typing import TYPE_CHECKING
2
+
3
+ from ...utils import (
4
+ DIFFUSERS_SLOW_IMPORT,
5
+ OptionalDependencyNotAvailable,
6
+ _LazyModule,
7
+ get_objects_from_module,
8
+ is_torch_available,
9
+ is_transformers_available,
10
+ )
11
+
12
+
13
+ _dummy_objects = {}
14
+ _import_structure = {}
15
+
16
+ try:
17
+ if not (is_transformers_available() and is_torch_available()):
18
+ raise OptionalDependencyNotAvailable()
19
+ except OptionalDependencyNotAvailable:
20
+ from ...utils import dummy_torch_and_transformers_objects # noqa F403
21
+
22
+ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
23
+ else:
24
+ _import_structure["marigold_image_processing"] = ["MarigoldImageProcessor"]
25
+ _import_structure["pipeline_marigold_depth"] = ["MarigoldDepthOutput", "MarigoldDepthPipeline"]
26
+ _import_structure["pipeline_marigold_normals"] = ["MarigoldNormalsOutput", "MarigoldNormalsPipeline"]
27
+
28
+ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
29
+ try:
30
+ if not (is_transformers_available() and is_torch_available()):
31
+ raise OptionalDependencyNotAvailable()
32
+
33
+ except OptionalDependencyNotAvailable:
34
+ from ...utils.dummy_torch_and_transformers_objects import *
35
+ else:
36
+ from .marigold_image_processing import MarigoldImageProcessor
37
+ from .pipeline_marigold_depth import MarigoldDepthOutput, MarigoldDepthPipeline
38
+ from .pipeline_marigold_normals import MarigoldNormalsOutput, MarigoldNormalsPipeline
39
+
40
+ else:
41
+ import sys
42
+
43
+ sys.modules[__name__] = _LazyModule(
44
+ __name__,
45
+ globals()["__file__"],
46
+ _import_structure,
47
+ module_spec=__spec__,
48
+ )
49
+ for name, value in _dummy_objects.items():
50
+ setattr(sys.modules[__name__], name, value)
@@ -0,0 +1,561 @@
1
+ from typing import List, Optional, Tuple, Union
2
+
3
+ import numpy as np
4
+ import PIL
5
+ import torch
6
+ import torch.nn.functional as F
7
+ from PIL import Image
8
+
9
+ from ... import ConfigMixin
10
+ from ...configuration_utils import register_to_config
11
+ from ...image_processor import PipelineImageInput
12
+ from ...utils import CONFIG_NAME, logging
13
+ from ...utils.import_utils import is_matplotlib_available
14
+
15
+
16
+ logger = logging.get_logger(__name__) # pylint: disable=invalid-name
17
+
18
+
19
+ class MarigoldImageProcessor(ConfigMixin):
20
+ config_name = CONFIG_NAME
21
+
22
+ @register_to_config
23
+ def __init__(
24
+ self,
25
+ vae_scale_factor: int = 8,
26
+ do_normalize: bool = True,
27
+ do_range_check: bool = True,
28
+ ):
29
+ super().__init__()
30
+
31
+ @staticmethod
32
+ def expand_tensor_or_array(images: Union[torch.Tensor, np.ndarray]) -> Union[torch.Tensor, np.ndarray]:
33
+ """
34
+ Expand a tensor or array to a specified number of images.
35
+ """
36
+ if isinstance(images, np.ndarray):
37
+ if images.ndim == 2: # [H,W] -> [1,H,W,1]
38
+ images = images[None, ..., None]
39
+ if images.ndim == 3: # [H,W,C] -> [1,H,W,C]
40
+ images = images[None]
41
+ elif isinstance(images, torch.Tensor):
42
+ if images.ndim == 2: # [H,W] -> [1,1,H,W]
43
+ images = images[None, None]
44
+ elif images.ndim == 3: # [1,H,W] -> [1,1,H,W]
45
+ images = images[None]
46
+ else:
47
+ raise ValueError(f"Unexpected input type: {type(images)}")
48
+ return images
49
+
50
+ @staticmethod
51
+ def pt_to_numpy(images: torch.Tensor) -> np.ndarray:
52
+ """
53
+ Convert a PyTorch tensor to a NumPy image.
54
+ """
55
+ images = images.cpu().permute(0, 2, 3, 1).float().numpy()
56
+ return images
57
+
58
+ @staticmethod
59
+ def numpy_to_pt(images: np.ndarray) -> torch.Tensor:
60
+ """
61
+ Convert a NumPy image to a PyTorch tensor.
62
+ """
63
+ if np.issubdtype(images.dtype, np.integer) and not np.issubdtype(images.dtype, np.unsignedinteger):
64
+ raise ValueError(f"Input image dtype={images.dtype} cannot be a signed integer.")
65
+ if np.issubdtype(images.dtype, np.complexfloating):
66
+ raise ValueError(f"Input image dtype={images.dtype} cannot be complex.")
67
+ if np.issubdtype(images.dtype, bool):
68
+ raise ValueError(f"Input image dtype={images.dtype} cannot be boolean.")
69
+
70
+ images = torch.from_numpy(images.transpose(0, 3, 1, 2))
71
+ return images
72
+
73
+ @staticmethod
74
+ def resize_antialias(
75
+ image: torch.Tensor, size: Tuple[int, int], mode: str, is_aa: Optional[bool] = None
76
+ ) -> torch.Tensor:
77
+ if not torch.is_tensor(image):
78
+ raise ValueError(f"Invalid input type={type(image)}.")
79
+ if not torch.is_floating_point(image):
80
+ raise ValueError(f"Invalid input dtype={image.dtype}.")
81
+ if image.dim() != 4:
82
+ raise ValueError(f"Invalid input dimensions; shape={image.shape}.")
83
+
84
+ antialias = is_aa and mode in ("bilinear", "bicubic")
85
+ image = F.interpolate(image, size, mode=mode, antialias=antialias)
86
+
87
+ return image
88
+
89
+ @staticmethod
90
+ def resize_to_max_edge(image: torch.Tensor, max_edge_sz: int, mode: str) -> torch.Tensor:
91
+ if not torch.is_tensor(image):
92
+ raise ValueError(f"Invalid input type={type(image)}.")
93
+ if not torch.is_floating_point(image):
94
+ raise ValueError(f"Invalid input dtype={image.dtype}.")
95
+ if image.dim() != 4:
96
+ raise ValueError(f"Invalid input dimensions; shape={image.shape}.")
97
+
98
+ h, w = image.shape[-2:]
99
+ max_orig = max(h, w)
100
+ new_h = h * max_edge_sz // max_orig
101
+ new_w = w * max_edge_sz // max_orig
102
+
103
+ if new_h == 0 or new_w == 0:
104
+ raise ValueError(f"Extreme aspect ratio of the input image: [{w} x {h}]")
105
+
106
+ image = MarigoldImageProcessor.resize_antialias(image, (new_h, new_w), mode, is_aa=True)
107
+
108
+ return image
109
+
110
+ @staticmethod
111
+ def pad_image(image: torch.Tensor, align: int) -> Tuple[torch.Tensor, Tuple[int, int]]:
112
+ if not torch.is_tensor(image):
113
+ raise ValueError(f"Invalid input type={type(image)}.")
114
+ if not torch.is_floating_point(image):
115
+ raise ValueError(f"Invalid input dtype={image.dtype}.")
116
+ if image.dim() != 4:
117
+ raise ValueError(f"Invalid input dimensions; shape={image.shape}.")
118
+
119
+ h, w = image.shape[-2:]
120
+ ph, pw = -h % align, -w % align
121
+
122
+ image = F.pad(image, (0, pw, 0, ph), mode="replicate")
123
+
124
+ return image, (ph, pw)
125
+
126
+ @staticmethod
127
+ def unpad_image(image: torch.Tensor, padding: Tuple[int, int]) -> torch.Tensor:
128
+ if not torch.is_tensor(image):
129
+ raise ValueError(f"Invalid input type={type(image)}.")
130
+ if not torch.is_floating_point(image):
131
+ raise ValueError(f"Invalid input dtype={image.dtype}.")
132
+ if image.dim() != 4:
133
+ raise ValueError(f"Invalid input dimensions; shape={image.shape}.")
134
+
135
+ ph, pw = padding
136
+ uh = None if ph == 0 else -ph
137
+ uw = None if pw == 0 else -pw
138
+
139
+ image = image[:, :, :uh, :uw]
140
+
141
+ return image
142
+
143
+ @staticmethod
144
+ def load_image_canonical(
145
+ image: Union[torch.Tensor, np.ndarray, Image.Image],
146
+ device: torch.device = torch.device("cpu"),
147
+ dtype: torch.dtype = torch.float32,
148
+ ) -> Tuple[torch.Tensor, int]:
149
+ if isinstance(image, Image.Image):
150
+ image = np.array(image)
151
+
152
+ image_dtype_max = None
153
+ if isinstance(image, (np.ndarray, torch.Tensor)):
154
+ image = MarigoldImageProcessor.expand_tensor_or_array(image)
155
+ if image.ndim != 4:
156
+ raise ValueError("Input image is not 2-, 3-, or 4-dimensional.")
157
+ if isinstance(image, np.ndarray):
158
+ if np.issubdtype(image.dtype, np.integer) and not np.issubdtype(image.dtype, np.unsignedinteger):
159
+ raise ValueError(f"Input image dtype={image.dtype} cannot be a signed integer.")
160
+ if np.issubdtype(image.dtype, np.complexfloating):
161
+ raise ValueError(f"Input image dtype={image.dtype} cannot be complex.")
162
+ if np.issubdtype(image.dtype, bool):
163
+ raise ValueError(f"Input image dtype={image.dtype} cannot be boolean.")
164
+ if np.issubdtype(image.dtype, np.unsignedinteger):
165
+ image_dtype_max = np.iinfo(image.dtype).max
166
+ image = image.astype(np.float32) # because torch does not have unsigned dtypes beyond torch.uint8
167
+ image = MarigoldImageProcessor.numpy_to_pt(image)
168
+
169
+ if torch.is_tensor(image) and not torch.is_floating_point(image) and image_dtype_max is None:
170
+ if image.dtype != torch.uint8:
171
+ raise ValueError(f"Image dtype={image.dtype} is not supported.")
172
+ image_dtype_max = 255
173
+
174
+ if not torch.is_tensor(image):
175
+ raise ValueError(f"Input type unsupported: {type(image)}.")
176
+
177
+ if image.shape[1] == 1:
178
+ image = image.repeat(1, 3, 1, 1) # [N,1,H,W] -> [N,3,H,W]
179
+ if image.shape[1] != 3:
180
+ raise ValueError(f"Input image is not 1- or 3-channel: {image.shape}.")
181
+
182
+ image = image.to(device=device, dtype=dtype)
183
+
184
+ if image_dtype_max is not None:
185
+ image = image / image_dtype_max
186
+
187
+ return image
188
+
189
+ @staticmethod
190
+ def check_image_values_range(image: torch.Tensor) -> None:
191
+ if not torch.is_tensor(image):
192
+ raise ValueError(f"Invalid input type={type(image)}.")
193
+ if not torch.is_floating_point(image):
194
+ raise ValueError(f"Invalid input dtype={image.dtype}.")
195
+ if image.min().item() < 0.0 or image.max().item() > 1.0:
196
+ raise ValueError("Input image data is partially outside of the [0,1] range.")
197
+
198
+ def preprocess(
199
+ self,
200
+ image: PipelineImageInput,
201
+ processing_resolution: Optional[int] = None,
202
+ resample_method_input: str = "bilinear",
203
+ device: torch.device = torch.device("cpu"),
204
+ dtype: torch.dtype = torch.float32,
205
+ ):
206
+ if isinstance(image, list):
207
+ images = None
208
+ for i, img in enumerate(image):
209
+ img = self.load_image_canonical(img, device, dtype) # [N,3,H,W]
210
+ if images is None:
211
+ images = img
212
+ else:
213
+ if images.shape[2:] != img.shape[2:]:
214
+ raise ValueError(
215
+ f"Input image[{i}] has incompatible dimensions {img.shape[2:]} with the previous images "
216
+ f"{images.shape[2:]}"
217
+ )
218
+ images = torch.cat((images, img), dim=0)
219
+ image = images
220
+ del images
221
+ else:
222
+ image = self.load_image_canonical(image, device, dtype) # [N,3,H,W]
223
+
224
+ original_resolution = image.shape[2:]
225
+
226
+ if self.config.do_range_check:
227
+ self.check_image_values_range(image)
228
+
229
+ if self.config.do_normalize:
230
+ image = image * 2.0 - 1.0
231
+
232
+ if processing_resolution is not None and processing_resolution > 0:
233
+ image = self.resize_to_max_edge(image, processing_resolution, resample_method_input) # [N,3,PH,PW]
234
+
235
+ image, padding = self.pad_image(image, self.config.vae_scale_factor) # [N,3,PPH,PPW]
236
+
237
+ return image, padding, original_resolution
238
+
239
+ @staticmethod
240
+ def colormap(
241
+ image: Union[np.ndarray, torch.Tensor],
242
+ cmap: str = "Spectral",
243
+ bytes: bool = False,
244
+ _force_method: Optional[str] = None,
245
+ ) -> Union[np.ndarray, torch.Tensor]:
246
+ """
247
+ Converts a monochrome image into an RGB image by applying the specified colormap. This function mimics the
248
+ behavior of matplotlib.colormaps, but allows the user to use the most discriminative color map "Spectral"
249
+ without having to install or import matplotlib. For all other cases, the function will attempt to use the
250
+ native implementation.
251
+
252
+ Args:
253
+ image: 2D tensor of values between 0 and 1, either as np.ndarray or torch.Tensor.
254
+ cmap: Colormap name.
255
+ bytes: Whether to return the output as uint8 or floating point image.
256
+ _force_method:
257
+ Can be used to specify whether to use the native implementation (`"matplotlib"`), the efficient custom
258
+ implementation of the "Spectral" color map (`"custom"`), or rely on autodetection (`None`, default).
259
+
260
+ Returns:
261
+ An RGB-colorized tensor corresponding to the input image.
262
+ """
263
+ if not (torch.is_tensor(image) or isinstance(image, np.ndarray)):
264
+ raise ValueError("Argument must be a numpy array or torch tensor.")
265
+ if _force_method not in (None, "matplotlib", "custom"):
266
+ raise ValueError("_force_method must be either `None`, `'matplotlib'` or `'custom'`.")
267
+
268
+ def method_matplotlib(image, cmap, bytes=False):
269
+ if is_matplotlib_available():
270
+ import matplotlib
271
+ else:
272
+ return None
273
+
274
+ arg_is_pt, device = torch.is_tensor(image), None
275
+ if arg_is_pt:
276
+ image, device = image.cpu().numpy(), image.device
277
+
278
+ if cmap not in matplotlib.colormaps:
279
+ raise ValueError(
280
+ f"Unexpected color map {cmap}; available options are: {', '.join(list(matplotlib.colormaps.keys()))}"
281
+ )
282
+
283
+ cmap = matplotlib.colormaps[cmap]
284
+ out = cmap(image, bytes=bytes) # [?,4]
285
+ out = out[..., :3] # [?,3]
286
+
287
+ if arg_is_pt:
288
+ out = torch.tensor(out, device=device)
289
+
290
+ return out
291
+
292
+ def method_custom(image, cmap, bytes=False):
293
+ arg_is_np = isinstance(image, np.ndarray)
294
+ if arg_is_np:
295
+ image = torch.tensor(image)
296
+ if image.dtype == torch.uint8:
297
+ image = image.float() / 255
298
+ else:
299
+ image = image.float()
300
+
301
+ if cmap != "Spectral":
302
+ raise ValueError("Only 'Spectral' color map is available without installing matplotlib.")
303
+
304
+ _Spectral_data = ( # Taken from matplotlib/_cm.py
305
+ (0.61960784313725492, 0.003921568627450980, 0.25882352941176473), # 0.0 -> [0]
306
+ (0.83529411764705885, 0.24313725490196078, 0.30980392156862746),
307
+ (0.95686274509803926, 0.42745098039215684, 0.2627450980392157),
308
+ (0.99215686274509807, 0.68235294117647061, 0.38039215686274508),
309
+ (0.99607843137254903, 0.8784313725490196, 0.54509803921568623),
310
+ (1.0, 1.0, 0.74901960784313726),
311
+ (0.90196078431372551, 0.96078431372549022, 0.59607843137254901),
312
+ (0.6705882352941176, 0.8666666666666667, 0.64313725490196083),
313
+ (0.4, 0.76078431372549016, 0.6470588235294118),
314
+ (0.19607843137254902, 0.53333333333333333, 0.74117647058823533),
315
+ (0.36862745098039218, 0.30980392156862746, 0.63529411764705879), # 1.0 -> [K-1]
316
+ )
317
+
318
+ cmap = torch.tensor(_Spectral_data, dtype=torch.float, device=image.device) # [K,3]
319
+ K = cmap.shape[0]
320
+
321
+ pos = image.clamp(min=0, max=1) * (K - 1)
322
+ left = pos.long()
323
+ right = (left + 1).clamp(max=K - 1)
324
+
325
+ d = (pos - left.float()).unsqueeze(-1)
326
+ left_colors = cmap[left]
327
+ right_colors = cmap[right]
328
+
329
+ out = (1 - d) * left_colors + d * right_colors
330
+
331
+ if bytes:
332
+ out = (out * 255).to(torch.uint8)
333
+
334
+ if arg_is_np:
335
+ out = out.numpy()
336
+
337
+ return out
338
+
339
+ if _force_method is None and torch.is_tensor(image) and cmap == "Spectral":
340
+ return method_custom(image, cmap, bytes)
341
+
342
+ out = None
343
+ if _force_method != "custom":
344
+ out = method_matplotlib(image, cmap, bytes)
345
+
346
+ if _force_method == "matplotlib" and out is None:
347
+ raise ImportError("Make sure to install matplotlib if you want to use a color map other than 'Spectral'.")
348
+
349
+ if out is None:
350
+ out = method_custom(image, cmap, bytes)
351
+
352
+ return out
353
+
354
+ @staticmethod
355
+ def visualize_depth(
356
+ depth: Union[
357
+ PIL.Image.Image,
358
+ np.ndarray,
359
+ torch.Tensor,
360
+ List[PIL.Image.Image],
361
+ List[np.ndarray],
362
+ List[torch.Tensor],
363
+ ],
364
+ val_min: float = 0.0,
365
+ val_max: float = 1.0,
366
+ color_map: str = "Spectral",
367
+ ) -> Union[PIL.Image.Image, List[PIL.Image.Image]]:
368
+ """
369
+ Visualizes depth maps, such as predictions of the `MarigoldDepthPipeline`.
370
+
371
+ Args:
372
+ depth (`Union[PIL.Image.Image, np.ndarray, torch.Tensor, List[PIL.Image.Image], List[np.ndarray],
373
+ List[torch.Tensor]]`): Depth maps.
374
+ val_min (`float`, *optional*, defaults to `0.0`): Minimum value of the visualized depth range.
375
+ val_max (`float`, *optional*, defaults to `1.0`): Maximum value of the visualized depth range.
376
+ color_map (`str`, *optional*, defaults to `"Spectral"`): Color map used to convert a single-channel
377
+ depth prediction into colored representation.
378
+
379
+ Returns: `PIL.Image.Image` or `List[PIL.Image.Image]` with depth maps visualization.
380
+ """
381
+ if val_max <= val_min:
382
+ raise ValueError(f"Invalid values range: [{val_min}, {val_max}].")
383
+
384
+ def visualize_depth_one(img, idx=None):
385
+ prefix = "Depth" + (f"[{idx}]" if idx else "")
386
+ if isinstance(img, PIL.Image.Image):
387
+ if img.mode != "I;16":
388
+ raise ValueError(f"{prefix}: invalid PIL mode={img.mode}.")
389
+ img = np.array(img).astype(np.float32) / (2**16 - 1)
390
+ if isinstance(img, np.ndarray) or torch.is_tensor(img):
391
+ if img.ndim != 2:
392
+ raise ValueError(f"{prefix}: unexpected shape={img.shape}.")
393
+ if isinstance(img, np.ndarray):
394
+ img = torch.from_numpy(img)
395
+ if not torch.is_floating_point(img):
396
+ raise ValueError(f"{prefix}: unexected dtype={img.dtype}.")
397
+ else:
398
+ raise ValueError(f"{prefix}: unexpected type={type(img)}.")
399
+ if val_min != 0.0 or val_max != 1.0:
400
+ img = (img - val_min) / (val_max - val_min)
401
+ img = MarigoldImageProcessor.colormap(img, cmap=color_map, bytes=True) # [H,W,3]
402
+ img = PIL.Image.fromarray(img.cpu().numpy())
403
+ return img
404
+
405
+ if depth is None or isinstance(depth, list) and any(o is None for o in depth):
406
+ raise ValueError("Input depth is `None`")
407
+ if isinstance(depth, (np.ndarray, torch.Tensor)):
408
+ depth = MarigoldImageProcessor.expand_tensor_or_array(depth)
409
+ if isinstance(depth, np.ndarray):
410
+ depth = MarigoldImageProcessor.numpy_to_pt(depth) # [N,H,W,1] -> [N,1,H,W]
411
+ if not (depth.ndim == 4 and depth.shape[1] == 1): # [N,1,H,W]
412
+ raise ValueError(f"Unexpected input shape={depth.shape}, expecting [N,1,H,W].")
413
+ return [visualize_depth_one(img[0], idx) for idx, img in enumerate(depth)]
414
+ elif isinstance(depth, list):
415
+ return [visualize_depth_one(img, idx) for idx, img in enumerate(depth)]
416
+ else:
417
+ raise ValueError(f"Unexpected input type: {type(depth)}")
418
+
419
+ @staticmethod
420
+ def export_depth_to_16bit_png(
421
+ depth: Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tensor]],
422
+ val_min: float = 0.0,
423
+ val_max: float = 1.0,
424
+ ) -> Union[PIL.Image.Image, List[PIL.Image.Image]]:
425
+ def export_depth_to_16bit_png_one(img, idx=None):
426
+ prefix = "Depth" + (f"[{idx}]" if idx else "")
427
+ if not isinstance(img, np.ndarray) and not torch.is_tensor(img):
428
+ raise ValueError(f"{prefix}: unexpected type={type(img)}.")
429
+ if img.ndim != 2:
430
+ raise ValueError(f"{prefix}: unexpected shape={img.shape}.")
431
+ if torch.is_tensor(img):
432
+ img = img.cpu().numpy()
433
+ if not np.issubdtype(img.dtype, np.floating):
434
+ raise ValueError(f"{prefix}: unexected dtype={img.dtype}.")
435
+ if val_min != 0.0 or val_max != 1.0:
436
+ img = (img - val_min) / (val_max - val_min)
437
+ img = (img * (2**16 - 1)).astype(np.uint16)
438
+ img = PIL.Image.fromarray(img, mode="I;16")
439
+ return img
440
+
441
+ if depth is None or isinstance(depth, list) and any(o is None for o in depth):
442
+ raise ValueError("Input depth is `None`")
443
+ if isinstance(depth, (np.ndarray, torch.Tensor)):
444
+ depth = MarigoldImageProcessor.expand_tensor_or_array(depth)
445
+ if isinstance(depth, np.ndarray):
446
+ depth = MarigoldImageProcessor.numpy_to_pt(depth) # [N,H,W,1] -> [N,1,H,W]
447
+ if not (depth.ndim == 4 and depth.shape[1] == 1):
448
+ raise ValueError(f"Unexpected input shape={depth.shape}, expecting [N,1,H,W].")
449
+ return [export_depth_to_16bit_png_one(img[0], idx) for idx, img in enumerate(depth)]
450
+ elif isinstance(depth, list):
451
+ return [export_depth_to_16bit_png_one(img, idx) for idx, img in enumerate(depth)]
452
+ else:
453
+ raise ValueError(f"Unexpected input type: {type(depth)}")
454
+
455
+ @staticmethod
456
+ def visualize_normals(
457
+ normals: Union[
458
+ np.ndarray,
459
+ torch.Tensor,
460
+ List[np.ndarray],
461
+ List[torch.Tensor],
462
+ ],
463
+ flip_x: bool = False,
464
+ flip_y: bool = False,
465
+ flip_z: bool = False,
466
+ ) -> Union[PIL.Image.Image, List[PIL.Image.Image]]:
467
+ """
468
+ Visualizes surface normals, such as predictions of the `MarigoldNormalsPipeline`.
469
+
470
+ Args:
471
+ normals (`Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tensor]]`):
472
+ Surface normals.
473
+ flip_x (`bool`, *optional*, defaults to `False`): Flips the X axis of the normals frame of reference.
474
+ Default direction is right.
475
+ flip_y (`bool`, *optional*, defaults to `False`): Flips the Y axis of the normals frame of reference.
476
+ Default direction is top.
477
+ flip_z (`bool`, *optional*, defaults to `False`): Flips the Z axis of the normals frame of reference.
478
+ Default direction is facing the observer.
479
+
480
+ Returns: `PIL.Image.Image` or `List[PIL.Image.Image]` with surface normals visualization.
481
+ """
482
+ flip_vec = None
483
+ if any((flip_x, flip_y, flip_z)):
484
+ flip_vec = torch.tensor(
485
+ [
486
+ (-1) ** flip_x,
487
+ (-1) ** flip_y,
488
+ (-1) ** flip_z,
489
+ ],
490
+ dtype=torch.float32,
491
+ )
492
+
493
+ def visualize_normals_one(img, idx=None):
494
+ img = img.permute(1, 2, 0)
495
+ if flip_vec is not None:
496
+ img *= flip_vec.to(img.device)
497
+ img = (img + 1.0) * 0.5
498
+ img = (img * 255).to(dtype=torch.uint8, device="cpu").numpy()
499
+ img = PIL.Image.fromarray(img)
500
+ return img
501
+
502
+ if normals is None or isinstance(normals, list) and any(o is None for o in normals):
503
+ raise ValueError("Input normals is `None`")
504
+ if isinstance(normals, (np.ndarray, torch.Tensor)):
505
+ normals = MarigoldImageProcessor.expand_tensor_or_array(normals)
506
+ if isinstance(normals, np.ndarray):
507
+ normals = MarigoldImageProcessor.numpy_to_pt(normals) # [N,3,H,W]
508
+ if not (normals.ndim == 4 and normals.shape[1] == 3):
509
+ raise ValueError(f"Unexpected input shape={normals.shape}, expecting [N,3,H,W].")
510
+ return [visualize_normals_one(img, idx) for idx, img in enumerate(normals)]
511
+ elif isinstance(normals, list):
512
+ return [visualize_normals_one(img, idx) for idx, img in enumerate(normals)]
513
+ else:
514
+ raise ValueError(f"Unexpected input type: {type(normals)}")
515
+
516
+ @staticmethod
517
+ def visualize_uncertainty(
518
+ uncertainty: Union[
519
+ np.ndarray,
520
+ torch.Tensor,
521
+ List[np.ndarray],
522
+ List[torch.Tensor],
523
+ ],
524
+ saturation_percentile=95,
525
+ ) -> Union[PIL.Image.Image, List[PIL.Image.Image]]:
526
+ """
527
+ Visualizes dense uncertainties, such as produced by `MarigoldDepthPipeline` or `MarigoldNormalsPipeline`.
528
+
529
+ Args:
530
+ uncertainty (`Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tensor]]`):
531
+ Uncertainty maps.
532
+ saturation_percentile (`int`, *optional*, defaults to `95`):
533
+ Specifies the percentile uncertainty value visualized with maximum intensity.
534
+
535
+ Returns: `PIL.Image.Image` or `List[PIL.Image.Image]` with uncertainty visualization.
536
+ """
537
+
538
+ def visualize_uncertainty_one(img, idx=None):
539
+ prefix = "Uncertainty" + (f"[{idx}]" if idx else "")
540
+ if img.min() < 0:
541
+ raise ValueError(f"{prefix}: unexected data range, min={img.min()}.")
542
+ img = img.squeeze(0).cpu().numpy()
543
+ saturation_value = np.percentile(img, saturation_percentile)
544
+ img = np.clip(img * 255 / saturation_value, 0, 255)
545
+ img = img.astype(np.uint8)
546
+ img = PIL.Image.fromarray(img)
547
+ return img
548
+
549
+ if uncertainty is None or isinstance(uncertainty, list) and any(o is None for o in uncertainty):
550
+ raise ValueError("Input uncertainty is `None`")
551
+ if isinstance(uncertainty, (np.ndarray, torch.Tensor)):
552
+ uncertainty = MarigoldImageProcessor.expand_tensor_or_array(uncertainty)
553
+ if isinstance(uncertainty, np.ndarray):
554
+ uncertainty = MarigoldImageProcessor.numpy_to_pt(uncertainty) # [N,1,H,W]
555
+ if not (uncertainty.ndim == 4 and uncertainty.shape[1] == 1):
556
+ raise ValueError(f"Unexpected input shape={uncertainty.shape}, expecting [N,1,H,W].")
557
+ return [visualize_uncertainty_one(img, idx) for idx, img in enumerate(uncertainty)]
558
+ elif isinstance(uncertainty, list):
559
+ return [visualize_uncertainty_one(img, idx) for idx, img in enumerate(uncertainty)]
560
+ else:
561
+ raise ValueError(f"Unexpected input type: {type(uncertainty)}")