shap_e model/pipeline review

Commit tested: 0f1abc4ae8b0eb2a3b40e82a310507281144c423

Review performed against the repository review rules.

Duplicate search status: checked GitHub Issues/PRs for shap_e, ShapE, ShapEImg2ImgPipeline, StratifiedRaySampler, mesh output, latent dtype, return_dict, and frame-size/ray batching. I found no likely open duplicates. Closed issue #4075 is only a docstring typo, closed issue #4808 is an old integration snapshot failure, and merged PR #4062 is relevant mesh precedent but does not cover the current gaps.

Issue 1: ShapEImg2ImgPipeline drops or rejects documented image batch inputs

Affected code:

Problem:
The docstring says image accepts np.ndarray and list[np.ndarray], but __call__ rejects those types. For list[PIL.Image.Image], _encode_image() calls the image processor and then indexes [0].unsqueeze(0), silently keeping only the first processed image.

Impact:
Documented batched image inputs fail with shape errors or are reduced to one image. This also leaves fast tests blind to PIL/NumPy image batches because they use tensor inputs.

Reproduction:

from PIL import Image
import numpy as np
import torch
from diffusers import ShapEImg2ImgPipeline
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModel

image_encoder = CLIPVisionModel(CLIPVisionConfig(hidden_size=8, image_size=32, intermediate_size=16, num_attention_heads=2, num_hidden_layers=1, patch_size=1))
image_processor = CLIPImageProcessor(do_resize=True, size={"shortest_edge": 32}, do_center_crop=True, crop_size={"height": 32, "width": 32}, do_normalize=False)
pipe = ShapEImg2ImgPipeline(None, image_encoder, image_processor, None, None)

images = [Image.fromarray(np.zeros((32, 32, 3), dtype=np.uint8)) for _ in range(2)]
embeds = pipe._encode_image(images, torch.device("cpu"), 1, False)
print(embeds.shape[0])  # 1, expected 2

Relevant precedent:
Stable Diffusion img2img preprocesses the whole input batch instead of indexing the first item:

Suggested fix:

elif isinstance(image, np.ndarray):
    batch_size = 1 if image.ndim == 3 else image.shape[0]
elif isinstance(image, list) and isinstance(image[0], (torch.Tensor, PIL.Image.Image, np.ndarray)):
    batch_size = len(image)

...

if not isinstance(image, torch.Tensor):
    image = self.image_processor(image, return_tensors="pt").pixel_values

Issue 2: Provided latents are not consistently validated or cast to model dtype

Affected code:

Problem:
prepare_latents() moves user-provided latents to the device but not the requested dtype. In img2img, provided latents skip prepare_latents() entirely, so they also skip shape validation and scheduler scaling.

Impact:
Half-precision pipelines fail with dtype mismatch when users pass normal float32 latents. Img2img can fail later inside PriorTransformer with cryptic batch/shape errors instead of raising at input validation.

Reproduction:

import types
import torch
from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEPipeline

prior = PriorTransformer(
    num_attention_heads=2, attention_head_dim=8, embedding_dim=8, num_embeddings=4,
    embedding_proj_dim=16, time_embed_dim=32, num_layers=1, clip_embed_dim=16,
    additional_embeddings=0, norm_in_type="layer", encoder_hid_proj_type=None, added_emb_type=None,
).to(dtype=torch.float16)
scheduler = HeunDiscreteScheduler(beta_schedule="exp", num_train_timesteps=8, prediction_type="sample")
pipe = ShapEPipeline(prior, None, None, scheduler, None)
pipe.set_progress_bar_config(disable=True)
pipe._encode_prompt = types.MethodType(lambda self, *args, **kwargs: torch.zeros(1, 16, dtype=torch.float16), pipe)

latents = torch.zeros(1, 4 * 8, dtype=torch.float32)
pipe("x", latents=latents, num_inference_steps=1, guidance_scale=1.0, output_type="latent")

Relevant precedent:
Flux2 casts provided latents to both device and dtype:

Suggested fix:

latents = latents.to(device=device, dtype=dtype)

For img2img, call prepare_latents(...) unconditionally, as the text pipeline does, so provided latents get the same validation/scaling path.

Issue 3: Fresh ShapERenderer mesh output has an empty marching-cubes lookup table

Affected code:

Problem:
MeshDecoder initializes cases and masks to zeros. The conversion script populates those buffers, but the runtime constructor does not. A renderer created from config or in fast tests cannot produce mesh geometry.

Impact:
output_type="mesh" is documented, but fresh/local renderer construction produces empty meshes and can later fail with empty texture batches. Tests do not catch this path.

Reproduction:

import torch
from diffusers.pipelines.shap_e.renderer import MeshDecoder

decoder = MeshDecoder()
print(int(decoder.cases.abs().sum()), int(decoder.masks.sum()))  # 0 0

field = torch.randn(4, 4, 4)
mesh = decoder(field, torch.tensor([-1., -1., -1.]), torch.tensor([2., 2., 2.]))
print(mesh.verts.shape, mesh.faces.shape)  # empty geometry

Relevant precedent:
The conversion script already has the lookup-table builder:

Suggested fix:
Move the marching-cubes table generation into runtime code, or store a static table used by MeshDecoder.__init__(). Keep loading checkpoint buffers for backwards compatibility, but the default constructor should initialize functional cases and masks.

Issue 4: decode_to_image() drops remainder rays and fails for many frame_size values

