fix: preserve 16-bit depth in LDM3D rgblike_to_depthmap

src/diffusers/image_processor.py

@@ -1049,27 +1049,21 @@ def rgblike_to_depthmap(image: np.ndarray | torch.Tensor) -> np.ndarray | torch.
 
         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

@@ -308,3 +308,27 @@ 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_rgblike_to_depthmap_preserves_uint16_range(self):
+        """Test that rgblike_to_depthmap returns uint16 values, not truncated to uint8."""
+        from diffusers.image_processor import VaeImageProcessorLDM3D
+
+        processor = VaeImageProcessorLDM3D()
+
+        # Create a test image where high/low bytes encode a depth > 255
+        # e.g., channel 1 (high byte) = 1, channel 2 (low byte) = 0 → depth = 256
+        h, w = 4, 4
+        img_np = np.zeros((h, w, 3), dtype=np.uint8)
+        img_np[:, :, 1] = 1  # high byte
+        img_np[:, :, 2] = 0  # low byte
+        depth_np = processor.rgblike_to_depthmap(img_np)
+        assert depth_np.dtype == np.uint16, f"Expected uint16, got {depth_np.dtype}"
+        assert depth_np[0, 0] == 256, f"Expected 256, got {depth_np[0, 0]}"
+
+        # Torch variant (H, W, C) layout
+        img_pt = torch.zeros(h, w, 3, dtype=torch.uint8)
+        img_pt[:, :, 1] = 1
+        img_pt[:, :, 2] = 0
+        depth_pt = processor.rgblike_to_depthmap(img_pt)
+        assert depth_pt.dtype == torch.uint16, f"Expected uint16, got {depth_pt.dtype}"
+        assert depth_pt[0, 0].item() == 256, f"Expected 256, got {depth_pt[0, 0].item()}"