shap_e model/pipeline review

[hlky]

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:

diffusers/src/diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py

Lines 149 to 153 in 0f1abc4

	if isinstance(image, list) and isinstance(image[0], torch.Tensor):
	image = torch.cat(image, axis=0) if image[0].ndim == 4 else torch.stack(image, axis=0)
	if not isinstance(image, torch.Tensor):
	image = self.image_processor(image, return_tensors="pt").pixel_values[0].unsqueeze(0)

diffusers/src/diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py

Lines 225 to 234 in 0f1abc4

	if isinstance(image, PIL.Image.Image):
	batch_size = 1
	elif isinstance(image, torch.Tensor):
	batch_size = image.shape[0]
	elif isinstance(image, list) and isinstance(image[0], (torch.Tensor, PIL.Image.Image)):
	batch_size = len(image)
	else:
	raise ValueError(
	f"`image` has to be of type `PIL.Image.Image`, `torch.Tensor`, `list[PIL.Image.Image]` or `list[torch.Tensor]` but is {type(image)}"
	)

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:

diffusers/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py

Line 1048 in 0f1abc4

image = self.image_processor.preprocess(image)

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:

diffusers/src/diffusers/pipelines/shap_e/pipeline_shap_e.py

Lines 130 to 138 in 0f1abc4

	def prepare_latents(self, shape, dtype, device, generator, latents, scheduler):
	if latents is None:
	latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
	else:
	if latents.shape != shape:
	raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
	latents = latents.to(device)
	latents = latents * scheduler.init_noise_sigma

diffusers/src/diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py

Lines 130 to 140 in 0f1abc4

	# Copied from diffusers.pipelines.deprecated.unclip.pipeline_unclip.UnCLIPPipeline.prepare_latents
	def prepare_latents(self, shape, dtype, device, generator, latents, scheduler):
	if latents is None:
	latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
	else:
	if latents.shape != shape:
	raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
	latents = latents.to(device)
	latents = latents * scheduler.init_noise_sigma
	return latents

diffusers/src/diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py

Lines 250 to 261 in 0f1abc4

	if latents is None:
	latents = self.prepare_latents(
	(batch_size, num_embeddings * embedding_dim),
	image_embeds.dtype,
	device,
	generator,
	latents,
	self.scheduler,
	)
	# YiYi notes: for testing only to match ldm, we can directly create a latents with desired shape: batch_size, num_embeddings, embedding_dim
	latents = latents.reshape(latents.shape[0], num_embeddings, embedding_dim)

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:

diffusers/src/diffusers/pipelines/flux2/pipeline_flux2.py

Line 644 in 0f1abc4

