Fix Stable Cascade checkpointing crash by using native diffusers checkpointing - #1539
Merged
Merged
Conversation
…kpointing enable_checkpointing_for_stable_cascade_blocks searched for a homogeneous nn.ModuleList of a single block type, but StableCascadeUNet's down_blocks/ up_blocks interleave SDCascadeResBlock, SDCascadeTimestepBlock, and SDCascadeAttnBlock in the same list, so the homogeneity assert always failed when checkpointing was enabled. diffusers' StableCascadeUNet already supports native gradient checkpointing with correct per-block-type dispatch, so use model.prior_prior.enable_gradient_checkpointing() instead, matching how Wuerstchen v2 is already handled. Fixes Nerogar#1529 Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
enable_checkpointing_for_stable_cascade_blockssearched for a homogeneousnn.ModuleListof a single block type (SDCascadeResBlock/SDCascadeAttnBlock/SDCascadeTimestepBlock), butStableCascadeUNet'sdown_blocks/up_blocksinterleave all three block types within the same list by design. The matcher's
assert all(isinstance(m, t) for m in child_module)therefore always failed onceit found the first matching list, crashing on startup whenever checkpointing was
enabled for the prior.
This workaround predates diffusers' own gradient-checkpointing support for Stable
Cascade (it was written when the model still relied on the older Wuerstchen pipeline
code without
enable_gradient_checkpointing()). diffusers'StableCascadeUNetnowsupports native gradient checkpointing with correct per-block-type dispatch
internally, so this switches to
model.prior_prior.enable_gradient_checkpointing(),the same call already used for Wuerstchen v2, and removes the now-unused custom
matcher function.
Fixes #1529
Drafted by Claude