huggingface / huggingface/diffusers

ddim model/pipeline review

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

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

Beitragsleitfaden

Beitragsleitfaden öffnen

Rechercherichtung

Beginne in src/diffusers/pipelines/ddim/pipeline_ddim.py und vergleiche die referenzierten Scheduler-, Pipeline-Utility- und Consistency-Model-Implementierungen. Führe tests/pipelines/ddim/test_ddim.py aus und ergänze Tests für die DDPM-Konvertierung mit Sigmoid und gelernter Varianz, die Bereinigung beim Offloading, die Validierung des Output-Typs und das autodoc-Beispiel. Als erledigt gilt die Aufgabe, wenn sich diese Fälle wie beschrieben verhalten und die bestehende Testsuite erfolgreich durchläuft.

Vom Indexierungsmodell aus dem Issue-Text verfasst.

Bewertung

Tech-Stack
python, pytorch
Bereich
machine-learning, testing-qa
Issue-Typ
Bug
Schwierigkeit
4/5
Geschätzter Aufwand
3-5 Tage
Aktivitätsstatus
Ruhig
Klarheit
Größtenteils klar
Anfängerfreundlichkeit
52/100

Neue Issues direkt in Ihr Postfach

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