stable_diffusion model/pipeline review

[hlky]

stable_diffusion model/pipeline review

Commit tested: 0f1abc4ae8b0eb2a3b40e82a310507281144c423

Review performed against the repository review rules. .ai/review-rules.md references AGENTS.md, but that file was absent in this checkout; all available referenced rule files were applied.

Duplicate search: searched GitHub Issues and PRs in huggingface/diffusers for stable_diffusion, affected class/function/file names, and failure modes. No duplicates found for Issues 1-7. Issue 8 has related coverage context in #11762 and #9371.

Issue 1: Stable Diffusion subpackage ONNX lazy imports are broken

Affected code:

diffusers/src/diffusers/pipelines/stable_diffusion/__init__.py

Lines 68 to 83 in 0f1abc4

	try:
	if not (is_transformers_available() and is_onnx_available()):
	raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
	from ...utils import dummy_onnx_objects # noqa F403
	_dummy_objects.update(get_objects_from_module(dummy_onnx_objects))
	else:
	_import_structure["pipeline_onnx_stable_diffusion"] = [
	"OnnxStableDiffusionPipeline",
	"StableDiffusionOnnxPipeline",
	]
	_import_structure["pipeline_onnx_stable_diffusion_img2img"] = ["OnnxStableDiffusionImg2ImgPipeline"]
	_import_structure["pipeline_onnx_stable_diffusion_inpaint"] = ["OnnxStableDiffusionInpaintPipeline"]
	_import_structure["pipeline_onnx_stable_diffusion_inpaint_legacy"] = ["OnnxStableDiffusionInpaintPipelineLegacy"]
	_import_structure["pipeline_onnx_stable_diffusion_upscale"] = ["OnnxStableDiffusionUpscalePipeline"]

Problem:
When ONNX is unavailable, diffusers.pipelines.stable_diffusion imports only dummy_onnx_objects, which contains OnnxRuntimeModel but not the Stable Diffusion ONNX pipeline dummies. Subpackage imports fail with ImportError instead of returning dummy classes with backend errors. The same block also advertises pipeline_onnx_stable_diffusion_inpaint_legacy, but that module does not exist under pipelines/stable_diffusion.

Impact:
Users importing from the public subpackage get inconsistent behavior compared with from diffusers import OnnxStableDiffusionPipeline. If ONNX is installed, the legacy lazy entry points at a nonexistent module.

Reproduction:

from diffusers.pipelines.stable_diffusion import OnnxStableDiffusionPipeline

Relevant precedent:
src/diffusers/pipelines/__init__.py uses dummy_torch_and_transformers_and_onnx_objects for these classes.

Suggested fix:

except OptionalDependencyNotAvailable:
 from ...utils import dummy_torch_and_transformers_and_onnx_objects
 _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_and_onnx_objects))
else:
 ...
 # Drop this nonexistent stable_diffusion module entry, or route it through deprecated exports.
 # _import_structure["pipeline_onnx_stable_diffusion_inpaint_legacy"] = [...]

Issue 2: ONNX upscaler ignores user-provided latents

Affected code:

diffusers/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_upscale.py

Lines 470 to 477 in 0f1abc4

	latents = self.prepare_latents(
	batch_size * num_images_per_prompt,
	self.config.num_latent_channels,
	height,
	width,
	latents_dtype,
	generator,
	)

Problem:
OnnxStableDiffusionUpscalePipeline.__call__ accepts latents, but does not forward it to prepare_latents, so custom latents are silently discarded.

Impact:
Users cannot reproduce or edit generations with externally prepared latents, unlike the Torch upscaler.

Reproduction:

import types, numpy as np, torch
from types import SimpleNamespace
from diffusers.pipelines.stable_diffusion.pipeline_onnx_stable_diffusion_upscale import OnnxStableDiffusionUpscalePipeline

class Fake:
 config = SimpleNamespace(num_latent_channels=4, num_unet_input_channels=7)
 safety_checker = None
 def check_inputs(self, *a, **k): pass
 def _encode_prompt(self, *a, **k): return np.zeros((1, 77, 8), dtype=np.float32)

pipe = Fake()
custom = np.zeros((1, 4, 64, 64), dtype=np.float32)
def prepare_latents(self, *args, latents=None):
 assert latents is custom, f"latents was dropped: {latents!r}"
pipe.prepare_latents = types.MethodType(prepare_latents, pipe)

OnnxStableDiffusionUpscalePipeline.__call__(
 pipe, prompt="x", image=torch.zeros(1, 3, 64, 64),
 num_inference_steps=1, generator=np.random.RandomState(0), latents=custom,
)

Relevant precedent:

diffusers/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py

Lines 712 to 721 in 0f1abc4

	latents = self.prepare_latents(
	batch_size * num_images_per_prompt,
	num_channels_latents,
	height,
	width,
	prompt_embeds.dtype,
	device,
	generator,
	latents,
	)

Suggested fix:

latents = self.prepare_latents(
 batch_size * num_images_per_prompt,
 self.config.num_latent_channels,
 height,
 width,
 latents_dtype,
 generator,
 latents=latents,
)

Issue 3: ONNX upscaler crashes when classifier-free guidance is disabled

Affected code:

diffusers/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion_upscale.py

Lines 543 to 546 in 0f1abc4

	# perform guidance
	if do_classifier_free_guidance:
	noise_pred_uncond, noise_pred_text = np.split(noise_pred, 2)
	noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

Problem:
noise_pred is recomputed from noise_pred_uncond and noise_pred_text outside the if do_classifier_free_guidance block. With guidance_scale <= 1.0, those variables are never assigned.

Impact:
OnnxStableDiffusionUpscalePipeline(..., guidance_scale=1.0) fails instead of running unconditional/no-CFG inference.

Reproduction:

import numpy as np, torch
from types import SimpleNamespace
from diffusers.pipelines.stable_diffusion.pipeline_onnx_stable_diffusion_upscale import OnnxStableDiffusionUpscalePipeline

class Bar:
 def __enter__(self): return self
 def __exit__(self, *a): pass
 def update(self): pass

class Scheduler:
 order = 1
 init_noise_sigma = 1.0
 def set_timesteps(self, n): self.timesteps = [np.array(1, dtype=np.float32)]
 def scale_model_input(self, sample, t): return sample

class Fake:
 config = SimpleNamespace(num_latent_channels=4, num_unet_input_channels=7)
 scheduler = Scheduler()
 low_res_scheduler = SimpleNamespace(add_noise=lambda image, noise, noise_level: image)
 unet = SimpleNamespace(
 model=SimpleNamespace(get_inputs=lambda: [SimpleNamespace(name="timestep", type="tensor(float)")]),
 __call__=lambda **kw: [np.zeros((1, 4, 64, 64), dtype=np.float32)],
 )
 safety_checker = None
 def check_inputs(self, *a, **k): pass
 def _encode_prompt(self, *a, **k): return np.zeros((1, 77, 8), dtype=np.float32)
 def prepare_latents(self, *a, **k): return np.zeros((1, 4, 64, 64), dtype=np.float32)
 def progress_bar(self, total): return Bar()

OnnxStableDiffusionUpscalePipeline.__call__(
 Fake(), prompt="x", image=torch.zeros(1, 3, 64, 64),
 num_inference_steps=1, guidance_scale=1.0, output_type="np",
)

Relevant precedent:

diffusers/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py

Lines 744 to 747 in 0f1abc4

	# concat latents, mask, masked_image_latents in the channel dimension
	latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
	latent_model_input = torch.cat([latent_model_input, image], dim=1)

Suggested fix:

if do_classifier_free_guidance:
 noise_pred_uncond, noise_pred_text = np.split(noise_pred, 2)
 noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

Issue 4: StableUnCLIPImageNormalizer breaks standard .to() kwargs

Affected code:

diffusers/src/diffusers/pipelines/stable_diffusion/stable_unclip_image_normalizer.py

Lines 40 to 47 in 0f1abc4

	def to(
	self,
	torch_device: str | torch.device | None = None,
	torch_dtype: torch.dtype | None = None,
	):
	self.mean = nn.Parameter(self.mean.to(torch_device).to(torch_dtype))
	self.std = nn.Parameter(self.std.to(torch_device).to(torch_dtype))
	return self

Problem:
The custom .to() signature accepts only torch_device and torch_dtype, so standard calls like .to(dtype=torch.float16) or .to(device="cuda") raise TypeError.

