huggingface / huggingface/diffusers

shap_e model/pipeline review

Open
#13,593 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
34.5k
Forks
7.3k
Avg merge
3d 3h
Merged PRs (30d)
91

Description

# `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 https://github.com/huggingface/diffusers/issues/4075 is only a docstring typo, closed issue https://github.com/huggingface/diffusers/issues/4808 is an old integration snapshot failure, and merged PR https://github.com/huggingface/diffusers/pull/4062 is relevant mesh precedent but does not cover the current gaps.

## Issue 1: `ShapEImg2ImgPipeline` drops or rejects documented image batch inputs

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py#L149-L153
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py#L225-L234

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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py#L1048

Suggested fix:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/shap_e/pipeline_shap_e.py#L130-L138
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py#L130-L140
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py#L250-L261

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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/flux2/pipeline_flux2.py#L644

Suggested fix:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/shap_e/renderer.py#L489-L495
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/docs/source/en/using-diffusers/shap-e.md#L135-L150

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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/scripts/convert_shap_e_to_diffusers.py#L830-L889

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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/shap_e/renderer.py#L921-L944

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:
```python
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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/shap_e/renderer.py#L393-L400
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/shap_e/renderer.py#L456-L460

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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/shap_e/pipeline_shap_e.py#L313-L314
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/shap_e/pipeline_shap_e_img2img.py#L300-L301

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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/flux/pipeline_flux.py#L1003-L1017

Suggested fix:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/shap_e/camera.py#L104-L118

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:
```python
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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/shap_e/test_shap_e.py#L225-L231
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/shap_e/test_shap_e_img2img.py#L248-L254

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:
```python
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.

Contributor guide

Open the contributing guide

Research direction

Start with the affected Shap-E files named in the report: pipeline_shap_e.py, pipeline_shap_e_img2img.py, renderer.py, camera.py, and their referenced tests or conversion script. Reproduce the listed failures one at a time, then trace the corresponding pipeline, renderer, and camera entry points. Done means each reported behavior is corrected with focused regression coverage without breaking existing checkpoint loading or output contracts.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.