Wan Animate Pipeline#367
Conversation
csgoogle
force-pushed
the
sagarchapara/wananimate-pipeline
branch
2 times, most recently
from
e281524 to
349d080
Compare
|
Please resolve conflicts and enable support for diagnostics and profiling as in this PR |
|
🤖 Hi @Perseus14, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
📋 Review Summary
The Pull Request introduces the Wan Animate pipeline, which includes the transformer model architecture, inference entry point, and necessary utilities. The implementation is comprehensive and follows the established patterns in the repository, including support for segment-based inference and parity with Diffusers.
🔍 General Feedback
- Performance Optimization: The current implementation of the transformer re-encodes the face video frames during every denoising step. Since the face video is static throughout the inference process, this encoding can be pre-computed once per segment to significantly reduce redundant computation and speed up generation.
- Compilation Efficiency: The generation script performs two full inference passes. For high-resolution video generation, this double work is expensive. Consider reducing the number of steps in the first (compile) pass.
- Robustness: Added checks for optional inputs in the transformer to prevent potential runtime errors when
face_pixel_valuesis not provided. - Code Quality: The reuse of the Wan attention operator and the integration with the existing configuration system is well-done. The use of
nnx.scanfor transformer blocks ensures memory efficiency during inference.
csgoogle
force-pushed
the
sagarchapara/wananimate-pipeline
branch
3 times, most recently
from
ef88d04 to
f6b4c22
Compare
|
🤖 Hi @Perseus14, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
|
🤖 Hi @Perseus14, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
|
🤖 I'm sorry @Perseus14, but I was unable to process your request. Please see the logs for more details. |
|
🤖 Hi @Perseus14, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
There was a problem hiding this comment.
This PR successfully implements the Wan Animate pipeline in MaxDiffusion, including the 3D transformer architecture, motion and face encoders, and the inference pipeline with temporal tiling. The implementation is robust, follows established patterns for sharding and JIT optimization, and is backed by comprehensive parity tests.
🔍 General Feedback
- Efficiency: The use of pre-computed motion vectors once per segment and the implementation of
scan_layersshow good attention to performance on TPU/GPU hardware. - Testing: The parity tests are exhaustive and provide high confidence in the implementation's correctness relative to the reference Diffusers implementation.
- Modularity: Reusing the Wan attention operator for face conditioning is a clean approach.
There was a problem hiding this comment.
Additional technical suggestions for the Wan Animate transformer implementation.
🔍 General Feedback
- Optimization: Minor optimizations in
MotionConv2dto avoid repeated array conversions. - Robustness: Defensive check in
WanAnimateFaceBlockCrossAttentionfor sequence length consistency.
|
🤖 I'm sorry @Perseus14, but I was unable to process your request. Please see the logs for more details. |
|
🤖 Hi @Perseus14, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
|
🤖 I'm sorry @Perseus14, but I was unable to process your request. Please see the logs for more details. |
|
🤖 Hi @Perseus14, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
|
🤖 I'm sorry @Perseus14, but I was unable to process your request. Please see the logs for more details. |
|
🤖 Hi @Perseus14, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
There was a problem hiding this comment.
The Pull Request introduces the Wan Animate pipeline and its associated 3D transformer model to the maxdiffusion library. The implementation is comprehensive, including checkpointing support, sharding-aware attention, and extensive parity tests against the reference implementation. The code follows the established architectural patterns of the project and integrates well with existing Wan and NNX utilities.
🔍 General Feedback
- Parity Testing: Excellent inclusion of detailed parity tests (
wan_animate_module_parity_test.py) which ensures the JAX/Flax implementation matches the reference torch model. - NNX Migration: The use of
flax.nnxfor the new models is consistent with the project's direction. - Robustness: Some minor improvements suggested for robustness (e.g., division by zero checks) and consistency in parameter access within
nnxmodules. - Documentation: The docstrings are informative and follow the project's style.
csgoogle
force-pushed
the
sagarchapara/wananimate-pipeline
branch
2 times, most recently
from
b2ea208 to
a2d4356
Compare
- Add WanAnimateTransformer3DModel with motion encoder, face encoder, and face adapter cross-attention blocks - Add WanAnimatePipeline supporting animate and replace modes with multi-segment temporal stitching - Add generate_wan_animate.py inference entrypoint - Add base_wan_animate.yml config for 720p inference - Pre-compute face motion vectors once per segment instead of every denoising step for faster inference - Simplify face block cross-attention forward pass: replace einops with jnp.reshape, remove redundant sharding constraints - Add parity tests for animate modules and diffusers comparison
csgoogle
force-pushed
the
sagarchapara/wananimate-pipeline
branch
from
a2d4356 to
eddfd4d
Compare
|
LGTM! Test case failed is unrelated to this change. |
2 participants
Wan Animate Pipeline
This CL publishes add the Wan Animate pipepline.
Links
Performance
Configuration