From c484e2db45718da6a5f5fa93049da6690952d64f Mon Sep 17 00:00:00 2001 From: Fabrice TIERCELIN Date: Tue, 2 Dec 2025 20:03:29 +0100 Subject: [PATCH] Optimize time & memory for Wan --- src/diffusers/models/transformers/transformer_wan.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/diffusers/models/transformers/transformer_wan.py b/src/diffusers/models/transformers/transformer_wan.py index f7693ec5d3ac..224b2c156927 100644 --- a/src/diffusers/models/transformers/transformer_wan.py +++ b/src/diffusers/models/transformers/transformer_wan.py @@ -106,7 +106,8 @@ def apply_rotary_emb( freqs_cos: torch.Tensor, freqs_sin: torch.Tensor, ): - x1, x2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1) + x1 = hidden_states[..., 0::2] + x2 = hidden_states[..., 1::2] cos = freqs_cos[..., 0::2] sin = freqs_sin[..., 1::2] out = torch.empty_like(hidden_states)