huggingface / huggingface/diffusers

shap_e model/pipeline review

Aperta
#13,593 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub
Lingua principale
Python
Stelle
34.5k
Fork
7.3k
Merge medio
3g 3h
PR unite (30g)
91

Descrizione

# `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.

Guida per i contributori

Apri la guida per i contributori

Direzione di ricerca

Inizia dai file Shap-E interessati indicati nel report: pipeline_shap_e.py, pipeline_shap_e_img2img.py, renderer.py, camera.py e i relativi test o script di conversione referenziati. Riproduci i malfunzionamenti elencati uno alla volta, quindi traccia i corrispondenti punti di ingresso della pipeline, del renderer e della camera. Il lavoro è completo quando ogni comportamento segnalato è corretto con una copertura di regressione mirata, senza compromettere il caricamento dei checkpoint esistenti o i contratti di output.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Valutazione

Stack tecnologico
python, pytorch
Ambito
machine-learning
Tipo di issue
Bug
Difficoltà
5/5
Tempo stimato
Più di una settimana
Stato di attività
Tranquilla
Chiarezza
Abbastanza chiara
Idoneità per principianti
25/100

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.