ddim model/pipeline review

[hlky]

ddim model/pipeline review

Commit tested: 0f1abc4ae8b0eb2a3b40e82a310507281144c423

Review performed against the repository review rules.

Issue 1: DDIMPipeline only partially supports DDPMScheduler configs

Affected code:

diffusers/src/diffusers/pipelines/ddim/pipeline_ddim.py

Lines 57 to 58 in 0f1abc4

	# make sure scheduler can always be converted to DDIM
	scheduler = DDIMScheduler.from_config(scheduler.config)

diffusers/src/diffusers/pipelines/ddim/pipeline_ddim.py

Lines 157 to 164 in 0f1abc4

	model_output = self.unet(image, t).sample
	# 2. predict previous mean of image x_t-1 and add variance depending on eta
	# eta corresponds to η in paper and should be between [0, 1]
	# do x_t -> x_t-1
	image = self.scheduler.step(
	model_output, t, image, eta=eta, use_clipped_model_output=use_clipped_model_output, generator=generator
	).prev_sample

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:

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:

diffusers/src/diffusers/schedulers/scheduling_ddpm.py

Lines 232 to 235 in 0f1abc4

	elif beta_schedule == "sigmoid":
	# GeoDiff sigmoid schedule
	betas = torch.linspace(-6, 6, num_train_timesteps)
	self.betas = torch.sigmoid(betas) * (beta_end - beta_start) + beta_start

Other pipelines explicitly handle learned-variance channel splitting before scheduler stepping:

diffusers/src/diffusers/pipelines/kandinsky/pipeline_kandinsky.py

Lines 375 to 379 in 0f1abc4

	if not (
	hasattr(self.scheduler.config, "variance_type")
	and self.scheduler.config.variance_type in ["learned", "learned_range"]
	):
	noise_pred, _ = noise_pred.split(latents.shape[1], dim=1)

Suggested fix:

# 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

# 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:

diffusers/src/diffusers/pipelines/ddim/pipeline_ddim.py

Lines 169 to 176 in 0f1abc4

	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)
	if not return_dict:
	return (image,)

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:

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:

diffusers/src/diffusers/pipelines/pipeline_utils.py

Lines 1293 to 1294 in 0f1abc4

	Make sure to add this function to the end of the `__call__` function of your pipeline so that it functions
	correctly when applying `enable_model_cpu_offload`.

diffusers/src/diffusers/pipelines/consistency_models/pipeline_consistency_models.py

Lines 273 to 277 in 0f1abc4

	# 6. Post-process image sample
	image = self.postprocess_image(sample, output_type=output_type)
	# Offload all models
	self.maybe_free_model_hooks()

diffusers/src/diffusers/pipelines/dit/pipeline_dit.py

Lines 240 to 244 in 0f1abc4

	if output_type == "pil":
	samples = self.numpy_to_pil(samples)
	# Offload all models
	self.maybe_free_model_hooks()

Suggested fix:

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:

diffusers/src/diffusers/pipelines/ddim/pipeline_ddim.py

Lines 97 to 117 in 0f1abc4

	Example:
	```py
	>>> from diffusers import DDIMPipeline
	>>> import PIL.Image
	>>> import numpy as np
	>>> # load model and scheduler
	>>> pipe = DDIMPipeline.from_pretrained("fusing/ddim-lsun-bedroom")
	>>> # run pipeline in inference (sample random noise and denoise)
	>>> image = pipe(eta=0.0, num_inference_steps=50)
	>>> # process image to PIL
	>>> image_processed = image.cpu().permute(0, 2, 3, 1)
	>>> image_processed = (image_processed + 1.0) * 127.5
	>>> image_processed = image_processed.numpy().astype(np.uint8)
	>>> image_pil = PIL.Image.fromarray(image_processed[0])
	>>> # save image
	>>> image_pil.save("test.png")

diffusers/src/diffusers/pipelines/ddim/pipeline_ddim.py

Lines 169 to 172 in 0f1abc4

	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)

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:

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:

diffusers/src/diffusers/pipelines/consistency_models/pipeline_consistency_models.py

Lines 109 to 126 in 0f1abc4

	return latents
	# Follows diffusers.VaeImageProcessor.postprocess
	def postprocess_image(self, sample: torch.Tensor, output_type: str = "pil"):
	if output_type not in ["pt", "np", "pil"]:
	raise ValueError(
	f"output_type={output_type} is not supported. Make sure to choose one of ['pt', 'np', or 'pil']"
	)
	# Equivalent to diffusers.VaeImageProcessor.denormalize
	sample = (sample / 2 + 0.5).clamp(0, 1)
	if output_type == "pt":
	return sample
	# Equivalent to diffusers.VaeImageProcessor.pt_to_numpy
	sample = sample.cpu().permute(0, 2, 3, 1).numpy()
	if output_type == "np":
	return sample

Suggested fix:

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:

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.