VaeImageProcessorLDM3D.preprocess fails for batched NumPy RGB/depth inputs

Start in src/diffusers/image_processor.py at VaeImageProcessorLDM3D.preprocess around line 1217, then inspect numpy_to_pt around line 188. Run the provided batched RGB/depth reproduction and add regression coverage for the expected tensor shapes. Done means the valid 4D inputs no longer raise the axis error while the single-image case remains working.

Describe the bug

VaeImageProcessorLDM3D.preprocess documents that NumPy RGB and depth inputs can be either a single image or a batch. A single 3D RGB/depth pair works, but a valid 4D batch raises ValueError: axes don't match array.

The RGB batch is concatenated and converted to a PyTorch tensor before the depth batching condition is evaluated. The condition then checks rgb[0].ndim instead of the original depth[0].ndim, so the 4D depth batch is stacked into a 5D array and passed to numpy_to_pt, which expects four axes.

Relevant source: https://github.com/huggingface/diffusers/blob/d6726f38a0c5ca6c06a8f227fb7bade3486ed98d/src/diffusers/image_processor.py#L1217

I would like to contribute the focused one-line correction and a regression test once a maintainer confirms the scope, following the repository's AI-assisted contribution policy.

Reproduction

import numpy as np

from diffusers import VaeImageProcessorLDM3D

processor = VaeImageProcessorLDM3D(vae_scale_factor=1)
rgb = np.zeros((2, 8, 8, 3), dtype=np.float32)
depth = np.zeros((2, 8, 8, 1), dtype=np.float32)

processed_rgb, processed_depth = processor.preprocess(rgb, depth)
print(processed_rgb.shape, processed_depth.shape)

Expected shapes:

torch.Size([2, 3, 8, 8]) torch.Size([2, 1, 8, 8])

Actual result:

ValueError: axes don't match array

The same reproduction fails on both Diffusers 0.39.0 and current main at d6726f38. Passing corresponding single-image arrays with shapes (8, 8, 3) and (8, 8, 1) succeeds.

Logs

Traceback (most recent call last):
  ...
  File "diffusers/image_processor.py", line 1218, in preprocess
    depth = self.numpy_to_pt(depth)
  File "diffusers/image_processor.py", line 188, in numpy_to_pt
    images = torch.from_numpy(images.transpose(0, 3, 1, 2))
ValueError: axes don't match array

System Info

• Diffusers version: 0.39.0; also reproduced on 0.40.0.dev0 at d6726f38
• Platform: Windows-10-10.0.26200-SP0
• Python version: 3.10.11
• PyTorch version: 2.13.0+cpu
• huggingface_hub version: 1.27.0
• Safetensors version: 0.8.0
• GPU used in script: No
• Distributed or parallel setup: No

Who can help?

@sayakpaul @DN6

AI disclosure: I used Codex to help identify and reproduce the behavior, verify it on the latest release and current main, search existing issues and PRs, and draft this report. I reviewed the reproduction and diagnosis and will personally handle any follow-up.