Repository: huggingface/diffusers
File: src/diffusers/image_processor.py (Python)

Code before:
        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.
        """
        # 1. Cast the tensor to a larger integer type (e.g., int32)
        #    to safely perform the multiplication by 256.
        # 2. Perform the 16-bit combination: High-byte * 256 + Low-byte.
        # 3. Cast the final result to the desired depth map type (uint16) if needed 
        #    before returning, though leaving it as int32/int64 is often safer 
        #    for return value from a library function.
        
        if isinstance(image, torch.Tensor):
            # Cast to a safe dtype (e.g., int32 or int64) for the calculation
            image_safe = image.to(torch.int32) 
            
            # Calculate the depth map
            depth_map = image_safe[:, :, 1] * 256 + image_safe[:, :, 2]
            
            # You may want to cast the final result to uint16, but casting to a 
            # larger int type (like int32) is sufficient to fix the overflow.
            # depth_map = depth_map.to(torch.uint16) # Uncomment if uint16 is strictly required
            return depth_map
            
        elif isinstance(image, np.ndarray):
            # NumPy equivalent: Cast to a safe dtype (e.g., np.int32)
            image_safe = image.astype(np.int32)
            
            # Calculate the depth map
            depth_map = image_safe[:, :, 1] * 256 + image_safe[:, :, 2]
            
            # depth_map = depth_map.astype(np.uint16) # Uncomment if uint16 is strictly required
            return depth_map
        else:
            raise TypeError("Input image must be a torch.Tensor or np.ndarray")

    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.
        """

Review comment (suggestion):
Change looks good. Could we just store the original dtype and cast back once the computation is done.
```python
original_dtype = image.dtype
# perform compute
return depth_map.to(original_dtype)
```

Code after:
        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.
        """
        # 1. Cast the tensor to a larger integer type (e.g., int32)
        #    to safely perform the multiplication by 256.
        # 2. Perform the 16-bit combination: High-byte * 256 + Low-byte.
        # 3. Cast the final result to the desired depth map type (uint16) if needed
        #    before returning, though leaving it as int32/int64 is often safer
        #    for return value from a library function.

        if isinstance(image, torch.Tensor):
            # Cast to a safe dtype (e.g., int32 or int64) for the calculation
            original_dtype = image.dtype
            image_safe = image.to(torch.int32)

            # Calculate the depth map
            depth_map = image_safe[:, :, 1] * 256 + image_safe[:, :, 2]

            # You may want to cast the final result to uint16, but casting to a
            # larger int type (like int32) is sufficient to fix the overflow.
            # depth_map = depth_map.to(torch.uint16) # Uncomment if uint16 is strictly required
            return depth_map.to(original_dtype)

        elif isinstance(image, np.ndarray):
            # NumPy equivalent: Cast to a safe dtype (e.g., np.int32)
            original_dtype = image.dtype
            image_safe = image.astype(np.int32)

            # Calculate the depth map
            depth_map = image_safe[:, :, 1] * 256 + image_safe[:, :, 2]

            # depth_map = depth_map.astype(np.uint16) # Uncomment if uint16 is strictly required
            return depth_map.astype(original_dtype)
        else:
            raise TypeError("Input image must be a torch.Tensor or np.ndarray")

    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]`: