class VaeImageProcessorLDM3D(VaeImageProcessor):
    """
    Image processor for VAE LDM3D.

    Args:
        do_resize (`bool`, *optional*, defaults to `True`):
            Whether to downscale the image's (height, width) dimensions to multiples of `vae_scale_factor`.
        vae_scale_factor (`int`, *optional*, defaults to `8`):
            VAE scale factor. If `do_resize` is `True`, the image is automatically resized to multiples of this factor.
        resample (`str`, *optional*, defaults to `lanczos`):
            Resampling filter to use when resizing the image.
        do_normalize (`bool`, *optional*, defaults to `True`):
            Whether to normalize the image to [-1,1].
    """

    config_name = CONFIG_NAME

    @register_to_config
    def __init__(
        self,
        do_resize: bool = True,
        vae_scale_factor: int = 8,
        resample: str = "lanczos",
        do_normalize: bool = True,
    ):
        super().__init__()

    @staticmethod
    def numpy_to_pil(images: np.ndarray) -> List[PIL.Image.Image]:
        r"""
        Convert a NumPy image or a batch of images to a list of PIL images.

        Args:
            images (`np.ndarray`):
                The input NumPy array of images, which can be a single image or a batch.

        Returns:
            `List[PIL.Image.Image]`:
                A list of PIL images converted from the input NumPy array.
        """
        if images.ndim == 3:
            images = images[None, ...]
        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]
        else:
            pil_images = [Image.fromarray(image[:, :, :3]) for image in images]

        return pil_images

    @staticmethod
    def depth_pil_to_numpy(images: Union[List[PIL.Image.Image], PIL.Image.Image]) -> np.ndarray:
        r"""
        Convert a PIL image or a list of PIL images to NumPy arrays.

        Args:
            images (`Union[List[PIL.Image.Image], PIL.Image.Image]`):
                The input image or list of images to be converted.

        Returns:
            `np.ndarray`:
                A NumPy array of the converted images.
        """
        if not isinstance(images, list):
            images = [images]

        images = [np.array(image).astype(np.float32) / (2**16 - 1) for image in images]
        images = np.stack(images, axis=0)
        return images

    @staticmethod
    def rgblike_to_depthmap(image: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]:
        r"""
        Convert an RGB-like depth image to a depth map.

        Args:
            image (`Union[np.ndarray, torch.Tensor]`):
                The RGB-like depth image to convert.

        Returns:
            `Union[np.ndarray, torch.Tensor]`:
                The corresponding depth map.
        """
        return image[:, :, 1] * 2**8 + image[:, :, 2]

    def numpy_to_depth(self, images: np.ndarray) -> List[PIL.Image.Image]:
        r"""
        Convert a NumPy depth image or a batch of images to a list of PIL images.

        Args:
            images (`np.ndarray`):
                The input NumPy array of depth images, which can be a single image or a batch.

        Returns:
            `List[PIL.Image.Image]`:
                A list of PIL images converted from the input NumPy depth images.
        """
        if images.ndim == 3:
            images = images[None, ...]
        images_depth = images[:, :, :, 3:]
        if images.shape[-1] == 6:
            images_depth = (images_depth * 255).round().astype("uint8")
            pil_images = [
                Image.fromarray(self.rgblike_to_depthmap(image_depth), mode="I;16") for image_depth in images_depth
            ]
        elif images.shape[-1] == 4:
            images_depth = (images_depth * 65535.0).astype(np.uint16)
            pil_images = [Image.fromarray(image_depth, mode="I;16") for image_depth in images_depth]
        else:
            raise Exception("Not supported")

        return pil_images

    def postprocess(
        self,
        image: torch.Tensor,
        output_type: str = "pil",
        do_denormalize: Optional[List[bool]] = None,
    ) -> Union[PIL.Image.Image, np.ndarray, torch.Tensor]:
        """
        Postprocess the image output from tensor to `output_type`.

        Args:
            image (`torch.Tensor`):
                The image input, should be a pytorch tensor with shape `B x C x H x W`.
            output_type (`str`, *optional*, defaults to `pil`):
                The output type of the image, can be one of `pil`, `np`, `pt`, `latent`.
            do_denormalize (`List[bool]`, *optional*, defaults to `None`):
                Whether to denormalize the image to [0,1]. If `None`, will use the value of `do_normalize` in the
                `VaeImageProcessor` config.

        Returns:
            `PIL.Image.Image`, `np.ndarray` or `torch.Tensor`:
                The postprocessed image.
        """
        if not isinstance(image, torch.Tensor):
            raise ValueError(
                f"Input for postprocessing is in incorrect format: {type(image)}. We only support pytorch tensor"
            )
        if output_type not in ["latent", "pt", "np", "pil"]:
            deprecation_message = (
                f"the output_type {output_type} is outdated and has been set to `np`. Please make sure to set it to one of these instead: "
                "`pil`, `np`, `pt`, `latent`"
            )
            deprecate("Unsupported output_type", "1.0.0", deprecation_message, standard_warn=False)
            output_type = "np"

        image = self._denormalize_conditionally(image, do_denormalize)

        image = self.pt_to_numpy(image)

        if output_type == "np":
            if image.shape[-1] == 6:
                image_depth = np.stack([self.rgblike_to_depthmap(im[:, :, 3:]) for im in image], axis=0)
            else:
                image_depth = image[:, :, :, 3:]
            return image[:, :, :, :3], image_depth

        if output_type == "pil":
            return self.numpy_to_pil(image), self.numpy_to_depth(image)
        else:
            raise Exception(f"This type {output_type} is not supported")

    def preprocess(
        self,
        rgb: Union[torch.Tensor, PIL.Image.Image, np.ndarray],
        depth: Union[torch.Tensor, PIL.Image.Image, np.ndarray],
        height: Optional[int] = None,
        width: Optional[int] = None,
        target_res: Optional[int] = None,
    ) -> torch.Tensor:
        r"""
        Preprocess the image input. Accepted formats are PIL images, NumPy arrays, or PyTorch tensors.

        Args:
            rgb (`Union[torch.Tensor, PIL.Image.Image, np.ndarray]`):
                The RGB input image, which can be a single image or a batch.
            depth (`Union[torch.Tensor, PIL.Image.Image, np.ndarray]`):
                The depth input image, which can be a single image or a batch.
            height (`Optional[int]`, *optional*, defaults to `None`):
                The desired height of the processed image. If `None`, defaults to the height of the input image.
            width (`Optional[int]`, *optional*, defaults to `None`):
                The desired width of the processed image. If `None`, defaults to the width of the input image.
            target_res (`Optional[int]`, *optional*, defaults to `None`):
                Target resolution for resizing the images. If specified, overrides height and width.

        Returns:
            `Tuple[torch.Tensor, torch.Tensor]`:
                A tuple containing the processed RGB and depth images as PyTorch tensors.
        """
        supported_formats = (PIL.Image.Image, np.ndarray, torch.Tensor)

        # Expand the missing dimension for 3-dimensional pytorch tensor or numpy array that represents grayscale image
        if self.config.do_convert_grayscale and isinstance(rgb, (torch.Tensor, np.ndarray)) and rgb.ndim == 3:
            raise Exception("This is not yet supported")

        if isinstance(rgb, supported_formats):
            rgb = [rgb]
            depth = [depth]
        elif not (isinstance(rgb, list) and all(isinstance(i, supported_formats) for i in rgb)):
            raise ValueError(
                f"Input is in incorrect format: {[type(i) for i in rgb]}. Currently, we only support {', '.join(supported_formats)}"
            )

        if isinstance(rgb[0], PIL.Image.Image):
            if self.config.do_convert_rgb:
                raise Exception("This is not yet supported")
                # rgb = [self.convert_to_rgb(i) for i in rgb]
                # depth = [self.convert_to_depth(i) for i in depth]  #TODO define convert_to_depth
            if self.config.do_resize or target_res:
                height, width = self.get_default_height_width(rgb[0], height, width) if not target_res else target_res
                rgb = [self.resize(i, height, width) for i in rgb]
                depth = [self.resize(i, height, width) for i in depth]
            rgb = self.pil_to_numpy(rgb)  # to np
            rgb = self.numpy_to_pt(rgb)  # to pt

            depth = self.depth_pil_to_numpy(depth)  # to np
            depth = self.numpy_to_pt(depth)  # to pt

        elif isinstance(rgb[0], np.ndarray):
            rgb = np.concatenate(rgb, axis=0) if rgb[0].ndim == 4 else np.stack(rgb, axis=0)
            rgb = self.numpy_to_pt(rgb)
            height, width = self.get_default_height_width(rgb, height, width)
            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 = self.numpy_to_pt(depth)
            height, width = self.get_default_height_width(depth, height, width)
            if self.config.do_resize:
                depth = self.resize(depth, height, width)

        elif isinstance(rgb[0], torch.Tensor):
            raise Exception("This is not yet supported")
            # rgb = torch.cat(rgb, axis=0) if rgb[0].ndim == 4 else torch.stack(rgb, axis=0)

            # if self.config.do_convert_grayscale and rgb.ndim == 3:
            #     rgb = rgb.unsqueeze(1)

            # channel = rgb.shape[1]

            # height, width = self.get_default_height_width(rgb, height, width)
            # if self.config.do_resize:
            #     rgb = self.resize(rgb, height, width)

            # depth = torch.cat(depth, axis=0) if depth[0].ndim == 4 else torch.stack(depth, axis=0)

            # if self.config.do_convert_grayscale and depth.ndim == 3:
            #     depth = depth.unsqueeze(1)

            # channel = depth.shape[1]
            # # don't need any preprocess if the image is latents
            # if depth == 4:
            #     return rgb, depth

            # height, width = self.get_default_height_width(depth, height, width)
            # if self.config.do_resize:
            #     depth = self.resize(depth, height, width)
        # expected range [0,1], normalize to [-1,1]
        do_normalize = self.config.do_normalize
        if rgb.min() < 0 and do_normalize:
            warnings.warn(
                "Passing `image` as torch tensor with value range in [-1,1] is deprecated. The expected value range for image tensor is [0,1] "
                f"when passing as pytorch tensor or numpy Array. You passed `image` with value range [{rgb.min()},{rgb.max()}]",
                FutureWarning,
            )
            do_normalize = False

        if do_normalize:
            rgb = self.normalize(rgb)
            depth = self.normalize(depth)

        if self.config.do_binarize:
            rgb = self.binarize(rgb)
            depth = self.binarize(depth)

        return rgb, depth
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/image_processor.py