Repository navigation
Conversation
There was a problem hiding this comment.
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) |
There was a problem hiding this comment.
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.
| 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) |
| 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, | ||
| ) |
There was a problem hiding this comment.
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.
| 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, | |
| ) |
| 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, | ||
| ) |
There was a problem hiding this comment.
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") |

No description provided.