vllm-omni/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py

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

import weakref
from contextlib import contextmanager, nullcontext
from types import SimpleNamespace
from typing import Any
from unittest.mock import Mock

import pytest
import torch

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


def test_h3_full_pipeline_profiler_includes_local_encoding():
    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline

    targets = MiniMaxH3Pipeline._PROFILER_TARGETS
    assert targets == [
        "encode_prompt",
        "_encode_local_media",
        "diffuse",
        "decode",
        "video_vae.decode_latent",
        "audio_vae.decode_latent",
        "prepare_encode",
        "denoise_step",
        "post_decode",
    ]


def test_decoded_output_is_unchanged_without_active_model_offload():
    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline

    pipeline = object.__new__(MiniMaxH3Pipeline)
    torch.nn.Module.__init__(pipeline)
    pipeline._model_cpu_offload_modules = []
    output = Mock()

    assert pipeline._offload_model_cpu_stage_output(output) is output
    output.cpu.assert_not_called()


def test_reply_rank_copies_decoded_output_to_host(monkeypatch):
    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline

    pipeline = object.__new__(MiniMaxH3Pipeline)
    torch.nn.Module.__init__(pipeline)
    pipeline._model_cpu_offload_modules = [Mock()]
    monkeypatch.setattr(pipeline, "_is_output_owner_rank", lambda: True)
    output = Mock()
    host_output = Mock()
    output.cpu.return_value = host_output

    assert pipeline._offload_model_cpu_stage_output(output) is host_output
    output.cpu.assert_called_once_with()


def test_non_reply_rank_drops_decoded_output_storage(monkeypatch):
    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline

    pipeline = object.__new__(MiniMaxH3Pipeline)
    torch.nn.Module.__init__(pipeline)
    pipeline._model_cpu_offload_modules = [Mock()]
    monkeypatch.setattr(pipeline, "_is_output_owner_rank", lambda: False)
    output = torch.ones((2, 3), dtype=torch.float16)

    placeholder = pipeline._offload_model_cpu_stage_output(output)

    assert placeholder.device.type == "meta"
    assert placeholder.shape == output.shape
    assert placeholder.dtype == output.dtype


def test_decode_releases_only_audio_before_return(monkeypatch):
    """Video is released by the callers of ``decode`` -- after quantization."""
    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module

    video = torch.ones((1, 1, 4, 4))
    audio = torch.ones((1, 8))
    pipeline = SimpleNamespace(
        device=torch.device("cpu"),
        video_vae=Mock(decode_latent=Mock(return_value=video)),
        audio_vae=Mock(decode_latent=Mock(return_value=audio)),
        _component_on_device=lambda component: nullcontext(),
        _offload_model_cpu_stage_output=Mock(side_effect=lambda value: value),
        _release_stage_cache=Mock(),
    )
    monkeypatch.setattr(
        module.current_omni_platform,
        "create_autocast_context",
        lambda **kwargs: nullcontext(),
    )

    actual_video, actual_audio = MiniMaxH3Pipeline.decode(
        pipeline,
        torch.zeros(1),
        torch.zeros(1),
        height=2,
        width=3,
    )

    assert actual_video.shape[-2:] == (2, 3)
    assert actual_audio is audio
    calls = pipeline._offload_model_cpu_stage_output.call_args_list
    assert len(calls) == 1
    assert calls[0].args[0] is audio


def _recording_release(events, releases):
    """Stand-in for ``_offload_model_cpu_stage_output`` that records what it got."""

    def release(value):
        events.append("release")
        releases.append(value)
        return value

    return release


