diff --git a/src/diffusers/image_processor.py b/src/diffusers/image_processor.py index 4f6f4bd52b9c..941b1df99150 100644 --- a/src/diffusers/image_processor.py +++ b/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") diff --git a/tests/others/test_image_processor.py b/tests/others/test_image_processor.py index 88e82ab54b82..a2a783651d89 100644 --- a/tests/others/test_image_processor.py +++ b/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"