huggingface / huggingface/diffusers

shap_e model/pipeline review

Offen
#13,593 0 Kommentare 0 Reaktionen 0 zugewiesene Personen Auf GitHub ansehen
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

Beitragsleitfaden öffnen

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

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.