Fix Qwen2.5-Omni Token2Wav DiT rotary embedding layout (interleaved cos/sin)

[HenryVarro666]

What does this PR do?

Fixes #47328.

The Token2Wav DiT applies an interleaved rotate — rotate_half_codec, which pairs adjacent channels (2i, 2i+1) — but since #39847 (🚨 [v5] Refactor RoPE for layer types) Qwen2_5OmniDiTRotaryEmbedding inherited LlamaRotaryEmbedding.forward, which builds cos/sin with the half-split layout torch.cat((freqs, freqs)). A half-split cos/sin combined with an interleaved rotate is not a valid rotation: RoPE's defining property — ⟨R_m·q, R_n·k⟩ depending only on the relative offset m − n — is destroyed. It is applied to head 0 of every DiT attention layer (query[:, :1], key[:, :1] = apply_rotary_pos_emb(...)), silently degrading the predicted mel-spectrogram and thus the synthesized audio for return_audio=True.

This is a regression, not the checkpoint design. Before #39847 (e.g. v4.57.0), Qwen2_5OmniDiTRotaryEmbedding.forward built the rotary embedding as

freqs = torch.stack((freqs, freqs), dim=-1)
freqs = freqs.reshape(*freqs.shape[:-2], -1)   # == repeat_interleave(freqs, 2): interleaved

i.e. the interleaved layout that matches rotate_half_codec. #39847 made the class inherit LlamaRotaryEmbedding and silently replaced this with the half-split torch.cat((freqs, freqs)), while leaving rotate_half_codec interleaved. The original add PR #36752 was integration-tested against Qwen/Qwen2.5-Omni-7B with the interleaved layout, so the released weights expect it. The code also carries the note # Modified from Llama with a different rotate function, will fixed in next release on apply_rotary_pos_emb, confirming the rotate is intentionally non-standard (interleaved).

This PR restores the original interleaved layout via repeat_interleave(freqs, 2). Why this direction, rather than switching the rotate to Llama's half-split? Both would be valid rotations in isolation, but rotate_half_codec is exactly the (2i, 2i+1) pairing the released weights were trained with, so it is left untouched — which makes interleaved cos/sin the only checkpoint-preserving choice. Rewriting the rotate to half-split would be a different rotation convention the checkpoint was never trained for. The text model's mrope (Qwen2_5OmniRotaryEmbedding) legitimately keeps cat with the half-split rotate_half and is untouched; qwen3_omni_moe has its own rotary and is unaffected.

Reproduction / verification (CPU, no weights)

Driving the real Qwen2_5OmniDiTRotaryEmbedding.forward + apply_rotary_pos_emb, and measuring the max drift of ⟨R_m·q, R_{m+Δ}·k⟩ across absolute position m (≈ 0 ⟺ valid RoPE):

	cos/sin layout	translation-invariance drift
main	half-split (cat)	10.74 ❌
this PR	interleaved (repeat_interleave)	1.0e-05 ✅

Tests

Adds Qwen2_5OmniDiTRotaryEmbeddingTest (CPU, weight-free): it asserts the interleaved cos/sin layout and RoPE translation invariance (checked across several offsets) for the DiT rotary. Both tests fail on main and pass with this fix.

Before submitting

• Was this discussed/approved via a GitHub issue? Qwen2.5-Omni Token2Wav DiT uses mismatched RoPE layout (half-split cos/sin with interleaved rotate) — regression from #39847 #47328
• Did you write any new necessary tests?

Who can review?

@eustlb @ebezzam

Thanks to @pepepeboom for independently reproducing the root cause and confirming the regression on main.

Disclosure: the root-cause analysis, regression bisect, reproduction, and tests in this PR were prepared with AI assistance and independently verified by me.

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: qwen2_5_omni

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 29661591086:2
Result: success | Jobs: 15 | Tests: 162,354 | Failures: 0 | Duration: 13h 56m