From 890803ca066fd159ce61072d9f75b032b1fd59a1 Mon Sep 17 00:00:00 2001 From: Kristoffer Date: Wed, 17 Jun 2026 09:48:22 +0200 Subject: [PATCH] [Wan2.2] Fix "uninitialized nn.Module of type RotaryEmbeddingWan" graph-break (#4053) Signed-off-by: Kristoffer Torp Co-authored-by: Hongsheng Liu --- vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py index cda01e083..3f4d3ec9a 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py @@ -406,6 +406,8 @@ class WanSelfAttention(nn.Module): ) self.dropout = nn.Dropout(dropout) + self.rotary_embedding = RotaryEmbeddingWan(is_neox_style=False, half_head_dim=True) + # Unified attention layer self.attn = Attention( num_heads=self.num_heads, @@ -441,7 +443,6 @@ class WanSelfAttention(nn.Module): # Apply rotary embeddings if rotary_emb is not None: - self.rotary_embedding = RotaryEmbeddingWan(is_neox_style=False, half_head_dim=True) freqs_cos, freqs_sin = rotary_emb query = self.rotary_embedding(query, freqs_cos, freqs_sin) key = self.rotary_embedding(key, freqs_cos, freqs_sin)