[LDM3D] Fix rgblike_to_depthmap truncating 16-bit depth back to 8 bits

src/diffusers/image_processor.py

@@ -1048,28 +1048,28 @@ def rgblike_to_depthmap(image: np.ndarray | torch.Tensor) -> np.ndarray | torch.
         #    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
+            # Cast to a wider integer type so the 16-bit combination does not
+            # overflow the 8-bit input dtype.
             image_safe = image.to(torch.int32)
 
-            # Calculate the depth map
+            # Calculate the 16-bit depth map: high-byte * 256 + low-byte.
             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)
+            # The result is 16-bit and must not be truncated back to the 8-bit
+            # input dtype; keep the wider integer type.
+            return depth_map
 
         elif isinstance(image, np.ndarray):
-            # NumPy equivalent: Cast to a safe dtype (e.g., np.int32)
-            original_dtype = image.dtype
+            # NumPy equivalent: cast to a wider dtype to avoid overflow.
             image_safe = image.astype(np.int32)
 
-            # Calculate the depth map
+            # Calculate the 16-bit depth map: high-byte * 256 + low-byte.
             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 a 16-bit-wide array as the ``mode="I;16"`` consumer in
+            # ``numpy_to_depth`` requires; truncating to the 8-bit input dtype
+            # would drop the high byte and break that consumer.
+            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,37 @@ 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}'."
         )
+
+
+class Ldm3dImageProcessorTest(unittest.TestCase):
+    def test_rgblike_to_depthmap_combines_two_bytes_into_16bit(self):
+        # green * 256 + blue must be a 16-bit value; casting the result back to
+        # the 8-bit input dtype drops the high byte (e.g. 1480 -> 200).
+        processor = VaeImageProcessorLDM3D()
+
+        rgb_np = np.zeros((2, 2, 3), dtype=np.uint8)
+        rgb_np[..., 1] = 5  # green (high byte)
+        rgb_np[..., 2] = 200  # blue (low byte)
+        depth_np = processor.rgblike_to_depthmap(rgb_np)
+        assert depth_np.dtype == np.uint16
+        assert int(depth_np.max()) == 5 * 256 + 200
+
+        rgb_pt = torch.zeros(2, 2, 3, dtype=torch.uint8)
+        rgb_pt[..., 1] = 5
+        rgb_pt[..., 2] = 200
+        depth_pt = processor.rgblike_to_depthmap(rgb_pt)
+        assert int(depth_pt.max()) == 5 * 256 + 200
+
+    def test_postprocess_pil_depth_does_not_crash(self):
+        # numpy_to_depth feeds rgblike_to_depthmap into Image.fromarray(mode="I;16"),
+        # which requires 16-bit-wide data; an 8-bit result raised
+        # "ValueError: buffer is not large enough".
+        processor = VaeImageProcessorLDM3D()
+        image = torch.zeros(1, 6, 4, 4)
+        image[:, 4] = 5 / 255.0
+        image[:, 5] = 200 / 255.0
+
+        rgb, depth = processor.postprocess(image, output_type="pil")
+
+        assert len(depth) == 1
+        assert depth[0].mode == "I;16"