vllm-omni/tests/diffusion/models/kandinsky6/test_pipeline_kandinsky6.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

import numpy as np
import pytest

pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]


def test_kandinsky6_pipeline_import_and_registry():
    from vllm_omni.diffusion.models.kandinsky6 import (
        AutoencoderKLHunyuanVideo,
        Kandinsky6AudioVAE,
        Kandinsky6TI2VAPipeline,
        Kandinsky6Transformer3DModel,
        KandinskyFlowMatchScheduler,
        get_kandinsky6_post_process_func,
        get_kandinsky6_pre_process_func,
    )
    from vllm_omni.diffusion.registry import (
        _DIFFUSION_MODELS,
        _DIFFUSION_POST_PROCESS_FUNCS,
        _DIFFUSION_PRE_PROCESS_FUNCS,
        _NO_CACHE_ACCELERATION,
    )

    assert Kandinsky6TI2VAPipeline is not None
    assert Kandinsky6Transformer3DModel is not None
    assert AutoencoderKLHunyuanVideo is not None
    assert Kandinsky6AudioVAE is not None
    assert KandinskyFlowMatchScheduler is not None
    assert get_kandinsky6_post_process_func is not None
    assert get_kandinsky6_pre_process_func is not None

    assert _DIFFUSION_MODELS["Kandinsky6TI2VAPipeline"] == (
        "kandinsky6",
        "pipeline_kandinsky6",
        "Kandinsky6TI2VAPipeline",
    )
    assert _DIFFUSION_POST_PROCESS_FUNCS["Kandinsky6TI2VAPipeline"] == "get_kandinsky6_post_process_func"
    assert _DIFFUSION_PRE_PROCESS_FUNCS["Kandinsky6TI2VAPipeline"] == "get_kandinsky6_pre_process_func"
    assert "Kandinsky6TI2VAPipeline" not in _NO_CACHE_ACCELERATION


def test_kandinsky6_component_discovery_declarations():
    from vllm_omni.diffusion.models.kandinsky6 import Kandinsky6TI2VAPipeline

    assert Kandinsky6TI2VAPipeline._dit_modules == ["transformer"]
    assert Kandinsky6TI2VAPipeline._encoder_modules == ["text_encoder", "text_encoder_2"]
    assert Kandinsky6TI2VAPipeline._vae_modules == ["vae", "audio_vae"]
    assert Kandinsky6TI2VAPipeline.supports_step_execution is True
    assert Kandinsky6TI2VAPipeline.support_audio_output is True
    assert Kandinsky6TI2VAPipeline.support_image_input is True


def test_kandinsky6_pipeline_satisfies_capability_protocols():
    """isinstance checks against the @runtime_checkable protocols the
    framework actually uses to detect these capabilities (io_support.py's
    supports_audio_output, etc.) — stronger than just checking the class
    attribute exists."""
    from vllm_omni.diffusion.models.interface import (
        SupportAudioOutput,
        SupportImageInput,
        SupportsComponentDiscovery,
    )
    from vllm_omni.diffusion.models.kandinsky6 import Kandinsky6TI2VAPipeline

    assert isinstance(Kandinsky6TI2VAPipeline, SupportAudioOutput)
    assert isinstance(Kandinsky6TI2VAPipeline, SupportImageInput)
    assert isinstance(Kandinsky6TI2VAPipeline, SupportsComponentDiscovery)


class _FakeAudioVAE:
    downsample_factor = 1024


