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