huggingface / huggingface/diffusers

ddim model/pipeline review

Aperta
#13,591 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

# `ddim` model/pipeline review

Commit tested: `0f1abc4ae8b0eb2a3b40e82a310507281144c423`

Review performed against the repository review rules.

## Issue 1: `DDIMPipeline` only partially supports `DDPMScheduler` configs

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/ddim/pipeline_ddim.py#L57-L58
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/ddim/pipeline_ddim.py#L157-L164

Problem:
The pipeline says the scheduler can be `DDPMScheduler` or `DDIMScheduler`, and the constructor always converts the incoming scheduler to `DDIMScheduler`. That conversion is incomplete for valid DDPM configs: `beta_schedule="sigmoid"` raises during construction, and learned-variance DDPM UNets produce 2x channels that `DDIMScheduler.step()` cannot consume.

Impact:
Some valid DDPM unconditional checkpoints cannot be sampled through `DDIMPipeline`, despite this being the advertised replacement path for faster inference.

Reproduction:
```python
import torch
from diffusers import DDIMPipeline, DDPMScheduler, UNet2DModel

unet = UNet2DModel(
block_out_channels=(8, 8),
layers_per_block=1,
norm_num_groups=4,
sample_size=8,
in_channels=3,
out_channels=3,
down_block_types=("DownBlock2D", "DownBlock2D"),
up_block_types=("UpBlock2D", "UpBlock2D"),
)

try:
DDIMPipeline(unet=unet, scheduler=DDPMScheduler(num_train_timesteps=10, beta_schedule="sigmoid"))
except Exception as e:
print(type(e).__name__, str(e))

learned_var_unet = UNet2DModel(
block_out_channels=(8, 8),
layers_per_block=1,
norm_num_groups=4,
sample_size=8,
in_channels=3,
out_channels=6,
down_block_types=("DownBlock2D", "DownBlock2D"),
up_block_types=("UpBlock2D", "UpBlock2D"),
)
pipe = DDIMPipeline(
unet=learned_var_unet,
scheduler=DDPMScheduler(num_train_timesteps=10, variance_type="learned_range"),
)
pipe.set_progress_bar_config(disable=True)
pipe(num_inference_steps=2, output_type="np", generator=torch.Generator(device="cpu").manual_seed(0))
```

Relevant precedent:
`DDPMScheduler` supports the sigmoid schedule:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/schedulers/scheduling_ddpm.py#L232-L235

Other pipelines explicitly handle learned-variance channel splitting before scheduler stepping:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/kandinsky/pipeline_kandinsky.py#L375-L379

Suggested fix:
```python
# In DDIMScheduler.__init__, mirror DDPMScheduler:
elif beta_schedule == "sigmoid":
betas = torch.linspace(-6, 6, num_train_timesteps)
self.betas = torch.sigmoid(betas) * (beta_end - beta_start) + beta_start
```

```python
# In DDIMPipeline.__call__, before self.scheduler.step(...):
if model_output.shape[1] == image.shape[1] * 2:
model_output, _ = model_output.split(image.shape[1], dim=1)
```

Duplicate check:
No matching existing issue or PR found for the learned-variance or sigmoid DDPM conversion failures. Related older issue #1918 was about DDIM accidentally keeping a DDPM scheduler and passing `eta`, which is a different closed bug.

## Issue 2: CPU offload hooks are not freed after `__call__`

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/ddim/pipeline_ddim.py#L169-L176

Problem:
`DDIMPipeline` sets `model_cpu_offload_seq = "unet"` but never calls `self.maybe_free_model_hooks()` at the end of `__call__`. The shared pipeline utility explicitly requires this for correct `enable_model_cpu_offload()` behavior.

Impact:
After an offloaded DDIM run, the final model can remain resident on the accelerator instead of being restored to the expected CPU-offloaded state, increasing VRAM pressure across repeated calls or when chaining pipelines.

Reproduction:
```python
import torch
from diffusers import DDIMPipeline, DDIMScheduler, UNet2DModel

if not torch.cuda.is_available():
raise SystemExit("This reproducer needs an accelerator.")

unet = UNet2DModel(
block_out_channels=(8, 8),
layers_per_block=1,
norm_num_groups=4,
sample_size=8,
in_channels=3,
out_channels=3,
down_block_types=("DownBlock2D", "DownBlock2D"),
up_block_types=("UpBlock2D", "UpBlock2D"),
)
pipe = DDIMPipeline(unet=unet, scheduler=DDIMScheduler(num_train_timesteps=10))
pipe.set_progress_bar_config(disable=True)
pipe.enable_model_cpu_offload(device="cuda")

pipe(num_inference_steps=1, output_type="np", generator=torch.Generator(device="cpu").manual_seed(0))

print(next(pipe.unet.parameters()).device)
assert next(pipe.unet.parameters()).device.type == "cpu"
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/pipeline_utils.py#L1293-L1294
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/consistency_models/pipeline_consistency_models.py#L273-L277
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/dit/pipeline_dit.py#L240-L244

Suggested fix:
```python
image = (image / 2 + 0.5).clamp(0, 1)
image = image.cpu().permute(0, 2, 3, 1).numpy()
if output_type == "pil":
image = self.numpy_to_pil(image)