Impact:
This violates nn.Module/ModelMixin behavior and makes the component harder to use or move independently.

Reproduction:

import torch
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer

StableUnCLIPImageNormalizer().to(dtype=torch.float16)

Relevant precedent:
Other ModelMixin modules rely on inherited nn.Module.to.

Suggested fix:

# Remove the override entirely; registered Parameters move with nn.Module.to.
# Or, if keeping it:
def to(self, *args, **kwargs):
 return super().to(*args, **kwargs)

Issue 5: UNet QKV unfuse state is not safe

Affected code:

diffusers/src/diffusers/models/unets/unet_2d_condition.py

Lines 821 to 849 in 0f1abc4

	def fuse_qkv_projections(self):
	"""
	Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value)
	are fused. For cross-attention modules, key and value projection matrices are fused.
	> [!WARNING] > This API is 🧪 experimental.
	"""
	self.original_attn_processors = None
	for _, attn_processor in self.attn_processors.items():
	if "Added" in str(attn_processor.__class__.__name__):
	raise ValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.")
	self.original_attn_processors = self.attn_processors
	for module in self.modules():
	if isinstance(module, Attention):
	module.fuse_projections(fuse=True)
	self.set_attn_processor(FusedAttnProcessor2_0())
	def unfuse_qkv_projections(self):
	"""Disables the fused QKV projection if enabled.
	> [!WARNING] > This API is 🧪 experimental.
	"""
	if self.original_attn_processors is not None:
	self.set_attn_processor(self.original_attn_processors)

Problem:
original_attn_processors is not initialized in __init__, so unfuse_qkv_projections() before fuse_qkv_projections() raises AttributeError. Calling fuse_qkv_projections() twice also overwrites the saved original processors with fused processors, so unfuse_qkv_projections() cannot restore the original state.

Impact:
A public optimization API is not idempotent and can leave the model fused permanently in common enable-twice-then-disable flows.

Reproduction:

from diffusers import UNet2DConditionModel

model = UNet2DConditionModel(
 block_out_channels=(4, 8), norm_num_groups=4,
 down_block_types=("CrossAttnDownBlock2D", "DownBlock2D"),
 up_block_types=("UpBlock2D", "CrossAttnUpBlock2D"),
 cross_attention_dim=8, attention_head_dim=2,
 out_channels=4, in_channels=4, layers_per_block=1, sample_size=16,
)

model.unfuse_qkv_projections() # AttributeError

model.fuse_qkv_projections()
model.fuse_qkv_projections()
model.unfuse_qkv_projections()
print({p.__class__.__name__ for p in model.attn_processors.values()}) # still fused

Relevant precedent:
The method doc says it disables fused projections “if enabled”.

Suggested fix:

# in __init__
self.original_attn_processors = None

# in fuse_qkv_projections
if self.original_attn_processors is None:
 self.original_attn_processors = self.attn_processors

# in unfuse_qkv_projections
if self.original_attn_processors is not None:
 self.set_attn_processor(self.original_attn_processors)
 self.original_attn_processors = None

Issue 6: Shorter UNet cross-attention masks crash instead of padding correctly

Affected code:

diffusers/src/diffusers/models/unets/unet_2d_condition.py

Lines 1073 to 1076 in 0f1abc4

	# convert encoder_attention_mask to a bias the same way we do for attention_mask
	if encoder_attention_mask is not None:
	encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0
	encoder_attention_mask = encoder_attention_mask.unsqueeze(1)

diffusers/src/diffusers/models/attention_processor.py

Lines 729 to 742 in 0f1abc4

	current_length: int = attention_mask.shape[-1]
	if current_length != target_length:
	if attention_mask.device.type == "mps":
	# HACK: MPS: Does not support padding by greater than dimension of input tensor.
	# Instead, we can manually construct the padding tensor.
	padding_shape = (attention_mask.shape[0], attention_mask.shape[1], target_length)
	padding = torch.zeros(padding_shape, dtype=attention_mask.dtype, device=attention_mask.device)
	attention_mask = torch.cat([attention_mask, padding], dim=2)
	else:
	# TODO: for pipelines such as stable-diffusion, padding cross-attn mask:
	# we want to instead pad by (0, remaining_length), where remaining_length is:
	# remaining_length: int = target_length - current_length
	# TODO: re-enable tests/models/test_models_unet_2d_condition.py#test_model_xattn_padding
	attention_mask = F.pad(attention_mask, (0, target_length), value=0.0)

