fix: preserve singleton dimensions in numpy_to_pil

examples/dreambooth/train_dreambooth_lora_flux2_img2img.py

@@ -875,7 +875,7 @@ def __init__(
 
                 if dest_image.shape[0] == 1:
                     # Gray scale image
-                    dest_image = Image.fromarray(dest_image.squeeze().numpy(), mode="L")
+                    dest_image = Image.fromarray(dest_image.squeeze(0).numpy(), mode="L")
                 else:
                     # RGB scale image: (C, H, W) -> (H, W, C)
                     dest_image = TF.to_pil_image(dest_image)

examples/dreambooth/train_dreambooth_lora_flux2_klein_img2img.py

@@ -871,7 +871,7 @@ def __init__(
 
                 if dest_image.shape[0] == 1:
                     # Gray scale image
-                    dest_image = Image.fromarray(dest_image.squeeze().numpy(), mode="L")
+                    dest_image = Image.fromarray(dest_image.squeeze(0).numpy(), mode="L")
                 else:
                     # RGB scale image: (C, H, W) -> (H, W, C)
                     dest_image = TF.to_pil_image(dest_image)

src/diffusers/image_processor.py

@@ -142,7 +142,7 @@ def numpy_to_pil(images: np.ndarray) -> list[PIL.Image.Image]:
         images = (images * 255).round().astype("uint8")
         if images.shape[-1] == 1:
             # special case for grayscale (single channel) images
-            pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
+            pil_images = [Image.fromarray(image.squeeze(-1), mode="L") for image in images]
         else:
             pil_images = [Image.fromarray(image) for image in images]
 
@@ -1009,7 +1009,7 @@ def numpy_to_pil(images: np.ndarray) -> list[PIL.Image.Image]:
         images = (images * 255).round().astype("uint8")
         if images.shape[-1] == 1:
             # special case for grayscale (single channel) images
-            pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
+            pil_images = [Image.fromarray(image.squeeze(-1), mode="L") for image in images]
         else:
             pil_images = [Image.fromarray(image[:, :, :3]) for image in images]
 

src/diffusers/utils/pil_utils.py

@@ -41,7 +41,7 @@ def numpy_to_pil(images):
     images = (images * 255).round().astype("uint8")
     if images.shape[-1] == 1:
         # special case for grayscale (single channel) images
-        pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
+        pil_images = [Image.fromarray(image.squeeze(-1), mode="L") for image in images]
     else:
         pil_images = [Image.fromarray(image) for image in images]
 

tests/others/test_image_processor.py

@@ -95,6 +95,18 @@ def test_vae_image_processor_pil(self):
                     f"decoded output does not match input for output_type {output_type}"
                 )
 
+    def test_numpy_to_pil_singleton_dimensions(self):
+        image_processor = VaeImageProcessor(do_resize=False, do_normalize=False)
+        # (1, 7, 1) -> PIL (7, 1)
+        res = image_processor.numpy_to_pil(np.zeros((1, 7, 1), dtype=np.float32))[0]
+        assert res.size == (7, 1)
+        # (7, 1, 1) -> PIL (1, 7)
+        res = image_processor.numpy_to_pil(np.zeros((7, 1, 1), dtype=np.float32))[0]
+        assert res.size == (1, 7)
+        # (1, 1, 1) -> PIL (1, 1)
+        res = image_processor.numpy_to_pil(np.zeros((1, 1, 1), dtype=np.float32))[0]
+        assert res.size == (1, 1)
+
     def test_preprocess_input_3d(self):
         image_processor = VaeImageProcessor(do_resize=False, do_normalize=False)