self.maybe_free_model_hooks()

if not return_dict:
return (image,)
```

Duplicate check:
No DDIM-specific existing issue or PR found. GitHub search found unrelated CPU-offload/context-parallelism issue #12533, but not this missing DDIM cleanup call.

## Issue 3: Output handling and the autodoc example are inconsistent

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/ddim/pipeline_ddim.py#L97-L117
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/ddim/pipeline_ddim.py#L169-L172

Problem:
The docstring example calls `.cpu()` on the pipeline return value, but `DDIMPipeline.__call__` returns `ImagePipelineOutput`, not a tensor. Separately, unsupported `output_type` values are silently treated as NumPy output because only `"pil"` is special-cased.

Impact:
The generated API docs contain an example that fails as written, and invalid output-type mistakes are accepted silently instead of producing an actionable error.

Reproduction:
```python
import torch
from diffusers import DDIMPipeline, DDIMScheduler, UNet2DModel

unet = UNet2DModel(
block_out_channels=(8, 8),
layers_per_block=1,
norm_num_groups=4,
sample_size=8,
in_channels=3,
out_channels=3,
down_block_types=("DownBlock2D", "DownBlock2D"),
up_block_types=("UpBlock2D", "UpBlock2D"),
)
pipe = DDIMPipeline(unet=unet, scheduler=DDIMScheduler(num_train_timesteps=10))
pipe.set_progress_bar_config(disable=True)

out = pipe(num_inference_steps=1, output_type="pt", generator=torch.Generator(device="cpu").manual_seed(0)).images
print(type(out)) # numpy.ndarray, despite requesting "pt"

pipe(num_inference_steps=1, generator=torch.Generator(device="cpu").manual_seed(0)).cpu()
```

Relevant precedent:
`ConsistencyModelPipeline` validates and handles postprocessing explicitly:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/consistency_models/pipeline_consistency_models.py#L109-L126

Suggested fix:
```python
if output_type not in ["pil", "np"]:
raise ValueError("`output_type` must be one of ['pil', 'np'].")

image = (image / 2 + 0.5).clamp(0, 1)
image = image.cpu().permute(0, 2, 3, 1).numpy()
if output_type == "pil":
image = self.numpy_to_pil(image)
```

Also update the example to use:
```python
image = pipe(eta=0.0, num_inference_steps=50).images[0]
image.save("test.png")
```

Duplicate check:
No matching existing issue or PR found for the broken DDIM docstring example or silent `output_type` fallback.

## Coverage / Search Status

Fast tests exist at `tests/pipelines/ddim/test_ddim.py` and slow tests exist under `DDIMPipelineIntegrationTests`, so slow coverage is not missing. Current tests do not cover DDPM sigmoid conversion, learned-variance DDPM conversion, or DDIM-specific offload cleanup.

I attempted `./.venv/Scripts/python.exe -m pytest tests/pipelines/ddim/test_ddim.py -q`, but collection failed in this environment because the installed PyTorch build lacks `torch._C._distributed_c10d`. The standalone CPU reproductions above were run with `.venv`.

Duplicate checks were run against GitHub Issues and PRs for `DDIMPipeline`, `pipeline_ddim.py`, learned variance, sigmoid beta schedule, CPU offload, and output type failure modes.

Guida per i contributori

Apri la guida per i contributori

Direzione di ricerca

Inizia in src/diffusers/pipelines/ddim/pipeline_ddim.py e confronta le implementazioni di scheduler, utility della pipeline e modello di consistenza a cui si fa riferimento. Esegui tests/pipelines/ddim/test_ddim.py, quindi aggiungi la copertura per la conversione DDPM con sigmoid e varianza appresa, la pulizia dell’offload, la validazione del tipo di output e l’esempio autodoc. Il lavoro è completo quando questi casi si comportano come descritto e la suite di test esistente passa.

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

Valutazione

Stack tecnologico
python, pytorch
Ambito
machine-learning, testing-qa
Tipo di issue
Bug
Difficoltà
4/5
Tempo stimato
3-5 giorni
Stato di attività
Tranquilla
Chiarezza
Abbastanza chiara
Idoneità per principianti
52/100

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.