From 6172013d67792613b1ee1f300373ebec82f77211 Mon Sep 17 00:00:00 2001 From: Andrew Chen <48723787+chuenchen309@users.noreply.github.com> Date: Thu, 16 Jul 2026 10:42:32 +0800 Subject: [PATCH] [LDM3D] Fix rgblike_to_depthmap truncating 16-bit depth back to 8 bits VaeImageProcessorLDM3D.rgblike_to_depthmap combines two 8-bit channels into a 16-bit depth value (green * 256 + blue, range 0-65535). It correctly casts up to a wider integer type for the multiplication, but then cast the result back to the input's dtype (uint8) via `.to(original_dtype)` / `.astype(original_dtype)`, dropping the entire high byte -- so it returned only `depth % 256` (e.g. 1480 came back as 200). numpy_to_depth feeds this into `Image.fromarray(..., mode="I;16")`, which requires 16-bit-wide data, so with a uint8 buffer `VaeImageProcessorLDM3D.postprocess(image, output_type="pil")` -- the normal way an LDM3D pipeline returns a depth image for a 6-channel output -- crashed with `ValueError: buffer is not large enough`. Keep the 16-bit result instead of truncating: return the wider integer type directly (numpy -> uint16 to match the `mode="I;16"` consumer, torch -> the int32 combination). The function's whole purpose is to produce a 16-bit depth map from 8-bit RGB, so casting back to the 8-bit input dtype was never correct. This regression was introduced by #12546 (which fixed a real NumPy 2.0 `uint8 * 256` overflow but replaced it with this truncation). The existing tests/others/test_image_processor.py had no VaeImageProcessorLDM3D coverage, which is why it shipped; this adds a small regression test class covering the 16-bit combination (numpy + torch) and the pil postprocess path. Verified against the real code before/after: rgblike_to_depthmap(green=5, blue=200) returned 200 (uint8) before and 1480 (uint16 / int32) after; postprocess(output_type="pil") raised ValueError before and returns a valid "I;16" depth image after. Full tests/others/test_image_processor.py: 12 passed. ruff 0.9.10 (repo-pinned) check + format: clean. Note (separate, not addressed here to keep this focused): the output_type="np" branch of postprocess passes float [0, 1] arrays into rgblike_to_depthmap, where the int cast collapses them to 0; that is an independent caller-side issue in the same function's contract and is better handled in its own change. --- src/diffusers/image_processor.py | 24 +++++++++---------- tests/others/test_image_processor.py | 36 +++++++++++++++++++++++++++- 2 files changed, 47 insertions(+), 13 deletions(-) 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"