[Wan2.2] Fix "uninitialized nn.Module of type RotaryEmbeddingWan" graph-break (#4053)

Signed-off-by: Kristoffer Torp <kristoffer.torp@amd.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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)