huggingface / huggingface/diffusers

ddim model/pipeline review

Open
#13,591 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
34.5k
Forks
7.3k
Avg merge
3d 3h
Merged PRs (30d)
91

Description

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

Contributor guide

Open the contributing guide

Research direction

Start in src/diffusers/pipelines/ddim/pipeline_ddim.py and compare the referenced scheduler, pipeline utility, and consistency-model implementations. Run tests/pipelines/ddim/test_ddim.py, then add coverage for sigmoid and learned-variance DDPM conversion, offload cleanup, output-type validation, and the autodoc example. Done means these cases behave as described and the existing test suite passes.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, testing-qa
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
52/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.