diff --git a/src/diffusers/image_processor.py b/src/diffusers/image_processor.py index 4f6f4bd52b9c..e69f6c3a40ff 100644 --- a/src/diffusers/image_processor.py +++ b/src/diffusers/image_processor.py @@ -1048,28 +1048,16 @@ 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 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 int32 — the 16-bit value does not fit back in the uint8 input dtype. + return depth_map 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) + # uint16 matches the mode="I;16" consumer in numpy_to_depth. + return depth_map.astype(np.uint16) else: raise TypeError("Input image must be a torch.Tensor or np.ndarray")