diff --git a/src/diffusers/image_processor.py b/src/diffusers/image_processor.py index 4f6f4bd52b9c..393f394ad0c5 100644 --- a/src/diffusers/image_processor.py +++ b/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") diff --git a/tests/others/test_image_processor.py b/tests/others/test_image_processor.py index 88e82ab54b82..367841b91ec5 100644 --- a/tests/others/test_image_processor.py +++ b/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()}"