def test_forward_releases_video_after_quantization_and_before_next_seed(monkeypatch):
    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline

    # MiniMaxH3VideoVAE.decode_latent ends in frames.float(), so this is what
    # the release hook would see if it ran before the quantizer.
    decoded_dtype = torch.float32
    events: list[str] = []
    releases: list[torch.Tensor] = []

    def diffuse(**kwargs):
        events.append(f"diffuse:{kwargs['seed']}")
        return torch.zeros(1), torch.zeros(1)

    def decode(video_latent, audio_latent, *, height, width):
        events.append("decode")
        return (
            torch.zeros((1, 3, 2, height, width), dtype=decoded_dtype),
            torch.zeros((1, 8)),
        )

    # ``forward`` indexes ``preencode_mp4`` (#7018): this test drives the
    # tensor-decode branch, which is the one the release hook wraps.
    context = {"num_outputs": 2, "seed": 41, "height": 4, "width": 4, "preencode_mp4": False}
    pipeline = SimpleNamespace(
        od_config=SimpleNamespace(),
        _extract_prompt=lambda raw: ("a prompt", {}),
        _extract_text_conditioning=lambda raw: None,
        _extract_prepared_reference_videos=lambda raw: None,
        _prepare_request_inputs=lambda *args, **kwargs: context,
        _denoise_kwargs=lambda ctx: {},
        diffuse=diffuse,
        decode=decode,
        _offload_model_cpu_stage_output=_recording_release(events, releases),
        _release_stage_cache=lambda: None,
    )
    request = SimpleNamespace(prompts=["a prompt"], sampling_params=SimpleNamespace())

    output = MiniMaxH3Pipeline.forward(pipeline, request)

    # The release must see the quantized frames, not the decoded floats.
    assert len(releases) == 2
    assert [tensor.dtype for tensor in releases] == [torch.uint8, torch.uint8]
    assert decoded_dtype is not torch.uint8
    # ... and it must still happen before the next seed reloads the DiT.
    assert events == ["diffuse:41", "decode", "release", "diffuse:42", "decode", "release"]
    assert output.output[0].dtype is torch.uint8


def test_forward_drops_the_decoded_video_before_the_next_seed_diffuses(monkeypatch):
    """Releasing the copy is not enough if a local still pins the original.

    ``_offload_model_cpu_stage_output`` hands back a *new* tensor, so the
    decoded frames stay alive for as long as anything still references them.
    The whole point of this branch is that they are gone before the next seed
    reloads the DiT, and ``diffuse`` is where that reload happens.
    """
    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline

    decoded_refs: list[weakref.ref] = []
    alive_at_diffuse: list[int] = []

    def diffuse(**kwargs):
        alive_at_diffuse.append(sum(1 for ref in decoded_refs if ref() is not None))
        return torch.zeros(1), torch.zeros(1)

    def decode(video_latent, audio_latent, *, height, width):
        video = torch.zeros((1, 3, 2, height, width), dtype=torch.float32)
        decoded_refs.append(weakref.ref(video))
        return video, torch.zeros((1, 8))

    # ``forward`` indexes ``preencode_mp4`` (#7018): this test drives the
    # tensor-decode branch, which is the one the release hook wraps.
    context = {"num_outputs": 2, "seed": 41, "height": 4, "width": 4, "preencode_mp4": False}
    pipeline = SimpleNamespace(
        od_config=SimpleNamespace(),
        _extract_prompt=lambda raw: ("a prompt", {}),
        _extract_text_conditioning=lambda raw: None,
        _extract_prepared_reference_videos=lambda raw: None,
        _prepare_request_inputs=lambda *args, **kwargs: context,
        _denoise_kwargs=lambda ctx: {},
        diffuse=diffuse,
        decode=decode,
        _offload_model_cpu_stage_output=lambda value: value.clone(),
        _release_stage_cache=lambda: None,
    )
    request = SimpleNamespace(prompts=["a prompt"], sampling_params=SimpleNamespace())

    MiniMaxH3Pipeline.forward(pipeline, request)

    assert len(decoded_refs) == 2
    # Seed 1's diffuse sees nothing decoded yet; seed 2's must not see seed 1's
    # decoded frames still resident.
    assert alive_at_diffuse == [0, 0]


def test_post_decode_releases_video_after_quantization(monkeypatch):
    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module

    events: list[str] = []
    releases: list[torch.Tensor] = []

    def decode(video_latent, audio_latent, *, height, width):
        events.append("decode")
        return (
            torch.zeros((1, 3, 2, height, width), dtype=torch.float16),
            torch.zeros((1, 8)),
        )

    pipeline = SimpleNamespace(
        od_config=SimpleNamespace(),
        _unpack_denoised_rows=lambda *args, **kwargs: (torch.zeros(1), torch.zeros(1)),
        decode=decode,
        _offload_model_cpu_stage_output=_recording_release(events, releases),
        _release_stage_cache=lambda: None,
    )
    state = SimpleNamespace(
        latents=torch.zeros(1),
        extra={
            module._STEP_BRANCH: object(),
            module._STEP_AUDIO_ROWS: torch.zeros(1),
            module._STEP_SHAPE: {
                "latent_t": 1,
                "latent_h": 1,
                "latent_w": 1,
                "audio_t": 1,
                "height": 4,
                "width": 4,
            },
        },
    )

    output = MiniMaxH3Pipeline.post_decode(pipeline, state)

    assert events == ["decode", "release"]
    assert len(releases) == 1
    assert releases[0].dtype is torch.uint8
    assert output.output[0] is releases[0]


