huggingface / huggingface/diffusers
shap_e model/pipeline review
- Vorherrschende Sprache
- Python
- Sterne
- 34.5k
- Forks
- 7.3k
- Ø Merge
- 3 T. 3 Std.
- Gemergte PRs (30 T.)
- 91
Beschreibung
# `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.
Beitragsleitfaden
Rechercherichtung
Beginne mit den im Bericht genannten betroffenen Shap-E-Dateien: pipeline_shap_e.py, pipeline_shap_e_img2img.py, renderer.py, camera.py und den von ihnen referenzierten Tests oder dem Konvertierungsskript. Reproduziere die aufgeführten Fehler nacheinander und verfolge anschließend die entsprechenden Einstiegspunkte der Pipeline, des Renderers und der Kamera. Als abgeschlossen gilt die Arbeit, wenn jedes gemeldete Verhalten durch gezielte Regressionstests korrigiert ist, ohne das Laden bestehender Checkpoints oder die Ausgabeverträge zu beeinträchtigen.
Vom Indexierungsmodell aus dem Issue-Text verfasst.
Bewertung
- Tech-Stack
- python, pytorch
- Bereich
- machine-learning
- Issue-Typ
- Bug
- Schwierigkeit
- 5/5
- Geschätzter Aufwand
- Über eine Woche
- Aktivitätsstatus
- Ruhig
- Klarheit
- Größtenteils klar
- Anfängerfreundlichkeit
- 25/100