Skip to content

Fix wan 2.1 i2v context parallel #12909

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sayakpaul merged 7 commits into from
Jan 6, 2026
Merged

Fix wan 2.1 i2v context parallel #12909

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
edc30c7
fix wan 2.1 i2v context parallel
DefTruth Jan 4, 2026
d4db358
fix wan 2.1 i2v context parallel
DefTruth Jan 4, 2026
bda855e
fix wan 2.1 i2v context parallel
DefTruth Jan 4, 2026
4535e3e
Merge branch 'main' into fix-wan-2.1-i2v-cp
DefTruth Jan 5, 2026
fcc688c
Merge branch 'main' into fix-wan-2.1-i2v-cp
DefTruth Jan 6, 2026
8875199
format
DefTruth Jan 6, 2026
55e472e
Merge branch 'fix-wan-2.1-i2v-cp' of https://github.com/xlite-dev/dif…
DefTruth Jan 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 9 additions & 5 deletions src/diffusers/models/transformers/transformer_wan.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -134,7 +134,8 @@ def apply_rotary_emb(
dropout_p=0.0,
is_causal=False,
backend=self._attention_backend,
parallel_config=self._parallel_config,
# Reference: https://github.com/huggingface/diffusers/pull/12909
parallel_config=None,
)
hidden_states_img = hidden_states_img.flatten(2, 3)
hidden_states_img = hidden_states_img.type_as(query)
Expand All @@ -147,7 +148,8 @@ def apply_rotary_emb(
dropout_p=0.0,
is_causal=False,
backend=self._attention_backend,
parallel_config=self._parallel_config,
# Reference: https://github.com/huggingface/diffusers/pull/12909
parallel_config=(self._parallel_config if encoder_hidden_states is None else None),
)
hidden_states = hidden_states.flatten(2, 3)
hidden_states = hidden_states.type_as(query)
Expand Down Expand Up @@ -552,9 +554,11 @@ class WanTransformer3DModel(
"blocks.0": {
"hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False),
},
"blocks.*": {
"encoder_hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False),
},
# Reference: https://github.com/huggingface/diffusers/pull/12909
Copy link
Member

@sayakpaul sayakpaul Jan 5, 2026 •
edited
Loading

Choose a reason for hiding this comment

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

Is this specific to I2V only? If so, then this change is probably a little to intrusive no?

Copy link
Contributor Author

@DefTruth DefTruth Jan 5, 2026 •
edited
Loading

Choose a reason for hiding this comment

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

@sayakpaul This is theoretically applicable to all wan series models and offers better performance. I conducted tests on the wan 2.1/2.2 t2v/i2v models, and the results were all correct. You can perform quick verification through examples in cache-dit. This patch has already been used in cache-dit to fix some precision issues. vipshop/cache-dit#639

pip3 install torch==2.9.1 transformers accelerate torchao bitsandbytes torchvision 
pip3 install opencv-python-headless einops imageio-ffmpeg ftfy 
pip3 install git+https://github.com/huggingface/diffusers.git # latest 
pip3 install git+https://github.com/vipshop/cache-dit.git # latest

git clone https://github.com/vipshop/cache-dit.git && cd cache-dit/examples

# use  `--cpu-offload` and `--parallel-text-encoder` for low VRAM device, e.g, < 48GiB
torchrun --nproc_per_node=4 generate.py wan2.1_t2v --parallel ulysses --parallel-text-encoder
torchrun --nproc_per_node=4 generate.py wan2.2_t2v --parallel ulysses --parallel-text-encoder --cpu-offload
torchrun --nproc_per_node=2 generate.py wan2.1_i2v --parallel ulysses --parallel-text-encoder --steps 16 --frames 21 --vae-tiling
torchrun --nproc_per_node=2 generate.py wan2.2_i2v --parallel ulysses --parallel-text-encoder --cpu-offload --steps 16 --frames 21 --vae-tiling

# We need to disable the splitting of encoder_hidden_states because the image_encoder
# (Wan 2.1 I2V) consistently generates 257 tokens for image_embed. This causes the shape
# of encoder_hidden_states—whose token count is always 769 (512 + 257) after concatenation
# —to be indivisible by the number of devices in the CP.
"proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
"": {
"timestep": ContextParallelInput(split_dim=1, expected_dims=2, split_output=False),
Expand Down
6 changes: 4 additions & 2 deletions src/diffusers/models/transformers/transformer_wan_animate.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -609,7 +609,8 @@ def apply_rotary_emb(
dropout_p=0.0,
is_causal=False,
backend=self._attention_backend,
parallel_config=self._parallel_config,
# Reference: https://github.com/huggingface/diffusers/pull/12909
parallel_config=None,
)
hidden_states_img = hidden_states_img.flatten(2, 3)
hidden_states_img = hidden_states_img.type_as(query)
Expand All @@ -622,7 +623,8 @@ def apply_rotary_emb(
dropout_p=0.0,
is_causal=False,
backend=self._attention_backend,
parallel_config=self._parallel_config,
# Reference: https://github.com/huggingface/diffusers/pull/12909
parallel_config=(self._parallel_config if encoder_hidden_states is None else None),
)
hidden_states = hidden_states.flatten(2, 3)
hidden_states = hidden_states.type_as(query)
Expand Down
Loading