diff --git a/examples/dreambooth/train_dreambooth_lora_flux2_img2img.py b/examples/dreambooth/train_dreambooth_lora_flux2_img2img.py index 2ba694405c48..c305afaff04b 100644 --- a/examples/dreambooth/train_dreambooth_lora_flux2_img2img.py +++ b/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) diff --git a/examples/dreambooth/train_dreambooth_lora_flux2_klein_img2img.py b/examples/dreambooth/train_dreambooth_lora_flux2_klein_img2img.py index df56f2557213..81f2ec9be996 100644 --- a/examples/dreambooth/train_dreambooth_lora_flux2_klein_img2img.py +++ b/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) diff --git a/src/diffusers/image_processor.py b/src/diffusers/image_processor.py index 4f6f4bd52b9c..c0e0714906a4 100644 --- a/src/diffusers/image_processor.py +++ b/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] diff --git a/src/diffusers/utils/pil_utils.py b/src/diffusers/utils/pil_utils.py index 72d4704fa945..625b834f1bfe 100644 --- a/src/diffusers/utils/pil_utils.py +++ b/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] diff --git a/tests/others/test_image_processor.py b/tests/others/test_image_processor.py index 0d358699f105..0021d36c817f 100644 --- a/tests/others/test_image_processor.py +++ b/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)