Skip to content

stable_cascade model/pipeline review #13589

Description

@hlky

stable_cascade model/pipeline review

Commit tested: 0f1abc4ae8b0eb2a3b40e82a310507281144c423

Review performed against the repository review rules.

Reviewed: public imports/lazy loading, config/loading/single-file surface, runtime dtype/device paths, offload/xFormers delegation, model forward behavior, docs, fast/slow tests, and duplicate issues/PRs.

Duplicate search status: searched huggingface/diffusers Issues/PRs for stable_cascade, class names, prompt_embeds_pooled, image-conditioning inputs, latents dtype, timesteps, xFormers, and coverage. No exact duplicates found. Related but not duplicate: #7355/#7644 for older arbitrary-resolution dtype errors, and #7598 for bf16 image tensor preprocessing.

Issue 1: Single image conditioning inputs are advertised but rejected

Affected code:

defencode_image(self, images, device, dtype, batch_size, num_images_per_prompt):
image_embeds= []
forimageinimages:
image=self.feature_extractor(image, return_tensors="pt").pixel_values
image=image.to(device=device, dtype=dtype)
image_embed=self.image_encoder(image).image_embeds.unsqueeze(1)
image_embeds.append(image_embed)
image_embeds=torch.cat(image_embeds, dim=1)
image_embeds=image_embeds.repeat(batch_size*num_images_per_prompt, 1, 1)
negative_image_embeds=torch.zeros_like(image_embeds)
returnimage_embeds, negative_image_embeds

ifimages:
fori, imageinenumerate(images):
ifnotisinstance(image, torch.Tensor) andnotisinstance(image, PIL.Image.Image):
raiseTypeError(
f"'images' must contain images of type 'torch.Tensor' or 'PIL.Image.Image, but got"
f"{type(image)} for image number {i}."
)

prompt: str|list[str] |None=None,
images: torch.Tensor|PIL.Image.Image|list[torch.Tensor] |list[PIL.Image.Image] =None,
height: int=512,

Problem:
StableCascadePriorPipeline and StableCascadeCombinedPipeline type-hint images as a single tensor/PIL image or a list, but check_inputs() uses if images: and then iterates images. A tensor raises ambiguous-bool errors, and a single PIL image raises 'Image' object is not iterable.

Impact:
Image-conditioned Stable Cascade calls fail for documented input forms unless users wrap the image in a list.

Reproduction:

fromPILimportImageimporttorchfromdiffusersimportDDPMWuerstchenScheduler, StableCascadePriorPipeline, StableCascadeUNetprior=StableCascadeUNet(
conditioning_dim=8, block_out_channels=(8,), num_attention_heads=(-1,),
down_num_layers_per_block=(1,), up_num_layers_per_block=(1,),
down_blocks_repeat_mappers=(1,), up_blocks_repeat_mappers=(1,),
block_types_per_layer=(("SDCascadeResBlock", "SDCascadeTimestepBlock"),),
clip_text_pooled_in_channels=8, clip_image_in_channels=8,
)
pipe=StableCascadePriorPipeline(None, None, prior, DDPMWuerstchenScheduler())
forimagesin [Image.new("RGB", (8, 8)), torch.zeros(3, 8, 8)]:
try:
pipe.check_inputs(prompt="cat", images=images)
exceptExceptionase:
print(type(e).__name__, str(e).splitlines()[0])

Relevant precedent:
Other image pipelines normalize single images to lists before iterating, for example pipeline_if_img2img.py.

Suggested fix:

ifimagesisnotNoneandnotisinstance(images, list):
images= [images]

Apply before validation and before encode_image().

Issue 2: Decoder precomputed prompt embeds break CFG unless negative pooled embeds are also supplied

Affected code:

defcheck_inputs(
self,
prompt,
negative_prompt=None,
prompt_embeds=None,
negative_prompt_embeds=None,
callback_on_step_end_tensor_inputs=None,
):
ifcallback_on_step_end_tensor_inputsisnotNoneandnotall(
kinself._callback_tensor_inputsforkincallback_on_step_end_tensor_inputs
):
raiseValueError(
f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[kforkincallback_on_step_end_tensor_inputsifknotinself._callback_tensor_inputs]}"
)
ifpromptisnotNoneandprompt_embedsisnotNone:
raiseValueError(
f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
" only forward one of the two."
)
elifpromptisNoneandprompt_embedsisNone:
raiseValueError(
"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
)
elifpromptisnotNoneand (notisinstance(prompt, str) andnotisinstance(prompt, list)):
raiseValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
ifnegative_promptisnotNoneandnegative_prompt_embedsisnotNone:
raiseValueError(
f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
)
ifprompt_embedsisnotNoneandnegative_prompt_embedsisnotNone:
ifprompt_embeds.shape!=negative_prompt_embeds.shape:
raiseValueError(
"`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
f" {negative_prompt_embeds.shape}."
)