latents = latents.to(device=device, dtype=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:

diffusers/src/diffusers/pipelines/shap_e/renderer.py

Lines 489 to 495 in 0f1abc4

	def __init__(self):
	super().__init__()
	cases = torch.zeros(256, 5, 3, dtype=torch.long)
	masks = torch.zeros(256, 5, dtype=torch.bool)
	self.register_buffer("cases", cases)
	self.register_buffer("masks", masks)

diffusers/docs/source/en/using-diffusers/shap-e.md

Lines 135 to 150 in 0f1abc4

	You can generate mesh outputs for both the [`ShapEPipeline`] and [`ShapEImg2ImgPipeline`] by specifying the `output_type` parameter as `"mesh"`:
	```py
	import torch
	from diffusers import ShapEPipeline
	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
	pipe = ShapEPipeline.from_pretrained("openai/shap-e", torch_dtype=torch.float16, variant="fp16")
	pipe = pipe.to(device)
	guidance_scale = 15.0
	prompt = "A birthday cupcake"
	images = pipe(prompt, guidance_scale=guidance_scale, num_inference_steps=64, frame_size=256, output_type="mesh").images
	```

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:

diffusers/scripts/convert_shap_e_to_diffusers.py

Lines 830 to 889 in 0f1abc4

	def create_mc_lookup_table():
	cases = torch.zeros(256, 5, 3, dtype=torch.long)
	masks = torch.zeros(256, 5, dtype=torch.bool)
	edge_to_index = {
	(0, 1): 0,
	(2, 3): 1,
	(4, 5): 2,
	(6, 7): 3,
	(0, 2): 4,
	(1, 3): 5,
	(4, 6): 6,
	(5, 7): 7,
	(0, 4): 8,
	(1, 5): 9,
	(2, 6): 10,
	(3, 7): 11,
	}
	for i, case in enumerate(MC_TABLE):
	for j, tri in enumerate(case):
	for k, (c1, c2) in enumerate(zip(tri[::2], tri[1::2])):
	cases[i, j, k] = edge_to_index[(c1, c2) if c1 < c2 else (c2, c1)]
	masks[i, j] = True
	return cases, masks
	RENDERER_CONFIG = {}
	def renderer_model_from_original_config():
	model = ShapERenderer(**RENDERER_CONFIG)
	return model
	RENDERER_MLP_ORIGINAL_PREFIX = "renderer.nerstf"
	RENDERER_PARAMS_PROJ_ORIGINAL_PREFIX = "encoder.params_proj"
	def renderer_model_original_checkpoint_to_diffusers_checkpoint(model, checkpoint):
	diffusers_checkpoint = {}
	diffusers_checkpoint.update(
	{f"mlp.{k}": checkpoint[f"{RENDERER_MLP_ORIGINAL_PREFIX}.{k}"] for k in model.mlp.state_dict().keys()}
	)
	diffusers_checkpoint.update(
	{
	f"params_proj.{k}": checkpoint[f"{RENDERER_PARAMS_PROJ_ORIGINAL_PREFIX}.{k}"]
	for k in model.params_proj.state_dict().keys()
	}
	)
	diffusers_checkpoint.update({"void.background": model.state_dict()["void.background"]})
	cases, masks = create_mc_lookup_table()
	diffusers_checkpoint.update({"mesh_decoder.cases": cases})
	diffusers_checkpoint.update({"mesh_decoder.masks": masks})

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:

diffusers/src/diffusers/pipelines/shap_e/renderer.py

Lines 921 to 944 in 0f1abc4

	# create cameras object
	camera = create_pan_cameras(size)
	rays = camera.camera_rays
	rays = rays.to(device)
	n_batches = rays.shape[1] // ray_batch_size
	coarse_sampler = StratifiedRaySampler()
	images = []
	for idx in range(n_batches):
	rays_batch = rays[:, idx * ray_batch_size : (idx + 1) * ray_batch_size]
	# render rays with coarse, stratified samples.
	_, fine_sampler, coarse_model_out = self.render_rays(rays_batch, coarse_sampler, n_coarse_samples)
	# Then, render with additional importance-weighted ray samples.
	channels, _, _ = self.render_rays(
	rays_batch, fine_sampler, n_fine_samples, prev_model_out=coarse_model_out
	)
	images.append(channels)
	images = torch.cat(images, dim=1)
	images = images.view(*camera.shape, camera.height, camera.width, -1).squeeze(0)

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:

diffusers/src/diffusers/pipelines/shap_e/renderer.py

Lines 393 to 400 in 0f1abc4

	mids = 0.5 * (ts[..., 1:] + ts[..., :-1])
	upper = torch.cat([mids, t1], dim=-1)
	lower = torch.cat([t0, mids], dim=-1)
	# yiyi notes: add a random seed here for testing, don't forget to remove
	torch.manual_seed(0)
	t_rand = torch.rand_like(ts)
	ts = lower + (upper - lower) * t_rand

diffusers/src/diffusers/pipelines/shap_e/renderer.py

Lines 456 to 460 in 0f1abc4

	t_rand = torch.rand(inds.shape, device=inds.device)
	lower_ = torch.gather(lower, -2, inds)
	upper_ = torch.gather(upper, -2, inds)
	ts = lower_ + (upper_ - lower_) * t_rand

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:

diffusers/src/diffusers/pipelines/shap_e/pipeline_shap_e.py

Lines 313 to 314 in 0f1abc4

	if output_type == "latent":
	return ShapEPipelineOutput(images=latents)

diffusers/src/diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py

Lines 300 to 301 in 0f1abc4

	if output_type == "latent":
	return ShapEPipelineOutput(images=latents)

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:

diffusers/src/diffusers/pipelines/flux/pipeline_flux.py

Lines 1003 to 1017 in 0f1abc4

	if output_type == "latent":
	image = latents
	else:
	latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
	latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
	image = self.vae.decode(latents, return_dict=False)[0]
	image = self.image_processor.postprocess(image, output_type=output_type)
	# Offload all models
	self.maybe_free_model_hooks()
	if not return_dict:
	return (image,)
	return FluxPipelineOutput(images=image)

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:

diffusers/src/diffusers/pipelines/shap_e/camera.py

Lines 104 to 118 in 0f1abc4

	def resize_image(self, width: int, height: int) -> "DifferentiableProjectiveCamera":
	"""
	Creates a new camera for the resized view assuming the aspect ratio does not change.
	"""
	assert width * self.height == height * self.width, "The aspect ratio should not change."
	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,
	)

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:

diffusers/tests/pipelines/shap_e/test_shap_e.py

Lines 225 to 231 in 0f1abc4

	@unittest.skip("Key error is raised with accelerate")
	def test_sequential_cpu_offload_forward_pass(self):
	pass
	@nightly
	@require_torch_accelerator

diffusers/tests/pipelines/shap_e/test_shap_e_img2img.py

Lines 248 to 254 in 0f1abc4

	@unittest.skip("Key error is raised with accelerate")
	def test_sequential_cpu_offload_forward_pass(self):
	pass
	@nightly
	@require_torch_accelerator

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.