Fix batch decoding for Wan-Video-VAE

[krahets]

The Wan-VAE decoding for batch sizes > 1 was broken after 830b1b7. We may roll it back.

values[
    :,
    :,
    :,
    target_h:target_h + hidden_states_batch.shape[3],
    target_w:target_w + hidden_states_batch.shape[4],
] += hidden_states_batch * mask
*** RuntimeError: output with shape [1, 3, 5, 416, 240] doesn't match the broadcast shape [2, 3, 5, 416, 240]

in which

values [1, 3, 5, 416, 240]
hidden_states_batch [2, 3, 5, 416, 240]

[gemini-code-assist]

Summary of Changes

Hello @krahets, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request resolves a critical bug in the Wan-Video-VAE model that previously caused decoding failures for video batches larger than one. The core of the fix involves modifying the decoding process to handle each video's hidden state individually, ensuring proper tensor shape alignment. This change restores the model's ability to efficiently process and decode multiple videos simultaneously, enhancing its robustness and usability for batch operations.

Highlights

• Fix Batch Decoding Error: This pull request addresses a critical RuntimeError that prevented the Wan-Video-VAE model from correctly decoding video batches when the batch size exceeded one. The error stemmed from a mismatch in tensor shapes during an addition operation.
• Iterative Decoding Implementation: The decode method has been refactored to process each hidden state sequentially within a loop. This ensures that single_decode or tiled_decode receive individual, correctly shaped tensors, resolving the dimension mismatch issue.
• CPU Transfer for Hidden States: Hidden states are now explicitly moved to the CPU before being processed in the decode method, which might be a necessary step for the sequential processing or to avoid device-related issues during the fix.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request addresses a critical bug in WanVideoVAE.decode that caused a RuntimeError with batch sizes greater than one. The fix correctly serializes the decoding process by iterating over the batch, ensuring that tiled_decode and single_decode receive inputs with a batch size of one. While this resolves the issue, I've provided a suggestion to refactor the implementation for improved readability and slightly better memory efficiency by avoiding the creation of an intermediate list and clarifying variable names.

[gemini-code-assist]

The implementation correctly fixes the batch decoding issue by iterating through the batch. However, the code can be made more readable and slightly more memory-efficient.

1. The current implementation reassigns hidden_states and then uses hidden_state as a loop variable which is also reassigned inside the loop. This can be confusing. Using new variable names for clarity is better.
2. It's slightly more memory-efficient to move each hidden_state to CPU inside the loop, rather than creating a new list of all hidden states on CPU at once.

Here is a suggested refactoring that addresses these points.

        videos = []
        for hs in hidden_states:
            hs_batch = hs.to("cpu").unsqueeze(0)
            if tiled:
                video = self.tiled_decode(hs_batch, device, tile_size, tile_stride)
            else:
                video = self.single_decode(hs_batch, device)
            videos.append(video.squeeze(0))
        return torch.stack(videos)