# 2. Encode caption
ifprompt_embedsisNoneandnegative_prompt_embedsisNone:
_, prompt_embeds_pooled, _, negative_prompt_embeds_pooled=self.encode_prompt(
prompt=prompt,
device=device,
batch_size=batch_size,
num_images_per_prompt=num_images_per_prompt,
do_classifier_free_guidance=self.do_classifier_free_guidance,
negative_prompt=negative_prompt,
prompt_embeds=prompt_embeds,
prompt_embeds_pooled=prompt_embeds_pooled,
negative_prompt_embeds=negative_prompt_embeds,
negative_prompt_embeds_pooled=negative_prompt_embeds_pooled,
)
# The pooled embeds from the prior are pooled again before being passed to the decoder
prompt_embeds_pooled= (
torch.cat([prompt_embeds_pooled, negative_prompt_embeds_pooled])
ifself.do_classifier_free_guidance

Problem:
The decoder only calls encode_prompt() when both prompt_embeds and negative_prompt_embeds are None. If a user supplies prompt_embeds/prompt_embeds_pooled and sets guidance_scale > 1, the pipeline never creates negative_prompt_embeds_pooled, then torch.cat() receives None.

Impact:
The documented precomputed-embedding path fails for classifier-free guidance instead of generating empty negative embeddings or raising a clear validation error.

Reproduction:

importtorchfromdiffusersimportDDPMWuerstchenScheduler, StableCascadeDecoderPipeline, StableCascadeUNetdecoder=StableCascadeUNet(
in_channels=4, out_channels=4, conditioning_dim=8, block_out_channels=(8,),
num_attention_heads=(-1,), down_num_layers_per_block=(1,), up_num_layers_per_block=(1,),
down_blocks_repeat_mappers=(1,), up_blocks_repeat_mappers=(1,),
block_types_per_layer=(("SDCascadeResBlock", "SDCascadeTimestepBlock"),),
clip_text_pooled_in_channels=8, effnet_in_channels=4,
)
pipe=StableCascadeDecoderPipeline(decoder, None, None, DDPMWuerstchenScheduler(), None, latent_dim_scale=1.0)
try:
pipe(
image_embeddings=torch.randn(1, 4, 1, 1),
prompt_embeds=torch.randn(1, 77, 8),
prompt_embeds_pooled=torch.randn(1, 1, 8),
guidance_scale=2.0,
num_inference_steps=1,
output_type="latent",
)
exceptExceptionase:
print(type(e).__name__, str(e).splitlines()[0])

Relevant precedent:
StableCascadePriorPipeline.check_inputs() already validates pooled prompt embeddings.

Suggested fix:

ifprompt_embedsisnotNoneandprompt_embeds_pooledisNone:
raiseValueError("If `prompt_embeds` are provided, `prompt_embeds_pooled` must also be provided.")
ifnegative_prompt_embedsisnotNoneandnegative_prompt_embeds_pooledisNone:
raiseValueError("If `negative_prompt_embeds` are provided, `negative_prompt_embeds_pooled` must also be provided.")
ifself.do_classifier_free_guidanceandnegative_prompt_embeds_pooledisNone:
_, _, _, negative_prompt_embeds_pooled=self.encode_prompt(...)

Issue 3: StableCascadeUNet.get_clip_embeddings() mishandles 2D pooled text embeddings

Affected code:

defget_clip_embeddings(self, clip_txt_pooled, clip_txt=None, clip_img=None):
iflen(clip_txt_pooled.shape) ==2:
clip_txt_pool=clip_txt_pooled.unsqueeze(1)
clip_txt_pool=self.clip_txt_pooled_mapper(clip_txt_pooled).view(
clip_txt_pooled.size(0), clip_txt_pooled.size(1) *self.config.clip_seq, -1
)
ifclip_txtisnotNoneandclip_imgisnotNone:
clip_txt=self.clip_txt_mapper(clip_txt)
iflen(clip_img.shape) ==2:
clip_img=clip_img.unsqueeze(1)
clip_img=self.clip_img_mapper(clip_img).view(
clip_img.size(0), clip_img.size(1) *self.config.clip_seq, -1
)
clip=torch.cat([clip_txt, clip_txt_pool, clip_img], dim=1)
else:
clip=clip_txt_pool
returnself.clip_norm(clip)

Problem:
The method detects 2D clip_txt_pooled and assigns clip_txt_pool = clip_txt_pooled.unsqueeze(1), but then ignores that normalized tensor and maps/views the original 2D tensor.

Impact:
Users passing normal CLIP pooled embeddings shaped (batch, dim) hit a runtime shape error, even though the method appears intended to support that shape.

Reproduction:

importtorchfromdiffusersimportStableCascadeUNetmodel=StableCascadeUNet(
conditioning_dim=8, block_out_channels=(8,), num_attention_heads=(-1,),
down_num_layers_per_block=(1,), up_num_layers_per_block=(1,),
down_blocks_repeat_mappers=(1,), up_blocks_repeat_mappers=(1,),
block_types_per_layer=(("SDCascadeResBlock", "SDCascadeTimestepBlock"),),
clip_text_pooled_in_channels=8,
)
try:
model.get_clip_embeddings(torch.randn(1, 8))
exceptExceptionase:
print(type(e).__name__, str(e).splitlines()[0])

Relevant precedent:
The pipelines produce pooled embeddings as (batch, 1, dim) internally; public direct model use should normalize 2D inputs to that same shape.

Suggested fix:

ifclip_txt_pooled.ndim==2:
clip_txt_pooled=clip_txt_pooled.unsqueeze(1)
clip_txt_pool=self.clip_txt_pooled_mapper(clip_txt_pooled).view(
clip_txt_pooled.size(0), clip_txt_pooled.size(1) *self.config.clip_seq, -1
)

Issue 4: sca/crp conditioning tensors crash batch forwards

Affected code:

timestep_ratio_embed=self.get_timestep_ratio_embedding(timestep_ratio)
forcinself.config.timestep_conditioning_type:
ifc=="sca":
cond=sca
elifc=="crp":
cond=crp
else:
cond=None
t_cond=condortorch.zeros_like(timestep_ratio)
timestep_ratio_embed=torch.cat([timestep_ratio_embed, self.get_timestep_ratio_embedding(t_cond)], dim=1)

Problem:
forward() uses t_cond = cond or torch.zeros_like(timestep_ratio). Tensor truthiness is invalid for batched tensors.

Impact:
The model exposes sca and crp conditioning parameters but cannot use them for normal batched inputs.

Reproduction:

importtorchfromdiffusersimportStableCascadeUNetmodel=StableCascadeUNet(
in_channels=4, out_channels=4, conditioning_dim=8, block_out_channels=(8,),
num_attention_heads=(-1,), down_num_layers_per_block=(1,), up_num_layers_per_block=(1,),
down_blocks_repeat_mappers=(1,), up_blocks_repeat_mappers=(1,),
block_types_per_layer=(("SDCascadeResBlock", "SDCascadeTimestepBlock"),),
clip_text_pooled_in_channels=8,
)
try:
model(
sample=torch.randn(2, 4, 8, 8),
timestep_ratio=torch.ones(2),
clip_text_pooled=torch.randn(2, 1, 8),
sca=torch.tensor([0.1, 0.2]),
crp=torch.tensor([0.3, 0.4]),
)
exceptExceptionase:
print(type(e).__name__, str(e).splitlines()[0])

Relevant precedent:
Standard tensor optional handling uses explicit is None checks.

Suggested fix:

t_cond=condifcondisnotNoneelsetorch.zeros_like(timestep_ratio)

Issue 5: User-supplied latents are moved to device but not cast to pipeline dtype

Affected code:

iflatentsisNone:
latents=randn_tensor(latent_shape, generator=generator, device=device, dtype=dtype)
else:
iflatents.shape!=latent_shape:
raiseValueError(f"Unexpected latents shape, got {latents.shape}, expected {latent_shape}")
latents=latents.to(device)
latents=latents*scheduler.init_noise_sigma

iflatentsisNone:
latents=randn_tensor(latents_shape, generator=generator, device=device, dtype=dtype)
else:
iflatents.shape!=latents_shape:
raiseValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}")
latents=latents.to(device)
latents=latents*scheduler.init_noise_sigma