@pytest.mark.parametrize(
    ("available", "initialized", "rank", "expected"),
    [(False, False, 7, True), (True, False, 7, True), (True, True, 0, True), (True, True, 1, False)],
)
def test_output_owner_rank_matches_executor_contract(monkeypatch, available, initialized, rank, expected):
    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module

    monkeypatch.setattr(module.dist, "is_available", lambda: available)
    monkeypatch.setattr(module.dist, "is_initialized", lambda: initialized)
    monkeypatch.setattr(module.dist, "get_rank", lambda: rank)

    assert MiniMaxH3Pipeline._is_output_owner_rank() is expected


@pytest.mark.parametrize("load_text_encoder", [True, False])
def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch, load_text_encoder):
    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module

    pipeline = object.__new__(MiniMaxH3Pipeline)
    torch.nn.Module.__init__(pipeline)
    pipeline.transformer = torch.nn.Linear(2, 2)
    pipeline.transformers_ref = torch.nn.Linear(2, 2)
    pipeline.video_vae = torch.nn.Linear(2, 2)
    pipeline.audio_vae = torch.nn.Linear(2, 2)
    pipeline._encoder_modules = ["text_encoder"] if load_text_encoder else []
    if load_text_encoder:
        pipeline.text_encoder = torch.nn.Linear(2, 2)
    apply_offload = Mock()
    remove_offload = Mock()
    monkeypatch.setattr(module, "apply_sequential_offload", apply_offload)
    monkeypatch.setattr(module, "remove_sequential_offload", remove_offload)

    pipeline.enable_omni_model_cpu_offload(
        device=torch.device("cpu"),
        pin_memory=False,
        use_hsdp=False,
    )

    dits = [pipeline.transformer, pipeline.transformers_ref]
    stages = [*([pipeline.text_encoder] if load_text_encoder else []), pipeline.video_vae, pipeline.audio_vae]
    apply_offload.assert_called_once_with(
        dit_modules=dits,
        encoder_modules=stages,
        device=torch.device("cpu"),
        pin_memory=False,
        use_hsdp=False,
        offload_initial_dits=True,
    )

    pipeline.disable_omni_model_cpu_offload()

    remove_offload.assert_called_once_with([*dits, *stages])


@pytest.mark.parametrize("decode_fails", [False, True])
def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails):
    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module

    pipeline = object.__new__(MiniMaxH3Pipeline)
    torch.nn.Module.__init__(pipeline)
    component = torch.nn.Linear(2, 2)
    events: list[tuple[Any, ...]] = []

    @contextmanager
    def record_component(value):
        events.append(("activate", value))
        try:
            yield
        finally:
            events.append(("offload", value))

    monkeypatch.setattr(module, "sequential_offload_component", record_component)
    pipeline._model_cpu_offload_modules = [component]

    def decode():
        with pipeline._component_on_device(component):
            events.append(("decode", component))
            if decode_fails:
                raise RuntimeError("decode failed")

    if decode_fails:
        with pytest.raises(RuntimeError, match="decode failed"):
            decode()
    else:
        decode()

    assert events == [
        ("activate", component),
        ("decode", component),
        ("offload", component),
    ]


def test_h3_model_cpu_offload_keeps_unselected_vaes_resident(monkeypatch):
    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module

    pipeline = object.__new__(MiniMaxH3Pipeline)
    torch.nn.Module.__init__(pipeline)
    pipeline.transformer = torch.nn.Linear(2, 2)
    pipeline.transformers_ref = torch.nn.Linear(2, 2)
    pipeline.text_encoder = torch.nn.Linear(2, 2)
    pipeline.video_vae = torch.nn.Linear(2, 2)
    pipeline.audio_vae = torch.nn.Linear(2, 2)
    apply_offload = Mock()
    monkeypatch.setattr(module, "apply_sequential_offload", apply_offload)

    pipeline.enable_omni_model_cpu_offload(
        device=torch.device("cpu"),
        pin_memory=False,
        use_hsdp=False,
        offload_components=frozenset({"dit", "text_encoder"}),
    )

    dits = [pipeline.transformer, pipeline.transformers_ref]
    stages = [pipeline.text_encoder, pipeline.video_vae, pipeline.audio_vae]
    apply_offload.assert_called_once_with(
        dit_modules=dits,
        encoder_modules=stages,
        device=torch.device("cpu"),
        pin_memory=False,
        use_hsdp=False,
        offload_initial_dits=True,
        offload_dit_modules=dits,
        offload_encoder_modules=[pipeline.text_encoder],
    )


def test_h3_model_cpu_offload_rejects_unloaded_selected_encoder(monkeypatch):
    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module

    pipeline = object.__new__(MiniMaxH3Pipeline)
    torch.nn.Module.__init__(pipeline)
    pipeline.transformer = torch.nn.Linear(2, 2)
    pipeline.transformers_ref = torch.nn.Linear(2, 2)
    pipeline.text_encoder = None
    pipeline.video_vae = torch.nn.Linear(2, 2)
    pipeline.audio_vae = torch.nn.Linear(2, 2)
    apply_offload = Mock()
    monkeypatch.setattr(module, "apply_sequential_offload", apply_offload)

    with pytest.raises(ValueError, match="no loaded text encoder"):
        pipeline.enable_omni_model_cpu_offload(
            device=torch.device("cpu"),
            pin_memory=False,
            use_hsdp=False,
            offload_components=frozenset({"text_encoder"}),
        )

    apply_offload.assert_not_called()


def _encode_media(pipeline, **kwargs):
    from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderMediaInput
    from vllm_omni.model_executor.models.minimax_h3.encoder_processing import encode_media

    media = MiniMaxH3EncoderMediaInput(
        task="ref2va", height=32, width=32, num_frames=124, latent_t=31, audio_t=208, **kwargs
    )
    return encode_media(
        media,
        video_vae=getattr(pipeline, "video_vae", None),
        audio_vae=getattr(pipeline, "audio_vae", None),
        emit_conditioning=True,
        component_scope=pipeline._component_on_device,
    )