Problem:
UNet2DConditionModel accepts encoder_attention_mask, but a mask shorter than encoder_hidden_states crashes because prepare_attention_mask pads by target_length instead of the remaining length. The repo already has this case as a skipped test.

Impact:
Callers cannot pass shortened cross-attention masks even though the attention code has padding logic for mismatched mask lengths.

Reproduction:

import torch
from diffusers import UNet2DConditionModel

torch.manual_seed(0)
model = UNet2DConditionModel(
 block_out_channels=(4, 8), norm_num_groups=4,
 down_block_types=("CrossAttnDownBlock2D", "DownBlock2D"),
 up_block_types=("UpBlock2D", "CrossAttnUpBlock2D"),
 cross_attention_dim=8, attention_head_dim=2,
 out_channels=4, in_channels=4, layers_per_block=1, sample_size=16,
).eval()

sample = torch.randn(1, 4, 16, 16)
cond = torch.randn(1, 4, 8)
short_mask = torch.zeros(1, 3, dtype=torch.bool)

with torch.no_grad():
 model(sample, torch.tensor([10]), cond, encoder_attention_mask=short_mask)

Relevant precedent:

diffusers/tests/models/unets/test_models_unet_2d_condition.py

Lines 678 to 705 in 0f1abc4

	# see diffusers.models.attention_processor::Attention#prepare_attention_mask
	# note: we may not need to fix mask padding to work for stable-diffusion cross-attn masks.
	# since the use-case (somebody passes in a too-short cross-attn mask) is pretty esoteric.
	# maybe it's fine that this only works for the unclip use-case.
	@mark.skip(
	reason="we currently pad mask by target_length tokens (what unclip needs), whereas stable-diffusion's cross-attn needs to instead pad by remaining_length."
	)
	def test_model_xattn_padding(self):
	init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
	model = self.model_class(**{**init_dict, "attention_head_dim": (8, 16)})
	model.to(torch_device)
	model.eval()
	cond = inputs_dict["encoder_hidden_states"]
	with torch.no_grad():
	full_cond_out = model(**inputs_dict).sample
	assert full_cond_out is not None
	batch, tokens, _ = cond.shape
	keeplast_mask = (torch.arange(tokens) == tokens - 1).expand(batch, -1).to(cond.device, torch.bool)
	keeplast_out = model(**{**inputs_dict, "encoder_attention_mask": keeplast_mask}).sample
	assert not keeplast_out.allclose(full_cond_out), "a 'keep last token' mask should change the result"
	trunc_mask = torch.zeros(batch, tokens - 1, device=cond.device, dtype=torch.bool)
	trunc_mask_out = model(**{**inputs_dict, "encoder_attention_mask": trunc_mask}).sample
	assert trunc_mask_out.allclose(keeplast_out), (
	"a mask with fewer tokens than condition, will be padded with 'keep' tokens. a 'discard-all' mask missing the final token is thus equivalent to a 'keep last' mask."

Suggested fix:
The fix is slightly risky because comments mention UnCLIP compatibility. Add separate tests for SD cross-attn masks and UnCLIP added-KV masks, then pad by target_length - current_length for the SD cross-attn case rather than by target_length.

Issue 7: checkpoint conversion ignores local_files_only in some SD branches

Affected code:

diffusers/src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py

Lines 1556 to 1562 in 0f1abc4

	elif pipeline_class == StableDiffusionUpscalePipeline:
	scheduler = DDIMScheduler.from_pretrained(
	"stabilityai/stable-diffusion-x4-upscaler", subfolder="scheduler"
	)
	low_res_scheduler = DDPMScheduler.from_pretrained(
	"stabilityai/stable-diffusion-x4-upscaler", subfolder="low_res_scheduler"
	)

diffusers/src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py

Lines 1588 to 1594 in 0f1abc4

	else:
	image_normalizer, image_noising_scheduler = stable_unclip_image_noising_components(
	original_config, clip_stats_path=clip_stats_path, device=device
	)
	if stable_unclip == "img2img":
	feature_extractor, image_encoder = stable_unclip_image_encoder(original_config)

Problem:
download_from_original_stable_diffusion_ckpt(..., local_files_only=True) still calls the x4 upscaler scheduler loaders without local_files_only, and the Stable UnCLIP img2img branch calls stable_unclip_image_encoder(original_config) without forwarding local_files_only.

Impact:
Offline/local-only conversion can unexpectedly hit the network or fail with less useful errors.

Reproduction:

import ast, inspect
import diffusers.pipelines.stable_diffusion.convert_from_ckpt as c

tree = ast.parse(inspect.getsource(c.download_from_original_stable_diffusion_ckpt))
calls = [n for n in ast.walk(tree) if isinstance(n, ast.Call)]
print(any(getattr(getattr(n, "func", None), "id", "") == "stable_unclip_image_encoder"
 and not any(k.arg == "local_files_only" for k in n.keywords) for n in calls))

Relevant precedent:
Nearby loaders in the same function pass local_files_only=local_files_only.

Suggested fix:

scheduler = DDIMScheduler.from_pretrained(
 "stabilityai/stable-diffusion-x4-upscaler", subfolder="scheduler", local_files_only=local_files_only
)
low_res_scheduler = DDPMScheduler.from_pretrained(
 "stabilityai/stable-diffusion-x4-upscaler", subfolder="low_res_scheduler", local_files_only=local_files_only
)

feature_extractor, image_encoder = stable_unclip_image_encoder(
 original_config, local_files_only=local_files_only
)

Issue 8: Missing or disabled slow coverage for parts of the target family

Affected code:

diffusers/tests/pipelines/stable_unclip/test_stable_unclip.py

Lines 202 to 204 in 0f1abc4

	@nightly
	@require_torch_accelerator
	class StableUnCLIPPipelineIntegrationTests(unittest.TestCase):

diffusers/tests/pipelines/stable_unclip/test_stable_unclip_img2img.py

Lines 219 to 221 in 0f1abc4

	@nightly
	@require_torch_accelerator
	class StableUnCLIPImg2ImgPipelineIntegrationTests(unittest.TestCase):

diffusers/tests/pipelines/stable_diffusion/test_onnx_stable_diffusion_upscale.py

Lines 45 to 51 in 0f1abc4

	# TODO: (Dhruv) Update hub_checkpoint repo_id
	@unittest.skip(
	"There is a potential backdoor vulnerability in the hub_checkpoint. Skip running this test until resolved"
	)
	class OnnxStableDiffusionUpscalePipelineFastTests(OnnxPipelineTesterMixin, unittest.TestCase):
	# TODO: is there an appropriate internal test set?
	hub_checkpoint = "ssube/stable-diffusion-x4-upscaler-onnx"

Problem:
Stable UnCLIP integration tests are @nightly but not @slow. Flax SD img2img/inpaint have no direct test files. ONNX upscaler fast tests are entirely skipped, and integration coverage is nightly-only.

Impact:
Release slow CI can miss regressions in these target pipelines. This likely allowed the ONNX upscaler no-CFG and ignored-latents regressions above to survive.

Reproduction:

from pathlib import Path

checks = {
 "stable_unclip": "tests/pipelines/stable_unclip/test_stable_unclip.py",
 "stable_unclip_img2img": "tests/pipelines/stable_unclip/test_stable_unclip_img2img.py",
 "onnx_upscale": "tests/pipelines/stable_diffusion/test_onnx_stable_diffusion_upscale.py",
}
for name, path in checks.items():
 text = Path(path).read_text()
 print(name, "@slow" in text, "@nightly" in text, "@unittest.skip" in text)

Relevant precedent:
Core Torch SD, img2img, inpaint, depth, latent-upscale, and x4-upscale have slow tests. Related existing issues: #11762 documents the risky ONNX upscaler checkpoint that led to skipped tests; #9371 is an open Flax img2img API cleanup request but does not cover slow coverage.

Suggested fix:
Add @slow integration tests for Stable UnCLIP and Stable UnCLIP img2img, add direct Flax img2img/inpaint coverage or explicitly deprecate those pipelines, and replace the skipped ONNX upscaler checkpoint with a safe internal tiny ONNX fixture that covers guidance_scale=1.0 and custom latents.

Local checks: utils/check_copies.py passed. Targeted pytest collection was blocked by this .venv Torch build missing torch._C._distributed_c10d when importing test utilities.