def test_kandinsky6_post_process_func_packages_video_and_audio():
    """Pure-function test: no model/weights needed. Confirms the post-process
    payload shape matches what io_support.py -> output_formatter.py ->
    media_utils.py's PyAV muxing expects (the same flat {"video", "audio",
    "audio_sample_rate", "fps"} shape MiniMax H3's own post-process function
    produces), and that int16 PCM from ``postprocess_audio`` is rescaled to
    the float32 [-1, 1] waveform the muxer consumes."""
    from vllm_omni.diffusion.models.kandinsky6 import get_kandinsky6_post_process_func

    post_process = get_kandinsky6_post_process_func(od_config=None)
    video = np.zeros((1, 4, 8, 8, 3), dtype=np.uint8)
    audio = np.array([0, 32767, -32767], dtype=np.int16)

    result = post_process({"video": video, "audio": audio, "audio_sample_rate": 44100}, output_type="np")

    assert result["video"] is video
    assert result["audio"].dtype == np.float32
    np.testing.assert_allclose(result["audio"], [0.0, 1.0, -1.0])
    assert result["audio_sample_rate"] == 44100
    assert result["fps"] == 24.0


def test_kandinsky6_post_process_sample_rate_survives_output_formatter():
    """The formatter only lifts ``audio_sample_rate`` into metadata from the
    flat payload form; an envelope would silently lose it and the MP4 would
    be muxed at the 24 kHz default."""
    from vllm_omni.diffusion.models.kandinsky6 import get_kandinsky6_post_process_func
    from vllm_omni.diffusion.output_formatter import normalize_diffusion_postprocess_output

    post_process = get_kandinsky6_post_process_func(od_config=None)
    video = np.zeros((1, 2, 4, 4, 3), dtype=np.uint8)
    audio = np.zeros((10,), dtype=np.int16)

    normalized = normalize_diffusion_postprocess_output(
        post_process({"video": video, "audio": audio, "audio_sample_rate": 44100}, output_type="np")
    )

    assert normalized.primary_key == "video"
    assert normalized.metadata["audio"]["sample_rate"] == 44100
    assert normalized.metadata["video"]["fps"] == 24.0
    assert "audio_sample_rate" not in normalized.outputs


def test_kandinsky6_post_process_func_omits_audio_when_none():
    from vllm_omni.diffusion.models.kandinsky6 import get_kandinsky6_post_process_func

    post_process = get_kandinsky6_post_process_func(od_config=None)
    video = np.zeros((1, 2, 4, 4, 3), dtype=np.uint8)

    result = post_process({"video": video, "audio": None, "audio_sample_rate": None}, output_type="np")

    assert result["video"] is video
    assert "audio" not in result
    assert "audio_sample_rate" not in result


def test_kandinsky6_post_process_func_unwraps_batched_audio_list():
    from vllm_omni.diffusion.models.kandinsky6 import get_kandinsky6_post_process_func

    post_process = get_kandinsky6_post_process_func(od_config=None)
    video = np.zeros((1, 2, 4, 4, 3), dtype=np.uint8)
    audio_item = np.full((50,), 16384, dtype=np.int16)

    result = post_process({"video": video, "audio": [audio_item], "audio_sample_rate": 44100}, output_type="np")

    assert result["audio"].shape == (50,)
    np.testing.assert_allclose(result["audio"], 16384 / 32767, rtol=1e-6)


class _StubAudioVAE:
    """Audio VAE stand-in: identity "decode" that returns the latent as a waveform."""

    scaling_factor = 0.5
    mean_value = 0.0
    device = "cpu"

    def wrapped_decode(self, latents):
        # (1, audio_dim, A) -> (A,) waveform: take the first channel.
        return latents[0, 0]


def test_postprocess_audio_normalizes_by_default_and_clips_on_request():
    """The default mode must match the k6_video production pipeline, which
    peak-normalizes each decoded waveform to full scale; ``clip`` keeps the
    raw amplitude (saturated to [-1, 1]) for callers that want it."""
    import torch

    from vllm_omni.diffusion.models.kandinsky6.pipeline_kandinsky6 import LatentBundle, postprocess_audio

    # Latent (A=4, audio_dim=1); the stub scales by 1/scaling_factor -> [0.2, -0.4, 0.6, 3.0].
    audio = torch.tensor([[0.1], [-0.2], [0.3], [1.5]], dtype=torch.float32)
    bundle = LatentBundle(
        video=None,
        audio=audio,
        video_cu_seqlens=None,
        audio_cu_seqlens=torch.tensor([0, 4], dtype=torch.int32),
    )
    vae = _StubAudioVAE()

    normalized = postprocess_audio(bundle, vae)
    assert normalized is not None and len(normalized) == 1
    assert normalized[0].dtype == np.int16
    np.testing.assert_allclose(
        normalized[0].astype(np.float32) / 32767,
        np.array([0.2, -0.4, 0.6, 3.0], dtype=np.float32) / 3.0,
        atol=1e-4,
    )

    clipped = postprocess_audio(bundle, vae, normalization_mode="clip")
    np.testing.assert_allclose(
        clipped[0].astype(np.float32) / 32767,
        np.array([0.2, -0.4, 0.6, 1.0], dtype=np.float32),
        atol=1e-4,
    )

    assert postprocess_audio(LatentBundle(None, None, None, None), vae) is None
    with pytest.raises(ValueError, match="normalization_mode"):
        postprocess_audio(bundle, vae, normalization_mode="loud")


def test_kandinsky6_pre_process_func_is_identity_for_now():
    """v1 scope: the pre-process hook is registered (matching the framework's
    convention for image-conditioned models like SanaImageToVideoPipeline)
    but does not yet validate/transform the request — documented follow-up,
    not a silent gap."""
    from vllm_omni.diffusion.models.kandinsky6 import get_kandinsky6_pre_process_func

    pre_process = get_kandinsky6_pre_process_func(od_config=None)
    sentinel = object()

    assert pre_process(sentinel) is sentinel


def test_hub_component_keys_map_onto_pipeline_modules():
    """Hub folder keys from Kandinsky-6.0-Pro-5s-Diffusers land on the modules."""
    from vllm_omni.diffusion.models.kandinsky6.pipeline_kandinsky6 import (
        _WEIGHT_SUBFOLDERS,
        _adapt_k6_weight_name,
    )

    assert _WEIGHT_SUBFOLDERS == (
        ("transformer", "transformer."),
        ("vae", "vae."),
        ("text_encoder", "text_encoder."),
        ("text_encoder_2", "text_encoder_2."),
        ("audio_vae", "audio_vae."),
    )
    assert _adapt_k6_weight_name("transformer.visual_transformer_blocks.0.videoT.self_attention.to_query.weight") == (
        "transformer.visual_transformer_blocks.0.video_dec_block.self_attention.to_query.weight"
    )
    assert _adapt_k6_weight_name("transformer.visual_transformer_blocks.0.audioT.feed_forward.in_layer.weight") == (
        "transformer.visual_transformer_blocks.0.audio_dec_block.feed_forward.in_layer.weight"
    )
    assert _adapt_k6_weight_name("vae.decoder.conv_in.conv.weight") == "vae.decoder.conv_in.conv.weight"
    assert _adapt_k6_weight_name("text_encoder_2.encoder.layers.0.mlp.fc1.weight") == (
        "text_encoder_2.encoder.layers.0.mlp.fc1.weight"
    )
    assert (
        _adapt_k6_weight_name("text_encoder.model.layers.0.input_layernorm.weight")
        == "text_encoder.model.language_model.layers.0.input_layernorm.weight"
    )
    assert _adapt_k6_weight_name("text_encoder.visual.blocks.0.norm1.weight") == (
        "text_encoder.model.visual.blocks.0.norm1.weight"
    )
    assert (
        _adapt_k6_weight_name("audio_vae.vae.decoder.conv_in.weight")
        == "audio_vae.native.tod.vae.decoder.conv_in.weight"
    )
    assert _adapt_k6_weight_name("audio_vae.vocoder.conv_pre.weight") == "audio_vae.native.tod.vocoder.conv_pre.weight"
    assert _adapt_k6_weight_name("audio_vae.mel_converter.hann_window") == "audio_vae.native.mel_converter.hann_window"