@pytest.mark.parametrize("encode_fails", [False, True])
def test_h3_model_cpu_offload_batches_visual_reference_scope(monkeypatch, encode_fails):
    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module

    pipeline = object.__new__(MiniMaxH3Pipeline)
    torch.nn.Module.__init__(pipeline)
    pipeline.video_vae = Mock()
    pipeline._model_cpu_offload_modules = [pipeline.video_vae]
    events = []

    @contextmanager
    def record_component(value):
        events.append(("activate", value))
        try:
            yield
        finally:
            events.append(("offload", value))

    def encode_image(image):
        events.append(("encode", image.size))
        if encode_fails and image.width == 32:
            raise RuntimeError("encode failed")
        return torch.ones(image.width // 16, 4)

    pipeline.video_vae.encode_image.side_effect = encode_image
    monkeypatch.setattr(module, "sequential_offload_component", record_component)
    images = (torch.zeros(16, 16, 3, dtype=torch.uint8), torch.zeros(16, 32, 3, dtype=torch.uint8))

    if encode_fails:
        with pytest.raises(RuntimeError, match="encode failed"):
            _encode_media(pipeline, images=images)
    else:
        result = _encode_media(pipeline, images=images)
        assert result.visual_condition.shape == (3, 4)
        assert result.visual_condition_shapes == ((1, 1, 1), (1, 1, 2))

    assert events == [
        ("activate", pipeline.video_vae),
        ("encode", (16, 16)),
        ("encode", (32, 16)),
        ("offload", pipeline.video_vae),
    ]


def test_h3_model_cpu_offload_shares_image_and_video_scope(monkeypatch):
    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module

    pipeline = object.__new__(MiniMaxH3Pipeline)
    torch.nn.Module.__init__(pipeline)
    pipeline.video_vae = Mock()
    pipeline._model_cpu_offload_modules = [pipeline.video_vae]
    events = []

    @contextmanager
    def record_component(value):
        events.append(("activate", value))
        try:
            yield
        finally:
            events.append(("offload", value))

    def encode_image(image):
        events.append(("image", image.size))
        return torch.ones(1, 4)

    def encode_video(frames):
        events.append(("video", frames.shape))
        return torch.ones(3, 4), (3, 2, 2)

    pipeline.video_vae.encode_image.side_effect = encode_image
    pipeline.video_vae.encode_video.side_effect = encode_video
    monkeypatch.setattr(module, "sequential_offload_component", record_component)
    result = _encode_media(
        pipeline,
        images=(torch.zeros(16, 16, 3, dtype=torch.uint8), torch.zeros(16, 32, 3, dtype=torch.uint8)),
        videos=(torch.zeros(12, 32, 32, 3, dtype=torch.uint8),),
        video_audios=(None,),
    )

    assert result.visual_condition.shape == (5, 4)
    assert result.visual_condition_shapes == ((1, 1, 1), (1, 1, 2), (3, 2, 2))
    assert events == [
        ("activate", pipeline.video_vae),
        ("image", (16, 16)),
        ("image", (32, 16)),
        ("video", (12, 32, 32, 3)),
        ("offload", pipeline.video_vae),
    ]


@pytest.mark.parametrize("encode_fails", [False, True])
def test_h3_model_cpu_offload_shares_embedded_and_standalone_audio_scope(monkeypatch, encode_fails):
    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module

    pipeline = object.__new__(MiniMaxH3Pipeline)
    torch.nn.Module.__init__(pipeline)
    pipeline.audio_vae = Mock()
    pipeline._model_cpu_offload_modules = [pipeline.audio_vae]
    events: list[tuple] = []

    @contextmanager
    def record_component(value):
        events.append(("activate", value))
        try:
            yield
        finally:
            events.append(("offload", value))

    def encode_waveform(waveform, sample_rate):
        events.append(("audio", int(waveform[0]), sample_rate))
        if encode_fails:
            raise RuntimeError("audio encode failed")
        return torch.ones(80, 32), 80

    pipeline.audio_vae.encode_waveform.side_effect = encode_waveform
    monkeypatch.setattr(module, "sequential_offload_component", record_component)

    def call():
        return _encode_media(
            pipeline,
            video_audios=(),
            audios=((torch.ones(32_000), 16_000),),
        )

    if encode_fails:
        with pytest.raises(RuntimeError, match="audio encode failed"):
            call()
    else:
        result = call()
        assert result.audio_condition.shape == (80, 32)
        assert result.audio_condition_lengths == (80,)
        assert events[1:-1] == [("audio", 1, 16_000)]

    assert events[0] == ("activate", pipeline.audio_vae)
    assert events[-1] == ("offload", pipeline.audio_vae)
    assert events.count(("activate", pipeline.audio_vae)) == 1
    assert events.count(("offload", pipeline.audio_vae)) == 1


@pytest.mark.parametrize("distributed", [False, True])
def test_manual_component_offload_wins_over_requested_model_offload(distributed):
    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline

    pipeline = object.__new__(MiniMaxH3Pipeline)
    torch.nn.Module.__init__(pipeline)
    pipeline.od_config = SimpleNamespace(
        enable_cpu_offload=True,
        enable_layerwise_offload=not distributed,
        enable_distributed_layerwise_offload=distributed,
    )
    pipeline._model_cpu_offload_modules = []
    pipeline.text_encoder = Mock()
    expected = torch.ones(2, 3)
    pipeline.text_encoder.encode_ids.return_value = expected
    component = Mock()

    actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
    with pipeline._component_on_device(component):
        component.load_to_device.assert_called_once_with()
        component.offload_to_cpu.assert_not_called()

    assert actual is expected
    pipeline.text_encoder.load_to_device.assert_called_once_with()
    pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
    component.offload_to_cpu.assert_called_once_with()


def test_requested_model_offload_without_backend_uses_resident_components():
    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline

    pipeline = object.__new__(MiniMaxH3Pipeline)
    torch.nn.Module.__init__(pipeline)
    pipeline.od_config = SimpleNamespace(
        enable_cpu_offload=True,
        enable_layerwise_offload=False,
        enable_distributed_layerwise_offload=False,
    )
    pipeline._model_cpu_offload_modules = []
    pipeline.text_encoder = Mock()
    expected = torch.ones(2, 3)
    pipeline.text_encoder.encode_ids.return_value = expected

    actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})

    assert actual is expected
    pipeline.text_encoder.load_to_device.assert_called_once_with()
    pipeline.text_encoder.encode_ids.assert_called_once()