Support grayscale images in `numpy_to_pil` (#1025)

fb38bb16 · Anton Lozhkov · GitHub · de00c632 · fb38bb16 · fb38bb16
Unverified Commit fb38bb16 authored Oct 27, 2022 by Anton Lozhkov Committed by GitHub Oct 27, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 10 additions and 2 deletions

src/diffusers/pipeline_flax_utils.py src/diffusers/pipeline_flax_utils.py +5 -1

src/diffusers/pipeline_utils.py src/diffusers/pipeline_utils.py +5 -1

No files found.
--- a/src/diffusers/pipeline_flax_utils.py
+++ b/src/diffusers/pipeline_flax_utils.py
@@ -444,7 +444,11 @@ class FlaxDiffusionPipeline(ConfigMixin):
        if images.ndim == 3:
            images = images[None, ...]
        images = (images * 255).round().astype("uint8")
-        pil_images = [Image.fromarray(image) for image in images]
+        if images.shape[-1] == 1:
+            # special case for grayscale (single channel) images
+            pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
+        else:
+            pil_images = [Image.fromarray(image) for image in images]
        return pil_images

--- a/src/diffusers/pipeline_utils.py
+++ b/src/diffusers/pipeline_utils.py
@@ -625,7 +625,11 @@ class DiffusionPipeline(ConfigMixin):
        if images.ndim == 3:
            images = images[None, ...]
        images = (images * 255).round().astype("uint8")
-        pil_images = [Image.fromarray(image) for image in images]
+        if images.shape[-1] == 1:
+            # special case for grayscale (single channel) images
+            pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
+        else:
+            pil_images = [Image.fromarray(image) for image in images]
        return pil_images