Adding tea cache wan2.2s2v by aviveise · Pull Request #1017 · modelscope/DiffSynth-Studio · GitHub
Skip to content

Adding tea cache wan2.2s2v - #1017

Open
aviveise wants to merge 23 commits into
modelscope:mainfrom
Smiti-AI:adding_tea_cache_wan2.2s2v
Open

aviveise wants to merge 23 commits into
modelscope:mainfrom
Smiti-AI:adding_tea_cache_wan2.2s2v

Conversation

@aviveise

Copy link
Copy Markdown

No description provided.

@gemini-code-assist

Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces a 'tea cache' mechanism for wan2.2s2v, refactors variable names for better clarity (e.g., x to latents/hidden_states), and adds support for offline preprocessing in the WanVideoPipeline. While the refactoring improves readability, I've identified several critical bugs. These include using undefined variables after renaming and passing incorrect sequence lengths in checkpointed functions when using unified sequence parallelism, which will likely cause runtime errors or incorrect outputs. Additionally, there are a couple of medium-severity issues: some dead code that should be removed and a debug print statement left in the code. Addressing these issues is crucial for the correctness and maintainability of the new features.


if self.has_image_input:
x = torch.cat([x, y], dim=1) # (b, c_x + c_y, f, h, w)
hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1) # (b, c_x + c_y, f, h, w)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

There seems to be a copy-paste error here. encoder_hidden_states is being concatenated with hidden_states, but based on the original code and the comment, it should be y (the reference image latents). Concatenating text embeddings with image latents along the channel dimension is likely incorrect.

Suggested change
hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1) # (b, c_x + c_y, f, h, w)
hidden_states = torch.cat([hidden_states, y], dim=1) # (b, c_x + c_y, f, h, w)

Comment on lines +393 to 403
hidden_states = torch.utils.checkpoint.checkpoint(
create_custom_forward(block),
x, context, t_mod, freqs,
use_reentrant=False,
)
else:
x = torch.utils.checkpoint.checkpoint(
hidden_states = torch.utils.checkpoint.checkpoint(
create_custom_forward(block),
x, context, t_mod, freqs,
use_reentrant=False,
)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

The variable x is used here in the arguments to torch.utils.checkpoint.checkpoint, but it is not defined in the scope of this function after the refactoring. The variable was renamed to hidden_states. You should use hidden_states instead of x.

Suggested change
hidden_states = torch.utils.checkpoint.checkpoint(
create_custom_forward(block),
x, context, t_mod, freqs,
use_reentrant=False,
)
else:
x = torch.utils.checkpoint.checkpoint(
hidden_states = torch.utils.checkpoint.checkpoint(
create_custom_forward(block),
x, context, t_mod, freqs,
use_reentrant=False,
)
hidden_states = torch.utils.checkpoint.checkpoint(
create_custom_forward(block),
hidden_states, context, t_mod, freqs,
use_reentrant=False,
)
else:
hidden_states = torch.utils.checkpoint.checkpoint(
create_custom_forward(block),
hidden_states, context, t_mod, freqs,
use_reentrant=False,
)

Comment on lines +580 to 610
if use_gradient_checkpointing_offload:
with torch.autograd.graph.save_on_cpu():
hidden_states = torch.utils.checkpoint.checkpoint(
create_custom_forward(block),
hidden_states,
encoder_hidden_states,
t_mod,
seq_len_x,
pre_compute_freqs[0],
use_reentrant=False,
)
hidden_states = torch.utils.checkpoint.checkpoint(
create_custom_forward(lambda x: self.after_transformer_block(block_id, x, audio_emb_global, merged_audio_emb, seq_len_x)),
hidden_states,
use_reentrant=False,
)
elif use_gradient_checkpointing:
hidden_states = torch.utils.checkpoint.checkpoint(
create_custom_forward(block),
x,
context,
hidden_states,
encoder_hidden_states,
t_mod,
seq_len_x,
pre_compute_freqs[0],
use_reentrant=False,
)
x = torch.utils.checkpoint.checkpoint(
hidden_states = torch.utils.checkpoint.checkpoint(
create_custom_forward(lambda x: self.after_transformer_block(block_id, x, audio_emb_global, merged_audio_emb, seq_len_x)),
x,
hidden_states,
use_reentrant=False,
)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

When using gradient checkpointing with unified sequence parallel (USP), you are passing seq_len_x to after_transformer_block. However, seq_len_x is the local sequence length for the current parallel rank. The after_transformer_block function performs an all_gather and requires the global sequence length (seq_len_x_global) to correctly process the gathered tensor. The non-checkpointed path correctly uses seq_len_x_global. This should be corrected for the checkpointed paths as well to avoid errors or incorrect behavior. Also, use_unified_sequence_parallel should be passed to after_transformer_block.

                if use_gradient_checkpointing_offload:
                    with torch.autograd.graph.save_on_cpu():
                        hidden_states = torch.utils.checkpoint.checkpoint(
                            create_custom_forward(block),
                            hidden_states,
                            encoder_hidden_states,
                            t_mod,
                            seq_len_x,
                            pre_compute_freqs[0],
                            use_reentrant=False,
                        )
                        hidden_states = torch.utils.checkpoint.checkpoint(
                            create_custom_forward(lambda x: self.after_transformer_block(block_id, x, audio_emb_global, merged_audio_emb, seq_len_x_global, use_unified_sequence_parallel=use_unified_sequence_parallel)),
                            hidden_states,
                            use_reentrant=False,
                        )
                elif use_gradient_checkpointing:
                    hidden_states = torch.utils.checkpoint.checkpoint(
                        create_custom_forward(block),
                        hidden_states,
                        encoder_hidden_states,
                        t_mod,
                        seq_len_x,
                        pre_compute_freqs[0],
                        use_reentrant=False,
                    )
                    hidden_states = torch.utils.checkpoint.checkpoint(
                        create_custom_forward(lambda x: self.after_transformer_block(block_id, x, audio_emb_global, merged_audio_emb, seq_len_x_global, use_unified_sequence_parallel=use_unified_sequence_parallel)),
                        hidden_states,
                        use_reentrant=False,
                    )

def process(self, pipe: WanVideoPipeline):
if hasattr(pipe, "use_unified_sequence_parallel"):
if pipe.use_unified_sequence_parallel:
print("use_unified_sequence_parallel true")

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

This print statement appears to be for debugging and should be removed from production code.

@Artiprocher

Copy link
Copy Markdown
Collaborator

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants