vllm-omni/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py

from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import numpy as np
import PIL.Image
import pytest
import torch
import torch.nn as nn

from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import (
    Wan22S2VPipeline,
    _make_clip_generators,
)
from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import WanS2VTransformer3DModel
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch

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


def _make_s2v_sampling(**overrides):
    values = {
        "height": 16,
        "width": 16,
        "num_frames": 8,
        "num_inference_steps": 1,
        "guidance_scale_provided": True,
        "guidance_scale": 1.0,
        "generator": None,
        "seed": None,
        "num_outputs_per_prompt": 1,
        "max_sequence_length": 8,
        "latents": None,
        "output_type": "latent",
        "extra_args": {},
    }
    values.update(overrides)
    return SimpleNamespace(**values)


def _make_s2v_validation_pipeline() -> Wan22S2VPipeline:
    pipeline = object.__new__(Wan22S2VPipeline)
    nn.Module.__init__(pipeline)
    pipeline.device = torch.device("cpu")
    pipeline.transformer = SimpleNamespace(dtype=torch.float32)
    pipeline.vae_scale_factor_spatial = 8
    pipeline.resolution_divisor = 16
    pipeline.motion_frames = 7
    pipeline.drop_first_motion = True
    pipeline._DEFAULT_INFER_FRAMES = 8
    pipeline._guidance_scale = None
    pipeline.check_inputs = lambda *args, **kwargs: None
    pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None)  # type: ignore[method-assign]
    return pipeline


def test_s2v_predict_noise_keeps_batch_dimension() -> None:
    pipeline = object.__new__(Wan22S2VPipeline)
    nn.Module.__init__(pipeline)
    pipeline.device = torch.device("cpu")
    transformer = MagicMock()
    transformer.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)])
    transformer.return_value = (torch.zeros(1, 16, 2, 2, 2),)
    pipeline.transformer = transformer

    result = pipeline.predict_noise(hidden_states=torch.zeros(1, 16, 2, 2, 2))

    assert result.shape == (1, 16, 2, 2, 2)


def test_s2v_clip_generators_preserve_main_seed_per_clip_behavior() -> None:
    explicit_generator = torch.Generator(device="cpu").manual_seed(999)

    generators = _make_clip_generators(
        seeds=[1234, 5678, None, None],
        request_generators=[
            torch.Generator(device="cpu").manual_seed(1234),
            torch.Generator(device="cpu").manual_seed(5678),
            explicit_generator,
            None,
        ],
        clip_index=2,
        device=torch.device("cpu"),
    )

    assert [generator.initial_seed() for generator in generators if generator is not None] == [1236, 5680, 999, 2]
    assert generators[2] is explicit_generator


def test_s2v_forward_rejects_different_num_repeat() -> None:
    pipeline = _make_s2v_validation_pipeline()
    pipeline.encode_audio = lambda *args, **kwargs: (torch.zeros(1, 1, 2, 16), 2, 16)  # type: ignore[method-assign]
    image = PIL.Image.new("RGB", (16, 16))
    audio = np.zeros(16000, dtype=np.float32)
    batch = DiffusionRequestBatch(
        requests=[
            SimpleNamespace(
                request_id="a",
                prompt={
                    "prompt": "first",
                    "multi_modal_data": {"image": image, "audio": audio},
                    "additional_information": {"num_repeat": 1},
                },
                sampling_params=_make_s2v_sampling(),
            ),
            SimpleNamespace(
                request_id="b",
                prompt={
                    "prompt": "second",
                    "multi_modal_data": {"image": image, "audio": audio},
                    "additional_information": {"num_repeat": 2},
                },
                sampling_params=_make_s2v_sampling(),
            ),
        ]
    )

    with pytest.raises(ValueError, match="num_repeat"):
        pipeline.forward(batch)


def test_s2v_forward_rejects_different_audio_time_lengths() -> None:
    pipeline = _make_s2v_validation_pipeline()
    pipeline.encode_audio = lambda audio, **kwargs: (  # type: ignore[method-assign]
        torch.zeros(1, 1, 2, 16 if audio[0] == 0 else 24),
        1,
        16 if audio[0] == 0 else 24,
    )
    image = PIL.Image.new("RGB", (16, 16))
    batch = DiffusionRequestBatch(
        requests=[
            SimpleNamespace(
                request_id="a",
                prompt={
                    "prompt": "first",
                    "multi_modal_data": {"image": image, "audio": np.zeros(16000, dtype=np.float32)},
                },
                sampling_params=_make_s2v_sampling(),
            ),
            SimpleNamespace(
                request_id="b",
                prompt={
                    "prompt": "second",
                    "multi_modal_data": {"image": image, "audio": np.ones(16000, dtype=np.float32)},
                },
                sampling_params=_make_s2v_sampling(),
            ),
        ]
    )

    with pytest.raises(ValueError, match="audio embedding time lengths"):
        pipeline.forward(batch)


def test_s2v_forward_rejects_different_init_first_frame() -> None:
    pipeline = _make_s2v_validation_pipeline()
    image = PIL.Image.new("RGB", (16, 16))
    audio = np.zeros(16000, dtype=np.float32)
    batch = DiffusionRequestBatch(
        requests=[
            SimpleNamespace(
                request_id="a",
                prompt={
                    "prompt": "first",
                    "multi_modal_data": {"image": image, "audio": audio},
                    "additional_information": {"init_first_frame": True},
                },
                sampling_params=_make_s2v_sampling(),
            ),
            SimpleNamespace(
                request_id="b",
                prompt={
                    "prompt": "second",
                    "multi_modal_data": {"image": image, "audio": audio},
                    "additional_information": {"init_first_frame": False},
                },
                sampling_params=_make_s2v_sampling(),
            ),
        ]
    )

    with pytest.raises(ValueError, match="init_first_frame"):
        pipeline.forward(batch)


def test_s2v_forward_rejects_different_raw_audio_shapes() -> None:
    pipeline = _make_s2v_validation_pipeline()
    pipeline.encode_audio = lambda *args, **kwargs: (torch.zeros(1, 1, 2, 16), 1, 16)  # type: ignore[method-assign]
    image = PIL.Image.new("RGB", (16, 16))
    batch = DiffusionRequestBatch(
        requests=[
            SimpleNamespace(
                request_id="a",
                prompt={
                    "prompt": "first",
                    "multi_modal_data": {"image": image, "audio": np.zeros(16000, dtype=np.float32)},
                },
                sampling_params=_make_s2v_sampling(),
            ),
            SimpleNamespace(
                request_id="b",
                prompt={
                    "prompt": "second",
                    "multi_modal_data": {"image": image, "audio": np.zeros(8000, dtype=np.float32)},
                },
                sampling_params=_make_s2v_sampling(),
            ),
        ]
    )

    with pytest.raises(ValueError, match="raw audio shapes and sample rates"):
        pipeline.forward(batch)


def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None:
    pipeline = object.__new__(Wan22S2VPipeline)
    nn.Module.__init__(pipeline)
    pipeline.device = torch.device("cpu")
    pipeline.transformer = MagicMock()
    pipeline.transformer.dtype = torch.float32
    pipeline.transformer.parameters.return_value = iter([torch.zeros(1, dtype=torch.float32)])
    pipeline.transformer.casual_audio_encoder = None
    pipeline.transformer.encode_audio.side_effect = lambda audio, _motion: {"audio_emb": audio}
    pipeline.vae = MagicMock()
    pipeline.vae.dtype = torch.float32
    pipeline.vae.decode.return_value = (torch.zeros(4, 3, 8, 16, 16),)
    pipeline.od_config = SimpleNamespace(
        enable_cpu_offload=False,
        parallel_config=SimpleNamespace(use_hsdp=False),
    )
    pipeline.scheduler = MagicMock(timesteps=torch.tensor([1.0]))
    pipeline.vae_scale_factor_spatial = 8
    pipeline.resolution_divisor = 16
    pipeline.motion_frames = 7
    pipeline.drop_first_motion = True
    pipeline._DEFAULT_INFER_FRAMES = 8
    pipeline._guidance_scale = None
    pipeline._num_timesteps = None
    pipeline.check_inputs = lambda *args, **kwargs: None
    pipeline.encode_prompt = MagicMock(return_value=(torch.zeros(4, 2, 3), torch.ones(4, 2, 3)))
    pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 1, 2, 8), 1, 8))
    pipeline.encode_ref_image = MagicMock(return_value=torch.zeros(1, 16, 1, 2, 2))
    pipeline.prepare_motion_latents = MagicMock(
        side_effect=lambda pixels, **_: torch.zeros(pixels.shape[0], 16, 2, 2, 2)
    )
    pipeline._denormalize_latents = lambda latents: latents
    pipeline.diffuse = MagicMock(side_effect=lambda **kwargs: kwargs["latents"])

    image = PIL.Image.new("RGB", (16, 16))
    audio_a = np.zeros(16000, dtype=np.float32)
    audio_b = np.ones(16000, dtype=np.float32)
    latents_a = torch.zeros(2, 16, 2, 2, 2)
    latents_b = torch.ones(2, 16, 2, 2, 2)
    generators_a = [torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)]
    generators_b = [torch.Generator().manual_seed(3), torch.Generator().manual_seed(4)]
    batch = DiffusionRequestBatch(
        requests=[
            SimpleNamespace(
                request_id="a",
                prompt={
                    "prompt": "first",
                    "negative_prompt": "negative first",
                    "multi_modal_data": {"image": image, "audio": audio_a},
                },
                sampling_params=_make_s2v_sampling(
                    num_outputs_per_prompt=2,
                    generator=generators_a,
                    latents=latents_a,
                ),
            ),
            SimpleNamespace(
                request_id="b",
                prompt={
                    "prompt": "second",
                    "negative_prompt": "negative second",
                    "multi_modal_data": {"image": image, "audio": audio_b},
                },
                sampling_params=_make_s2v_sampling(
                    num_outputs_per_prompt=2,
                    generator=generators_b,
                    latents=latents_b,
                ),
            ),
        ]
    )

    with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as platform:
        platform.is_available.return_value = False
        outputs = pipeline.forward(batch)

    assert len(outputs) == 2
    assert outputs[0].output[0].shape[0] == 2
    assert outputs[1].output[0].shape[0] == 2
    assert outputs[0].output[1].shape == (16000,)
    assert outputs[1].output[1].shape == (16000,)
    np.testing.assert_array_equal(outputs[0].output[1], audio_a)
    np.testing.assert_array_equal(outputs[1].output[1], audio_b)
    assert np.shares_memory(outputs[0].output[1], audio_a)
    assert np.shares_memory(outputs[1].output[1], audio_b)
    torch.testing.assert_close(pipeline.diffuse.call_args.kwargs["latents"], torch.cat([latents_a, latents_b]))
    assert pipeline.diffuse.call_args.kwargs["clip_generator"] == generators_a + generators_b
    assert pipeline.encode_prompt.call_args.kwargs["prompt"] == ["first", "second"]
    assert pipeline.encode_prompt.call_args.kwargs["negative_prompt"] == ["negative first", "negative second"]


def test_s2v_exposes_hsdp_shard_conditions_for_transformer_blocks():
    model = object.__new__(WanS2VTransformer3DModel)
    nn.Module.__init__(model)
    model.blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(3)])

    conditions = getattr(model, "_hsdp_shard_conditions", None)

    assert conditions is not None
    assert len(conditions) == 1

    matched = []
    for name, module in model.named_modules():
        if any(cond(name, module) for cond in conditions):
            matched.append(name)

    assert matched == ["blocks.0", "blocks.1", "blocks.2"]


def test_s2v_hsdp_shard_condition_does_not_match_non_block_modules():
    model = object.__new__(WanS2VTransformer3DModel)
    nn.Module.__init__(model)
    model.blocks = nn.ModuleList([nn.Linear(4, 4)])
    model.head_indicator = nn.Linear(4, 4)
    model.casual_audio_encoder = nn.Linear(4, 4)

    conditions = model._hsdp_shard_conditions
    non_block_matched = []
    for name, module in model.named_modules():
        if name and "blocks" not in name:
            if any(cond(name, module) for cond in conditions):
                non_block_matched.append(name)

    assert non_block_matched == []


def test_encode_audio_calls_unshard_reshard_when_fsdp_managed():
    model = object.__new__(WanS2VTransformer3DModel)
    nn.Module.__init__(model)
    model.enable_adain = False
    model.casual_audio_encoder = MagicMock(return_value=torch.zeros(1, 10, 64))

    model.unshard = MagicMock()
    model.reshard = MagicMock()

    audio_input = torch.randn(1, 1, 64, 5)
    motion_frames = [2, 2]

    result = model.encode_audio(audio_input, motion_frames)

    model.unshard.assert_called_once()
    model.reshard.assert_called_once()
    assert "audio_emb" in result


def test_encode_audio_reshard_called_on_exception():
    """Test that reshard() is always called even when encode_audio logic raises."""
    model = object.__new__(WanS2VTransformer3DModel)
    nn.Module.__init__(model)
    model.enable_adain = False
    model.casual_audio_encoder = MagicMock(side_effect=RuntimeError("encoder failed"))

    model.unshard = MagicMock()
    model.reshard = MagicMock()

    audio_input = torch.randn(1, 1, 64, 5)
    motion_frames = [2, 2]

    with pytest.raises(RuntimeError, match="encoder failed"):
        model.encode_audio(audio_input, motion_frames)

    model.unshard.assert_called_once()
    model.reshard.assert_called_once()


def test_encode_audio_skips_unshard_reshard_when_not_fsdp():
    model = object.__new__(WanS2VTransformer3DModel)
    nn.Module.__init__(model)
    model.enable_adain = False
    model.casual_audio_encoder = MagicMock(return_value=torch.zeros(1, 10, 64))

    audio_input = torch.randn(1, 1, 64, 5)
    motion_frames = [2, 2]

    result = model.encode_audio(audio_input, motion_frames)

    assert not hasattr(model, "unshard")
    assert not hasattr(model, "reshard")
    assert "audio_emb" in result


def test_s2v_pipeline_skips_cpu_offload_when_hsdp_enabled():
    """Test that transformer.to('cpu') is NOT called when HSDP is active."""
    from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import Wan22S2VPipeline

    pipeline = object.__new__(Wan22S2VPipeline)
    nn.Module.__init__(pipeline)

    od_config = MagicMock()
    od_config.enable_cpu_offload = True
    parallel_config = MagicMock()
    parallel_config.use_hsdp = True
    od_config.parallel_config = parallel_config
    pipeline.od_config = od_config

    mock_transformer = MagicMock()
    pipeline.transformer = mock_transformer

    # Simulate the offload decision from the forward loop
    if pipeline.od_config.enable_cpu_offload and not getattr(pipeline.od_config.parallel_config, "use_hsdp", False):
        pipeline.transformer.to("cpu")

    mock_transformer.to.assert_not_called()


def test_s2v_pipeline_allows_cpu_offload_when_hsdp_disabled():
    """Test that transformer.to('cpu') IS called when HSDP is not active."""
    from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import Wan22S2VPipeline

    pipeline = object.__new__(Wan22S2VPipeline)
    nn.Module.__init__(pipeline)

    od_config = MagicMock()
    od_config.enable_cpu_offload = True
    parallel_config = MagicMock()
    parallel_config.use_hsdp = False
    od_config.parallel_config = parallel_config
    pipeline.od_config = od_config

    mock_transformer = MagicMock()
    pipeline.transformer = mock_transformer

    # Simulate the offload decision from the forward loop
    if pipeline.od_config.enable_cpu_offload and not getattr(pipeline.od_config.parallel_config, "use_hsdp", False):
        pipeline.transformer.to("cpu")

    mock_transformer.to.assert_called_once_with("cpu")


def test_s2v_pipeline_hsdp_forward_complete_process():
    """Integration-level mock test verifying the complete forward path of the
    S2V pipeline under HSDP mode.

    Verifies:
    - Text encoding is invoked
    - Audio encoding invokes unshard/reshard on the transformer
    - Reference image encoding via VAE
    - Denoising loop calls the transformer
    - VAE decode is invoked
    - transformer.to('cpu') is NOT called (HSDP mode)
    - Final output is a DiffusionOutput with video + audio
    """
    import numpy as np
    import PIL.Image

    from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import Wan22S2VPipeline

    pipeline = object.__new__(Wan22S2VPipeline)
    nn.Module.__init__(pipeline)

    # -- Config --
    od_config = MagicMock()
    od_config.enable_cpu_offload = True
    od_config.enable_diffusion_pipeline_profiler = False
    parallel_config = MagicMock()
    parallel_config.use_hsdp = True
    od_config.parallel_config = parallel_config
    pipeline.od_config = od_config

    # -- Pipeline attributes --
    pipeline.device = torch.device("cpu")
    pipeline._guidance_scale = 4.5
    pipeline._num_timesteps = None
    pipeline._current_timestep = None
    pipeline.vae_scale_factor_spatial = 8
    pipeline.vae_scale_factor_temporal = 4
    pipeline.resolution_divisor = 16
    pipeline.motion_frames = 73
    pipeline.drop_first_motion = True
    pipeline.fps = 16
    pipeline.audio_sample_m = 0
    pipeline._DEFAULT_INFER_FRAMES = 80

    # -- Mock text encoder --
    pipeline.tokenizer = MagicMock()
    pipeline.tokenizer.return_value = MagicMock(
        input_ids=torch.zeros(1, 512, dtype=torch.long),
        attention_mask=torch.ones(1, 512, dtype=torch.long),
    )
    mock_text_encoder = MagicMock()
    mock_text_encoder.dtype = torch.bfloat16
    mock_text_encoder.return_value = MagicMock(last_hidden_state=torch.zeros(1, 512, 4096))
    pipeline.text_encoder = mock_text_encoder

    # -- Mock transformer --
    mock_transformer = MagicMock()
    mock_transformer.dtype = torch.bfloat16
    mock_transformer.parameters = MagicMock(return_value=iter([torch.zeros(1)]))

    # encode_audio returns dict with audio_emb
    mock_transformer.encode_audio = MagicMock(return_value={"audio_emb": torch.zeros(1, 10, 64)})
    # Forward returns noise prediction
    mock_transformer.return_value = (torch.zeros(1, 16, 20, 88, 128),)
    pipeline.transformer = mock_transformer

    # -- Mock VAE --
    mock_vae = MagicMock()
    mock_vae.dtype = torch.bfloat16
    mock_vae.config = MagicMock()
    mock_vae.config.scale_factor_temporal = 4
    mock_vae.config.scale_factor_spatial = 8
    mock_vae.config.latents_mean = [0.0] * 16
    mock_vae.config.latents_std = [1.0] * 16
    mock_vae.config.z_dim = 16
    # encode returns mock with latent_dist
    mock_encode_result = MagicMock()
    mock_encode_result.latent_dist = MagicMock()
    mock_encode_result.latent_dist.mode = MagicMock(return_value=torch.zeros(1, 16, 1, 88, 128))
    mock_vae.encode = MagicMock(return_value=mock_encode_result)
    # decode returns video
    mock_vae.decode = MagicMock(return_value=(torch.zeros(1, 3, 80, 704, 1024),))
    pipeline.vae = mock_vae

    # -- Mock audio model --
    mock_audio_model = MagicMock()
    mock_audio_model.device = torch.device("cpu")
    mock_audio_param = torch.zeros(1)
    mock_audio_model.parameters = MagicMock(return_value=iter([mock_audio_param]))
    mock_audio_model.return_value = MagicMock(hidden_states=[torch.zeros(1, 100, 1024)] * 25)
    pipeline.audio_model = mock_audio_model

    pipeline.audio_processor = MagicMock()
    pipeline.audio_processor.return_value = MagicMock(input_values=torch.zeros(1, 16000))

    # -- Mock scheduler --
    mock_scheduler = MagicMock()
    mock_scheduler.timesteps = torch.linspace(999, 0, 5)
    mock_scheduler.step = MagicMock(return_value=(torch.zeros(1, 16, 20, 88, 128),))
    pipeline.scheduler = mock_scheduler

    # -- Bind methods from the real class --
    pipeline.encode_prompt = Wan22S2VPipeline.encode_prompt.__get__(pipeline)
    pipeline.encode_ref_image = Wan22S2VPipeline.encode_ref_image.__get__(pipeline)
    pipeline.prepare_motion_latents = Wan22S2VPipeline.prepare_motion_latents.__get__(pipeline)
    pipeline.prepare_latents = Wan22S2VPipeline.prepare_latents.__get__(pipeline)
    pipeline.check_inputs = Wan22S2VPipeline.check_inputs.__get__(pipeline)
    pipeline.diffuse = Wan22S2VPipeline.diffuse.__get__(pipeline)
    pipeline._normalize_latents = Wan22S2VPipeline._normalize_latents.__get__(pipeline)
    pipeline._denormalize_latents = Wan22S2VPipeline._denormalize_latents.__get__(pipeline)
    pipeline._prompt_clean = Wan22S2VPipeline._prompt_clean

    # -- Build request --
    sampling_params = _make_s2v_sampling(
        height=704,
        width=1024,
        num_frames=80,
        num_inference_steps=5,
        guidance_scale=4.5,
        generator=torch.Generator(device="cpu").manual_seed(42),
        max_sequence_length=512,
    )

    ref_image = PIL.Image.new("RGB", (1024, 704))
    audio_data = np.zeros(16000, dtype=np.float32)

    req = DiffusionRequestBatch(
        requests=[
            SimpleNamespace(
                request_id="test",
                prompt={
                    "prompt": "test prompt",
                    "negative_prompt": "bad quality",
                    "multi_modal_data": {"image": ref_image, "audio": audio_data},
                    "additional_information": {
                        "audio_path": audio_data,
                        "pose_video": None,
                        "init_first_frame": False,
                    },
                },
                sampling_params=sampling_params,
            )
        ]
    )

    # -- Mock methods that use complex internal state --
    pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 25, 64, 80), 1, 80))

    # Mock progress bar context
    pipeline.progress_bar = MagicMock()
    pipeline.progress_bar.return_value.__enter__ = MagicMock(return_value=MagicMock())
    pipeline.progress_bar.return_value.__exit__ = MagicMock(return_value=False)

    # Mock predict_noise_maybe_with_cfg from CFGParallelMixin
    pipeline.predict_noise_maybe_with_cfg = MagicMock(return_value=torch.zeros(1, 16, 20, 88, 128))
    pipeline.predict_noise = Wan22S2VPipeline.predict_noise.__get__(pipeline)

    # Mock platform methods
    with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as mock_platform:
        mock_platform.empty_cache = MagicMock()
        mock_platform.is_available = MagicMock(return_value=False)

        with patch(
            "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.load_audio", return_value=(audio_data, 16000)
        ):
            result = Wan22S2VPipeline.forward(pipeline, req=req)[0]

    # -- Assertions --
    # Text encoder was called
    mock_text_encoder.assert_called()

    # Audio encoding was invoked
    pipeline.encode_audio.assert_called_once()

    # VAE encode was called (for ref image and motion latents)
    mock_vae.encode.assert_called()

    # VAE decode was called
    mock_vae.decode.assert_called()

    # transformer.to("cpu") must NOT be called under HSDP
    mock_transformer.to.assert_not_called()

    # Output is a DiffusionOutput tuple with (video, audio_waveform, sample_rate)
    from vllm_omni.diffusion.data import DiffusionOutput

    assert isinstance(result, DiffusionOutput)
    video, audio_waveform, audio_sr = result.output
    assert video.shape[0] == 1  # batch
    assert video.shape[1] == 3  # channels
    assert audio_waveform is not None
    assert audio_sr == 16000