56 lines
2.1 KiB
Python
56 lines
2.1 KiB
Python
|
# 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")
|