diff --git a/src/diffusers/image_processor.py b/src/diffusers/image_processor.py index 4f6f4bd52b9c..9f7c2eca66a2 100644 --- a/src/diffusers/image_processor.py +++ b/src/diffusers/image_processor.py @@ -989,7 +989,12 @@ def __init__( resample: str = "lanczos", do_normalize: bool = True, ): - super().__init__() + super().__init__( + do_resize=do_resize, + vae_scale_factor=vae_scale_factor, + resample=resample, + do_normalize=do_normalize, + ) @staticmethod def numpy_to_pil(images: np.ndarray) -> list[PIL.Image.Image]: @@ -1214,7 +1219,7 @@ def preprocess( if self.config.do_resize: rgb = self.resize(rgb, height, width) - depth = np.concatenate(depth, axis=0) if rgb[0].ndim == 4 else np.stack(depth, axis=0) + depth = np.concatenate(depth, axis=0) if depth[0].ndim == 4 else np.stack(depth, axis=0) depth = self.numpy_to_pt(depth) height, width = self.get_default_height_width(depth, height, width) if self.config.do_resize: diff --git a/tests/others/test_image_processor.py b/tests/others/test_image_processor.py index 0d358699f105..472a6ca11b3f 100644 --- a/tests/others/test_image_processor.py +++ b/tests/others/test_image_processor.py @@ -17,7 +17,7 @@ import PIL.Image import torch -from diffusers.image_processor import VaeImageProcessor +from diffusers.image_processor import VaeImageProcessor, VaeImageProcessorLDM3D class TestImageProcessor: @@ -306,3 +306,26 @@ def test_vae_image_processor_resize_np(self): assert out_np.shape == exp_np_shape, ( f"resized image output shape '{out_np.shape}' didn't match expected shape '{exp_np_shape}'." ) + + def test_vae_image_processor_ldm3d_config(self): + image_processor = VaeImageProcessorLDM3D( + do_resize=False, + vae_scale_factor=4, + resample="nearest", + do_normalize=False, + ) + + assert image_processor.config.do_resize is False + assert image_processor.config.vae_scale_factor == 4 + assert image_processor.config.resample == "nearest" + assert image_processor.config.do_normalize is False + + def test_vae_image_processor_ldm3d_np_batch(self): + image_processor = VaeImageProcessorLDM3D(do_resize=False, do_normalize=False) + rgb = np.zeros((2, 8, 8, 3), dtype=np.float32) + depth = np.zeros((2, 8, 8, 1), dtype=np.float32) + + processed_rgb, processed_depth = image_processor.preprocess(rgb, depth) + + assert processed_rgb.shape == (2, 3, 8, 8) + assert processed_depth.shape == (2, 1, 8, 8)