fix(image_processor): return uint16 from LDM3D rgblike_to_depthmap

src/diffusers/image_processor.py

@@ -1039,37 +1039,22 @@ def depth_pil_to_numpy(images: list[PIL.Image.Image] | PIL.Image.Image) -> np.nd
     def rgblike_to_depthmap(image: np.ndarray | torch.Tensor) -> 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.
 
+        The depth is reconstructed by combining two 8-bit channels into a 16-bit
+        value, so the result is always returned as `uint16` regardless of the
+        input dtype. Casting back to the 8-bit input dtype would truncate the
+        value (see #14206) and make the subsequent `numpy_to_depth` build of the
+        `mode="I;16"` PIL image fail with "buffer is not large enough".
+        """
         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)
+            return depth_map.to(torch.uint16)
 
         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)
+            return depth_map.astype(np.uint16)
         else:
             raise TypeError("Input image must be a torch.Tensor or np.ndarray")
 

tests/others/test_image_processor.py

@@ -19,7 +19,7 @@
 import PIL.Image
 import torch
 
-from diffusers.image_processor import VaeImageProcessor
+from diffusers.image_processor import VaeImageProcessor, VaeImageProcessorLDM3D
 
 
 class ImageProcessorTest(unittest.TestCase):
@@ -308,3 +308,35 @@ 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_ldm3d_rgblike_to_depthmap_returns_uint16(self):
+        # VaeImageProcessorLDM3D.rgblike_to_depthmap combines two 8-bit channels
+        # into a 16-bit depth value, so it must return uint16 regardless of the
+        # input dtype. Casting back to the 8-bit input dtype truncates the value
+        # and breaks the downstream `numpy_to_depth` build of the I;16 PIL image.
+        high, low = 200, 100
+        expected = high * 256 + low
+
+        np_image = np.zeros((8, 8, 6), dtype=np.uint8)
+        np_image[:, :, 1] = high
+        np_image[:, :, 2] = low
+        np_depth = VaeImageProcessorLDM3D.rgblike_to_depthmap(np_image)
+        assert np_depth.dtype == np.uint16, f"expected uint16, got {np_depth.dtype}"
+        assert int(np_depth[0, 0]) == expected
+
+        pt_image = torch.zeros((8, 8, 6), dtype=torch.uint8)
+        pt_image[:, :, 1] = high
+        pt_image[:, :, 2] = low
+        pt_depth = VaeImageProcessorLDM3D.rgblike_to_depthmap(pt_image)
+        assert pt_depth.dtype == torch.uint16, f"expected uint16, got {pt_depth.dtype}"
+        assert int(pt_depth[0, 0]) == expected
+
+        # The depth map feeds numpy_to_depth, which builds a mode="I;16" PIL image;
+        # that must not raise "buffer is not large enough".
+        image_processor = VaeImageProcessorLDM3D()
+        batch = np.zeros((1, 8, 8, 6), dtype=np.uint8)
+        batch[:, :, :, 1] = high
+        batch[:, :, :, 2] = low
+        pil_images = image_processor.numpy_to_depth(batch)
+        assert pil_images[0].mode == "I;16"
+        assert pil_images[0].size == (8, 8)