Problem:
Both prepare_latents() methods cast generated latents to dtype, but pre-generated latents only call .to(device). Passing float32 latents to bf16/fp16 weights preserves float32 and can fail in convolutions.

Impact:
Reusing deterministic latents with half/bfloat16 Stable Cascade can error or silently change compute behavior. Related dtype-error issues exist (#7355/#7644), but those concerned arbitrary-resolution interpolation, not this supplied-latents path.

Reproduction:

importtorchfromdiffusersimportDDPMWuerstchenScheduler, StableCascadePriorPipeline, StableCascadeUNetprior=StableCascadeUNet(
conditioning_dim=8, block_out_channels=(8,), num_attention_heads=(-1,),
down_num_layers_per_block=(1,), up_num_layers_per_block=(1,),
down_blocks_repeat_mappers=(1,), up_blocks_repeat_mappers=(1,),
block_types_per_layer=(("SDCascadeResBlock", "SDCascadeTimestepBlock"),),
clip_text_pooled_in_channels=8, clip_image_in_channels=8,
)
pipe=StableCascadePriorPipeline(None, None, prior, DDPMWuerstchenScheduler())
latents=torch.randn(1, prior.config.in_channels, 1, 1, dtype=torch.float32)
out=pipe.prepare_latents(1, 42, 42, 1, torch.bfloat16, torch.device("cpu"), None, latents, pipe.scheduler)
print(out.dtype) # torch.float32, expected torch.bfloat16

Relevant precedent:

iflatentsisnotNone:
latent_image_ids=self._prepare_latent_image_ids(batch_size, height//2, width//2, device, dtype)
returnlatents.to(device=device, dtype=dtype), latent_image_ids

Suggested fix:

latents=latents.to(device=device, dtype=dtype)

Issue 6: StableCascadePriorPipeline.timesteps is public but ignored

Affected code:


# 4. Prepare and set timesteps
self.scheduler.set_timesteps(num_inference_steps, device=device)
timesteps=self.scheduler.timesteps

prior_num_inference_steps (`int | dict[float, int]`, *optional*, defaults to 60):
The number of prior denoising steps. More denoising steps usually lead to a higher quality image at the
expense of slower inference. For more specific timestep spacing, you can pass customized
`prior_timesteps`
num_inference_steps (`int`, *optional*, defaults to 12):
The number of decoder denoising steps. More denoising steps usually lead to a higher quality image at
the expense of slower inference. For more specific timestep spacing, you can pass customized
`timesteps`

Problem:
The prior pipeline accepts timesteps, and the combined docstring advertises customized prior_timesteps/timesteps, but the prior always calls self.scheduler.set_timesteps(num_inference_steps, device=device).

Impact:
Users cannot control the prior schedule despite the public API suggesting they can.

Reproduction:

importinspectfromdiffusersimportStableCascadePriorPipelinesrc=inspect.getsource(StableCascadePriorPipeline.__call__)
print("timesteps parameter exists:", "timesteps: list[float]"insrc)
print("passed to scheduler:", "timesteps=timesteps"insrc)

Relevant precedent:

defset_timesteps(
self,
num_inference_steps: int=None,
timesteps: list[int] |None=None,
device: str|torch.device=None,
):
"""
Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference.
Args:
num_inference_steps (`dict[float, int]`):
the number of diffusion steps used when generating samples with a pre-trained model. If passed, then
`timesteps` must be `None`.
device (`str` or `torch.device`, optional):
the device to which the timesteps are moved to. {2 / 3: 20, 0.0: 10}
"""
iftimestepsisNone:
timesteps=torch.linspace(1.0, 0.0, num_inference_steps+1, device=device)
ifnotisinstance(timesteps, torch.Tensor):
timesteps=torch.Tensor(timesteps).to(device)

Suggested fix:

self.scheduler.set_timesteps(num_inference_steps, timesteps=timesteps, device=device)

Also either add matching decoder/combined args or remove the combined docstring claims.

Issue 7: Combined xFormers enable only reaches the decoder

Affected code:

defenable_xformers_memory_efficient_attention(self, attention_op: Callable|None=None):
self.decoder_pipe.enable_xformers_memory_efficient_attention(attention_op)

Problem:
StableCascadeCombinedPipeline.enable_xformers_memory_efficient_attention() delegates only to decoder_pipe, leaving the prior unchanged.

Impact:
Users enabling xFormers on the combined pipeline do not get memory-efficient attention for the prior stage.

Reproduction:

fromunittest.mockimportMockfromdiffusersimportStableCascadeCombinedPipelinepipe=StableCascadeCombinedPipeline.__new__(StableCascadeCombinedPipeline)
pipe.prior_pipe=Mock()
pipe.decoder_pipe=Mock()
pipe.enable_xformers_memory_efficient_attention()
print(pipe.prior_pipe.enable_xformers_memory_efficient_attention.call_count) # 0print(pipe.decoder_pipe.enable_xformers_memory_efficient_attention.call_count) # 1

Relevant precedent:
The prior and decoder both contain StableCascadeUNet attention modules and both support recursive xFormers enabling through ModelMixin.

Suggested fix:

defenable_xformers_memory_efficient_attention(self, attention_op: Callable|None=None):
self.prior_pipe.enable_xformers_memory_efficient_attention(attention_op)
self.decoder_pipe.enable_xformers_memory_efficient_attention(attention_op)

Issue 8: Test coverage misses combined slow coverage and dedicated model fast tests

Affected code:

classStableCascadeCombinedPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
pipeline_class=StableCascadeCombinedPipeline
params= ["prompt"]
batch_params= ["prompt", "negative_prompt"]
required_optional_params= [
"generator",
"height",
"width",
"latents",
"prior_guidance_scale",
"decoder_guidance_scale",
"negative_prompt",
"num_inference_steps",
"return_dict",
"prior_num_inference_steps",
"output_type",
]
test_xformers_attention=True

@unittest.skip(reason="no callback test for combined pipeline")
deftest_callback_inputs(self):
super().test_callback_inputs()

logger=logging.get_logger(__name__)
enable_full_determinism()
@slow
@require_torch_accelerator
classStableCascadeUNetSingleFileTest:

Problem:
Prior and decoder have slow integration tests, but the combined pipeline has no @slow test. StableCascadeUNet also lacks a dedicated fast tests/models/... test; it is mostly covered indirectly and by slow single-file config checks.

Impact:
The issues above are not covered by the current fast/slow matrix, especially direct model input-shape behavior and combined pipeline delegation behavior.

Reproduction:

frompathlibimportPathprint("model fast tests:", list(Path("tests/models").glob("**/*stable*cascade*.py")))
forpathinsorted(Path("tests/pipelines/stable_cascade").glob("test_*.py")):
print(path.name, "@slow"inpath.read_text())

Relevant precedent:
Most active model families have dedicated ModelTesterMixin coverage for forward, save/load, dtype, gradient checkpointing, and attention/offload behavior.

Suggested fix:
Add:

# tests/models/unets/test_models_unet_stable_cascade.pyclassStableCascadeUNetTests(ModelTesterMixin, unittest.TestCase):
model_class=StableCascadeUNet

and a @slow combined pipeline integration test that loads stabilityai/stable-cascade and exercises the full prior+decoder path.

Issue 9: Pipeline docstring examples are not runnable

Affected code:

EXAMPLE_DOC_STRING="""
Examples:
```py
>>> import torch
>>> from diffusers import StableCascadePriorPipeline
>>> prior_pipe = StableCascadePriorPipeline.from_pretrained(
... "stabilityai/stable-cascade-prior", torch_dtype=torch.bfloat16
... ).to("cuda")
>>> prompt = "an image of a shiba inu, donning a spacesuit and helmet"
>>> prior_output = pipe(prompt)
```

EXAMPLE_DOC_STRING="""
Examples:
```py
>>> import torch
>>> from diffusers import StableCascadePriorPipeline, StableCascadeDecoderPipeline
>>> prior_pipe = StableCascadePriorPipeline.from_pretrained(
... "stabilityai/stable-cascade-prior", torch_dtype=torch.bfloat16
... ).to("cuda")
>>> gen_pipe = StableCascadeDecoderPipeline.from_pretrain(
... "stabilityai/stable-cascade", torch_dtype=torch.float16
... ).to("cuda")
>>> prompt = "an image of a shiba inu, donning a spacesuit and helmet"
>>> prior_output = pipe(prompt)
>>> images = gen_pipe(prior_output.image_embeddings, prompt=prompt)
```

Problem:
The prior example assigns prior_pipe but calls pipe(prompt). The decoder example calls StableCascadeDecoderPipeline.from_pretrain instead of from_pretrained and also calls pipe(prompt).

Impact:
Autodoc examples copied by users fail immediately.

Reproduction:

frompathlibimportPathforpathin [
Path("src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade_prior.py"),
Path("src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py"),
]:
text=path.read_text()
print(path.name, "from_pretrain("intext, "prior_output = pipe(prompt)"intext)

Relevant precedent:
The docs page docs/source/en/api/pipelines/stable_cascade.md uses the correct top-level imports and from_pretrained.

Suggested fix:

prior_output=prior_pipe(prompt)
gen_pipe=StableCascadeDecoderPipeline.from_pretrained(...)

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions