[Model] Add native MAGI-2 Preview diffusion support (#5918)

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

tests/diffusion/models/magi2/test_audio_decoder.py

@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.magi2.audio_decoder import (
+    convert_stable_audio_decoder_key,
+    convert_stable_audio_decoder_state_dict,
+)
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+@pytest.mark.parametrize(
+    ("source", "target"),
+    [
+        ("pretransform.model.decoder.layers.0.weight_v", "conv1.weight_v"),
+        ("pretransform.model.decoder.layers.1.layers.0.alpha", "block.0.snake1.alpha"),
+        ("pretransform.model.decoder.layers.3.layers.1.weight_g", "block.2.conv_t1.weight_g"),
+        (
+            "pretransform.model.decoder.layers.5.layers.4.layers.2.beta",
+            "block.4.res_unit3.snake2.beta",
+        ),
+        (
+            "pretransform.model.decoder.layers.2.layers.2.layers.3.weight_v",
+            "block.1.res_unit1.conv2.weight_v",
+        ),
+        ("pretransform.model.decoder.layers.6.alpha", "snake1.alpha"),
+        ("pretransform.model.decoder.layers.7.weight_g", "conv2.weight_g"),
+    ],
+)
+def test_stable_audio_decoder_key_conversion(source: str, target: str):
+    assert convert_stable_audio_decoder_key(source) == target
+
+
+def test_stable_audio_decoder_conversion_filters_and_reshapes_snake():
+    alpha = torch.arange(4, dtype=torch.float32)
+    weight = torch.ones(2, 4, 7)
+    state = {
+        "pretransform.model.decoder.layers.6.alpha": alpha,
+        "pretransform.model.decoder.layers.7.weight_v": weight,
+        "model.model.transformer.layers.0.weight": torch.empty(1),
+    }
+    converted = convert_stable_audio_decoder_state_dict(state)
+    assert set(converted) == {"snake1.alpha", "conv2.weight_v"}
+    assert converted["snake1.alpha"].shape == (1, 4, 1)
+    torch.testing.assert_close(converted["snake1.alpha"].flatten(), alpha)
+    assert converted["conv2.weight_v"] is weight
+
+
+def test_unknown_decoder_key_fails_closed():
+    with pytest.raises(ValueError, match="Unsupported Stable Audio decoder key"):
+        convert_stable_audio_decoder_key("pretransform.model.decoder.future.weight")