Affected code:

Problem:
n_batches = rays.shape[1] // ray_batch_size uses floor division. If total rays are smaller than ray_batch_size, images remains empty. If total rays are not exactly divisible, the tail rays are dropped and the final .view(...) shape is invalid.

Impact:
Users can pass arbitrary frame_size, but many sizes fail at render time. Current tests only use latent output in fast tests and frame_size=64 in nightly tests, so this is uncovered.

Reproduction:

import torch
from diffusers.pipelines.shap_e import ShapERenderer

renderer = ShapERenderer(
    param_shapes=((8, 93), (8, 8), (8, 8), (8, 8)),
    d_latent=16,
    d_hidden=8,
    n_output=12,
)
latents = torch.zeros(1, 32, 16)
renderer.decode_to_image(latents, torch.device("cpu"), size=9)

Relevant precedent:
No duplicate found.

Suggested fix:

for start in range(0, rays.shape[1], ray_batch_size):
    rays_batch = rays[:, start : start + ray_batch_size]
    _, fine_sampler, coarse_model_out = self.render_rays(rays_batch, coarse_sampler, n_coarse_samples)
    channels, _, _ = self.render_rays(rays_batch, fine_sampler, n_fine_samples, prev_model_out=coarse_model_out)
    images.append(channels)

Issue 5: Renderer sampling resets the global PyTorch RNG

Affected code:

Problem:
StratifiedRaySampler.sample() calls torch.manual_seed(0) inside production rendering. This mutates global RNG state, and the adjacent comment is explicitly temporary/debug context.

Impact:
Calling Shap-E rendering changes the caller’s global RNG sequence. It can also make each ray batch reuse the same sampling pattern.

Reproduction:

import torch
from diffusers.pipelines.shap_e.renderer import StratifiedRaySampler

torch.manual_seed(123)
sampler = StratifiedRaySampler()
sampler.sample(torch.zeros(1, 1), torch.ones(1, 1), 4)
print(torch.initial_seed())  # 0, expected to remain 123

Relevant precedent:
No duplicate found.

Suggested fix:
Thread a local torch.Generator through decode_to_image(), render_rays(), StratifiedRaySampler.sample(), ImportanceRaySampler.sample(), and sample_pmf(). If snapshot stability is required, seed that local generator with 0 without touching global state.

Issue 6: return_dict=False is ignored for latent outputs

Affected code:

Problem:
Both pipelines return ShapEPipelineOutput immediately for output_type="latent", before the shared if not return_dict branch.

Impact:
The public return_dict=False contract changes based on output type.

Reproduction:

import types
import torch
from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEPipeline

prior = PriorTransformer(num_attention_heads=2, attention_head_dim=8, embedding_dim=8, num_embeddings=4, embedding_proj_dim=16, time_embed_dim=32, num_layers=1, clip_embed_dim=16, additional_embeddings=0, norm_in_type="layer", encoder_hid_proj_type=None, added_emb_type=None)
scheduler = HeunDiscreteScheduler(beta_schedule="exp", num_train_timesteps=8, prediction_type="sample")
pipe = ShapEPipeline(prior, None, None, scheduler, None)
pipe.set_progress_bar_config(disable=True)
pipe._encode_prompt = types.MethodType(lambda self, *args, **kwargs: torch.zeros(1, 16), pipe)

out = pipe("x", num_inference_steps=1, guidance_scale=1.0, output_type="latent", return_dict=False)
print(type(out).__name__, isinstance(out, tuple))  # ShapEPipelineOutput False

Relevant precedent:
Flux keeps latent handling before the final return_dict branch:

Suggested fix:

if output_type == "latent":
    images = latents
else:
    ...

if not return_dict:
    return (images,)

return ShapEPipelineOutput(images=images)

Issue 7: DifferentiableProjectiveCamera.resize_image() cannot construct the resized camera

Affected code:

Problem:
DifferentiableProjectiveCamera requires shape, but resize_image() omits it when constructing the replacement camera.

Impact:
The camera helper returned by create_pan_cameras() exposes a broken resize method.

Reproduction:

from diffusers.pipelines.shap_e import create_pan_cameras

cam = create_pan_cameras(32)
cam.resize_image(64, 64)

Relevant precedent:
No duplicate found.

Suggested fix:

return DifferentiableProjectiveCamera(
    origin=self.origin,
    x=self.x,
    y=self.y,
    z=self.z,
    width=width,
    height=height,
    x_fov=self.x_fov,
    y_fov=self.y_fov,
    shape=self.shape,
)

Issue 8: Slow tests are missing, and offload coverage is skipped

Affected code:

Problem:
The target has fast tests and nightly integration tests, but no @slow tests. Both sequential CPU offload tests are explicitly skipped with “Key error is raised with accelerate.”

Impact:
RUN_SLOW does not exercise Shap-E, mesh output is untested, non-latent rendering is only covered nightly, and offload behavior is known-uncovered despite being in scope for pipeline review.

Reproduction:

from pathlib import Path

for path in sorted(Path("tests/pipelines/shap_e").glob("test_*.py")):
    text = path.read_text()
    print(path, "@slow" in text, "@nightly" in text)

Relevant precedent:
No duplicate found.

Suggested fix:
Add slow imports and @slow coverage for at least one pretrained text-to-3D and img2img run, plus a small output_type="mesh" assertion. Unskip or replace the sequential CPU offload tests with a current accelerate-compatible offload regression test.