[feat] Extend disaggregated MiniMax-H3 stage with VAE encoders (#6939)
Signed-off-by: asukaqaq-s <1311722138@qq.com> Signed-off-by: Jader <yjader@foxmail.com> Co-authored-by: Jader <yjader@foxmail.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
27 files changed
+3663
-2399
No files matched your search
@@ -0,0 +1,406 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from contextlib import contextmanager
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.data_entry_keys import flatten_payload
|
||||
from vllm_omni.model_executor.models.minimax_h3.conditioning import (
|
||||
MiniMaxH3EncoderConditioning,
|
||||
MiniMaxH3EncoderMediaConditioning,
|
||||
MiniMaxH3EncoderMediaInput,
|
||||
)
|
||||
from vllm_omni.model_executor.stage_input_processors.minimax_h3 import encoder2diffusion
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_pipeline_owns_three_encoders_then_dit_and_decoders() -> None:
|
||||
from vllm_omni.model_executor.models.minimax_h3.pipeline import MINIMAX_H3_PIPELINE
|
||||
|
||||
encoder, diffusion = MINIMAX_H3_PIPELINE.stages
|
||||
assert encoder.model_arch == "MiniMaxH3Encoder"
|
||||
assert encoder.custom_process_next_stage_input_func is None
|
||||
assert diffusion.model_arch == "MiniMaxH3Pipeline"
|
||||
assert diffusion.requires_multimodal_data is False
|
||||
assert diffusion.requires_full_payload_input is False
|
||||
|
||||
|
||||
def _patch_encoder_constructors(monkeypatch, *, rank: int):
|
||||
from vllm_omni.model_executor.models.minimax_h3 import encoder as encoder_module
|
||||
|
||||
group = SimpleNamespace(rank_in_group=rank, world_size=2, device_group=object())
|
||||
|
||||
class FakeVideoVAE(torch.nn.Module):
|
||||
def __init__(self, path, *, device, encode_only, trust_remote_code):
|
||||
super().__init__()
|
||||
self.path = path
|
||||
self.parallel_args = None
|
||||
self.trust_remote_code = trust_remote_code
|
||||
|
||||
def set_parallel_size(self, size, *, process_group):
|
||||
self.parallel_args = (size, process_group)
|
||||
|
||||
class FakeAudioVAE(torch.nn.Module):
|
||||
def __init__(self, path, *, device, encode_only, trust_remote_code):
|
||||
super().__init__()
|
||||
self.path = path
|
||||
self.trust_remote_code = trust_remote_code
|
||||
|
||||
def init_backbone(self, *, vllm_config, prefix):
|
||||
torch.nn.Module.__init__(self)
|
||||
|
||||
monkeypatch.setattr(encoder_module.MiniMaxH3TextEncoderBackbone, "__init__", init_backbone)
|
||||
monkeypatch.setattr(encoder_module, "get_tp_group", lambda: group)
|
||||
monkeypatch.setattr(encoder_module, "get_local_device", lambda: torch.device("cpu"))
|
||||
monkeypatch.setattr(encoder_module, "MiniMaxH3VideoVAE", FakeVideoVAE)
|
||||
monkeypatch.setattr(encoder_module, "MiniMaxH3AudioVAE", FakeAudioVAE)
|
||||
return encoder_module, group
|
||||
|
||||
|
||||
def _component_root(tmp_path):
|
||||
root = tmp_path / "FL2VA"
|
||||
for name in ("text_encoder", "video_vae", "audio_vae"):
|
||||
(root / name).mkdir(parents=True)
|
||||
return root
|
||||
|
||||
|
||||
def _encoder_config(root, *, video_mode="patch", roles=None, trust_remote_code=True):
|
||||
components = (
|
||||
{
|
||||
"text_encoder": {"parallel_mode": "tp"},
|
||||
"video_vae": {"parallel_mode": video_mode},
|
||||
"audio_vae": {"parallel_mode": "leader"},
|
||||
}
|
||||
if roles is None
|
||||
else roles
|
||||
)
|
||||
return SimpleNamespace(
|
||||
model_config=SimpleNamespace(
|
||||
model=str(root / "text_encoder"),
|
||||
hf_config=SimpleNamespace(minimax_h3_encoder_components=components),
|
||||
trust_remote_code=trust_remote_code,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_encoder_requires_exactly_three_components(monkeypatch, tmp_path) -> None:
|
||||
encoder_module, _ = _patch_encoder_constructors(monkeypatch, rank=0)
|
||||
root = tmp_path / "FL2VA"
|
||||
(root / "text_encoder").mkdir(parents=True)
|
||||
(root / "video_vae").mkdir()
|
||||
config = _encoder_config(root)
|
||||
|
||||
with pytest.raises(RuntimeError, match=r"requires all three encoders.*audio_vae"):
|
||||
encoder_module.MiniMaxH3Encoder(vllm_config=config)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("rank", "has_audio_vae"), [(0, True), (1, False)])
|
||||
@pytest.mark.parametrize("trust_remote_code", [False, True])
|
||||
def test_encoder_uses_tp_patch_and_leader_on_the_same_rank_set(
|
||||
monkeypatch,
|
||||
tmp_path,
|
||||
rank,
|
||||
has_audio_vae,
|
||||
trust_remote_code,
|
||||
) -> None:
|
||||
encoder_module, group = _patch_encoder_constructors(monkeypatch, rank=rank)
|
||||
root = _component_root(tmp_path)
|
||||
config = _encoder_config(root, trust_remote_code=trust_remote_code)
|
||||
|
||||
model = encoder_module.MiniMaxH3Encoder(vllm_config=config)
|
||||
|
||||
assert model.component_config.text_parallel_mode == "tp"
|
||||
assert model.component_config.video_parallel_mode == "patch"
|
||||
assert model.component_config.audio_parallel_mode == "leader"
|
||||
assert model.video_vae.parallel_args == (2, group.device_group)
|
||||
assert model.video_vae.trust_remote_code is trust_remote_code
|
||||
assert (model.audio_vae is not None) is has_audio_vae
|
||||
if has_audio_vae:
|
||||
assert model.audio_vae.trust_remote_code is trust_remote_code
|
||||
|
||||
|
||||
@pytest.mark.parametrize("component", ["video", "audio"])
|
||||
@pytest.mark.parametrize("trusted", [False, True])
|
||||
def test_encoder_only_vae_requires_remote_code_trust(monkeypatch, component, trusted) -> None:
|
||||
from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
|
||||
|
||||
config = {"auto_map": {"AutoModel": "modeling.RemoteVAE"}, "sample_rate": 32000}
|
||||
monkeypatch.setattr(vae_module, "_load_component_config", lambda _path: config)
|
||||
remote = torch.nn.Module()
|
||||
remote.model = torch.nn.Module()
|
||||
calls = []
|
||||
|
||||
def load_encoder(path, config_dict):
|
||||
calls.append((path, config_dict))
|
||||
return remote
|
||||
|
||||
monkeypatch.setattr(vae_module, f"_load_{component}_vae_encoder", load_encoder)
|
||||
vae_cls = vae_module.MiniMaxH3VideoVAE if component == "video" else vae_module.MiniMaxH3AudioVAE
|
||||
if trusted:
|
||||
model = vae_cls("unused", device=torch.device("cpu"), encode_only=True, trust_remote_code=True)
|
||||
assert model.remote is remote
|
||||
assert calls == [("unused", config)]
|
||||
else:
|
||||
with pytest.raises(ValueError, match="trust-remote-code"):
|
||||
vae_cls("unused", device=torch.device("cpu"), encode_only=True)
|
||||
assert calls == []
|
||||
|
||||
|
||||
def test_encoder_requires_all_role_policies(monkeypatch, tmp_path) -> None:
|
||||
encoder_module, _ = _patch_encoder_constructors(monkeypatch, rank=0)
|
||||
root = _component_root(tmp_path)
|
||||
config = _encoder_config(
|
||||
root,
|
||||
roles={
|
||||
"text_encoder": {"parallel_mode": "tp"},
|
||||
"video_vae": {"parallel_mode": "patch"},
|
||||
},
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match=r"requires exactly.*audio_vae"):
|
||||
encoder_module.MiniMaxH3Encoder(vllm_config=config)
|
||||
|
||||
|
||||
def test_encoder_does_not_infer_role_policies(monkeypatch, tmp_path) -> None:
|
||||
encoder_module, _ = _patch_encoder_constructors(monkeypatch, rank=0)
|
||||
root = _component_root(tmp_path)
|
||||
config = SimpleNamespace(
|
||||
model_config=SimpleNamespace(
|
||||
model=str(root / "text_encoder"),
|
||||
hf_config=SimpleNamespace(),
|
||||
)
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="requires explicit minimax_h3_encoder_components"):
|
||||
encoder_module.MiniMaxH3Encoder(vllm_config=config)
|
||||
|
||||
|
||||
def test_video_leader_policy_skips_nonleader(monkeypatch, tmp_path) -> None:
|
||||
encoder_module, _ = _patch_encoder_constructors(monkeypatch, rank=1)
|
||||
root = _component_root(tmp_path)
|
||||
model = encoder_module.MiniMaxH3Encoder(vllm_config=_encoder_config(root, video_mode="leader"))
|
||||
|
||||
assert model.video_vae is None
|
||||
assert model.audio_vae is None
|
||||
assert (
|
||||
model._encode_media(
|
||||
MiniMaxH3EncoderMediaInput(
|
||||
task="ref2va",
|
||||
height=32,
|
||||
width=32,
|
||||
num_frames=48,
|
||||
latent_t=1,
|
||||
audio_t=80,
|
||||
images=(torch.zeros(32, 32, 3, dtype=torch.uint8),),
|
||||
)
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
def test_encoder_runs_video_and_audio_components_on_ar_model() -> None:
|
||||
from vllm_omni.model_executor.models.minimax_h3.encoder import MiniMaxH3Encoder
|
||||
|
||||
class VideoVAE:
|
||||
def encode_image(self, _image):
|
||||
return torch.ones(1, 96)
|
||||
|
||||
def encode_video(self, _frames):
|
||||
return torch.full((1, 96), 2.0), (1, 2, 2)
|
||||
|
||||
class AudioVAE:
|
||||
def encode_waveform(self, waveform, _sample_rate):
|
||||
return torch.ones(160, 32), waveform.shape[-1] // 800
|
||||
|
||||
model = MiniMaxH3Encoder.__new__(MiniMaxH3Encoder)
|
||||
torch.nn.Module.__init__(model)
|
||||
model._component_leader = True
|
||||
model.video_vae = VideoVAE()
|
||||
model.audio_vae = AudioVAE()
|
||||
media = MiniMaxH3EncoderMediaInput(
|
||||
task="ref2va",
|
||||
height=32,
|
||||
width=32,
|
||||
num_frames=48,
|
||||
latent_t=1,
|
||||
audio_t=80,
|
||||
images=(torch.zeros(32, 32, 3, dtype=torch.uint8),),
|
||||
videos=(torch.zeros(1, 32, 32, 3, dtype=torch.uint8),),
|
||||
video_audios=((torch.zeros(64_000), 32_000),),
|
||||
)
|
||||
|
||||
conditioning = model._encode_media(media)
|
||||
|
||||
assert conditioning is not None
|
||||
assert conditioning.visual_condition_shapes == ((1, 2, 2), (1, 2, 2))
|
||||
assert conditioning.audio_condition_lengths == (80,)
|
||||
assert conditioning.ref_blocks[1]["kind"] == "video_audio"
|
||||
|
||||
|
||||
def test_prepare_encoder_inputs_keeps_reference_audio_budgets_separate(monkeypatch) -> None:
|
||||
from vllm_omni.model_executor.models.minimax_h3 import encoder_processing as processing
|
||||
|
||||
frames = torch.zeros(1, 32, 32, 3, dtype=torch.uint8).numpy()
|
||||
waveform = torch.zeros(320_000)
|
||||
monkeypatch.setattr(processing, "resolve_minimax_h3_shape", lambda *_args: (32, 32, 240, 60, 400))
|
||||
monkeypatch.setattr(
|
||||
processing,
|
||||
"prepare_reference_videos",
|
||||
lambda *_args, **_kwargs: [
|
||||
{"prepared_path": "prepared.mp4", "original_path": "reference.mp4", "input_has_audio": True}
|
||||
],
|
||||
)
|
||||
monkeypatch.setattr(processing, "load_video_frames", lambda _path: frames)
|
||||
monkeypatch.setattr(
|
||||
processing,
|
||||
"sample_reference_video_frames",
|
||||
lambda *_args, **_kwargs: {"frames": [frames[0]], "block_timestamps": [[0.0]]},
|
||||
)
|
||||
monkeypatch.setattr(processing, "load_video_audio", lambda *_args, **_kwargs: (waveform, 32_000))
|
||||
|
||||
prepared = processing.prepare_encoder_inputs(
|
||||
{
|
||||
"prompt": "reference",
|
||||
"multi_modal_data": {"video": "reference.mp4", "audio": (waveform, 32_000)},
|
||||
},
|
||||
SimpleNamespace(extra_args={"task": "ref2va"}),
|
||||
)
|
||||
|
||||
assert prepared.media.video_audios[0][0].shape[-1] == 320_000
|
||||
assert prepared.media.audios[0][0].shape[-1] == 320_000
|
||||
assert prepared.condition_labels == [("audio", 1), ("video", 1), ("audio", 2)]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("embedded_lengths", "standalone_lengths", "valid"),
|
||||
[((400,), (400,), True), ((320, 320), (400,), False), ((400,), (320, 320), False)],
|
||||
)
|
||||
def test_encode_media_keeps_audio_budgets_and_component_residency_separate(
|
||||
embedded_lengths, standalone_lengths, valid
|
||||
) -> None:
|
||||
from vllm_omni.model_executor.models.minimax_h3.encoder_processing import encode_media
|
||||
|
||||
active = None
|
||||
scopes = []
|
||||
|
||||
class VideoVAE:
|
||||
def encode_image(self, _image):
|
||||
assert active is self
|
||||
return torch.ones(1, 96)
|
||||
|
||||
def encode_video(self, _frames):
|
||||
assert active is self
|
||||
return torch.ones(1, 96), (1, 2, 2)
|
||||
|
||||
class AudioVAE:
|
||||
def encode_waveform(self, waveform, _sample_rate):
|
||||
assert active is self
|
||||
length = waveform.shape[-1] // 800
|
||||
return torch.ones(2 * length, 32), length
|
||||
|
||||
@contextmanager
|
||||
def component_scope(component):
|
||||
nonlocal active
|
||||
assert active is None
|
||||
active = component
|
||||
scopes.append(component)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
active = None
|
||||
|
||||
video_vae = VideoVAE()
|
||||
audio_vae = AudioVAE()
|
||||
media = MiniMaxH3EncoderMediaInput(
|
||||
task="ref2va",
|
||||
height=32,
|
||||
width=32,
|
||||
num_frames=360,
|
||||
latent_t=90,
|
||||
audio_t=600,
|
||||
images=(torch.zeros(32, 32, 3, dtype=torch.uint8),),
|
||||
videos=tuple(torch.zeros(1, 32, 32, 3, dtype=torch.uint8) for _ in embedded_lengths),
|
||||
video_audios=tuple((torch.zeros(length * 800), 32_000) for length in embedded_lengths),
|
||||
audios=tuple((torch.zeros(length * 800), 32_000) for length in standalone_lengths),
|
||||
)
|
||||
kwargs = dict(video_vae=video_vae, audio_vae=audio_vae, emit_conditioning=True, component_scope=component_scope)
|
||||
if valid:
|
||||
conditioning = encode_media(media, **kwargs)
|
||||
assert conditioning.audio_condition_lengths == embedded_lengths + standalone_lengths
|
||||
assert [block["kind"] for block in conditioning.ref_blocks] == ["image", "video_audio", "audio"]
|
||||
conditioning.to_omni_components()
|
||||
else:
|
||||
with pytest.raises(ValueError, match="at most 15 seconds"):
|
||||
encode_media(media, **kwargs)
|
||||
assert scopes == [video_vae, audio_vae]
|
||||
|
||||
|
||||
def test_encoder_output_reuses_encoder_handoff_and_round_trips() -> None:
|
||||
expected = MiniMaxH3EncoderConditioning(
|
||||
hidden_states=torch.randn(3, 5120, dtype=torch.bfloat16),
|
||||
token_tags=torch.tensor([1, 0, 1], dtype=torch.int64),
|
||||
task="t2va",
|
||||
height=256,
|
||||
width=448,
|
||||
num_frames=17,
|
||||
latent_t=5,
|
||||
audio_t=10,
|
||||
)
|
||||
|
||||
source = SimpleNamespace(
|
||||
finished=True,
|
||||
request_id="request-1",
|
||||
outputs=[SimpleNamespace(multimodal_output=flatten_payload(expected.to_omni_payload()))],
|
||||
)
|
||||
prompt = {
|
||||
"prompt": "hello",
|
||||
"additional_information": {"global_request_id": ["request-1"]},
|
||||
}
|
||||
|
||||
result = encoder2diffusion([source], prompt)
|
||||
|
||||
assert result is not None
|
||||
payload = result["additional_information"]["encoder_output"]
|
||||
actual = MiniMaxH3EncoderConditioning.from_omni_payload(payload)
|
||||
torch.testing.assert_close(actual.hidden_states, expected.hidden_states)
|
||||
torch.testing.assert_close(actual.token_tags, expected.token_tags)
|
||||
|
||||
|
||||
def test_encoder_releases_workspace_after_cpu_payload(monkeypatch) -> None:
|
||||
from vllm_omni.model_executor.models.minimax_h3 import encoder as encoder_module
|
||||
|
||||
model = encoder_module.MiniMaxH3Encoder.__new__(encoder_module.MiniMaxH3Encoder)
|
||||
torch.nn.Module.__init__(model)
|
||||
model._component_leader = True
|
||||
model._token_tags = torch.tensor([1])
|
||||
info = {
|
||||
"meta": {"minimax_h3_encoder_request": {}},
|
||||
"_minimax_h3_encoder_media_conditioning": MiniMaxH3EncoderMediaConditioning(
|
||||
task="t2va",
|
||||
height=256,
|
||||
width=448,
|
||||
num_frames=17,
|
||||
latent_t=5,
|
||||
audio_t=10,
|
||||
),
|
||||
}
|
||||
empty_cache_calls = 0
|
||||
|
||||
def empty_cache() -> None:
|
||||
nonlocal empty_cache_calls
|
||||
empty_cache_calls += 1
|
||||
|
||||
monkeypatch.setattr(torch.accelerator, "empty_cache", empty_cache)
|
||||
output = model.make_omni_output(
|
||||
torch.zeros(1, 5120),
|
||||
model_intermediate_buffer=[info],
|
||||
request_sample_eligible=[True],
|
||||
)
|
||||
|
||||
assert empty_cache_calls == 1
|
||||
assert output.multimodal_outputs is not None
|
||||
assert output.multimodal_outputs["embed"]["embedding"][0].device.type == "cpu"
|
||||
@@ -2,8 +2,6 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Regression tests for MiniMax H3's disaggregated text-encoder contract."""
|
||||
|
||||
import os
|
||||
import shutil
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
@@ -12,39 +10,37 @@ import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from vllm_omni.data_entry_keys import OmniPayload, deserialize_payload, serialize_payload
|
||||
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
|
||||
_load_audio,
|
||||
resolve_minimax_h3_diffusion_model_path,
|
||||
from vllm_omni.data_entry_keys import flatten_payload
|
||||
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import resolve_minimax_h3_diffusion_model_path
|
||||
from vllm_omni.engine.serialization import (
|
||||
deserialize_additional_information,
|
||||
serialize_additional_information,
|
||||
)
|
||||
from vllm_omni.errors import OmniClientError
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.model_executor.models.minimax_h3.checkpoint import (
|
||||
resolve_minimax_h3_model_root,
|
||||
resolve_minimax_h3_encoder_model_root,
|
||||
resolve_minimax_h3_partition,
|
||||
)
|
||||
from vllm_omni.model_executor.models.minimax_h3.conditioning import (
|
||||
MINIMAX_H3_CONDITION_LABELS_KEY,
|
||||
MINIMAX_H3_ENCODER_REQUEST_KEY,
|
||||
MINIMAX_H3_PRESENTATION_TASK_KEY,
|
||||
MINIMAX_H3_TEXT_CONDITIONING_SCHEMA,
|
||||
MiniMaxH3TextConditioning,
|
||||
MiniMaxH3EncoderConditioning,
|
||||
)
|
||||
from vllm_omni.model_executor.models.minimax_h3.encoder_processing import _audio_items, _load_audio
|
||||
from vllm_omni.model_executor.models.minimax_h3.preprocessing import (
|
||||
minimax_h3_ref2va_presentation,
|
||||
minimax_h3_ref2va_video_presentation,
|
||||
)
|
||||
from vllm_omni.model_executor.models.minimax_h3.reference_video import (
|
||||
MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY,
|
||||
deserialize_prepared_reference_videos,
|
||||
)
|
||||
from vllm_omni.model_executor.models.minimax_h3.text_encoder import (
|
||||
MiniMaxH3MultiModalProcessor,
|
||||
_build_minimax_h3_presentation,
|
||||
)
|
||||
from vllm_omni.model_executor.stage_input_processors.minimax_h3 import (
|
||||
_audio_items,
|
||||
_diffusion_sampling_params,
|
||||
prepare_text_encoder_prompt,
|
||||
text_encoder2diffusion,
|
||||
encoder2diffusion,
|
||||
prepare_encoder_prompt,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
@@ -66,13 +62,14 @@ class _SegmentTokenizer:
|
||||
return self._special_ids[token]
|
||||
|
||||
|
||||
def test_h3_processor_reprocesses_media_instead_of_using_partial_sender_cache(mocker):
|
||||
def test_h3_processor_reprocesses_media_instead_of_using_partial_sender_cache(monkeypatch):
|
||||
processor = object.__new__(MiniMaxH3MultiModalProcessor)
|
||||
sentinel = ([1, 2, 3], object(), True)
|
||||
apply_processor = mocker.patch.object(
|
||||
apply_processor = Mock(return_value=sentinel)
|
||||
monkeypatch.setattr(
|
||||
MiniMaxH3MultiModalProcessor,
|
||||
"_apply_hf_processor",
|
||||
return_value=sentinel,
|
||||
apply_processor,
|
||||
)
|
||||
inputs = object()
|
||||
timing_ctx = object()
|
||||
@@ -119,9 +116,11 @@ def test_fused_audio_loader_accepts_list_waveform_pair():
|
||||
def test_prepare_ref2va_keeps_original_text_and_exact_condition_order():
|
||||
prompt = {
|
||||
"prompt": "hello",
|
||||
"additional_information": {"global_request_id": ["request-1"]},
|
||||
"model_intermediate_buffer": {"private": "preserved"},
|
||||
"multi_modal_data": {
|
||||
"image": Image.new("RGB", (256, 256)),
|
||||
"audio": [np.zeros(16), 16_000],
|
||||
"audio": [np.zeros(32_000), 16_000],
|
||||
},
|
||||
}
|
||||
sampling = OmniDiffusionSamplingParams(
|
||||
@@ -130,7 +129,7 @@ def test_prepare_ref2va_keeps_original_text_and_exact_condition_order():
|
||||
extra_args={"task": "ref2va"},
|
||||
)
|
||||
|
||||
transformed = prepare_text_encoder_prompt(prompt, [sampling])
|
||||
transformed = prepare_encoder_prompt(prompt, [sampling])
|
||||
|
||||
assert transformed["prompt"] == "hello"
|
||||
assert len(transformed["multi_modal_data"]["image"]) == 1
|
||||
@@ -140,223 +139,148 @@ def test_prepare_ref2va_keeps_original_text_and_exact_condition_order():
|
||||
("image", 1),
|
||||
("audio", 1),
|
||||
]
|
||||
assert transformed["model_intermediate_buffer"] == {"private": "preserved"}
|
||||
runner_info = transformed["additional_information"]
|
||||
for _ in range(2):
|
||||
wire = serialize_additional_information(runner_info)
|
||||
runner_info = deserialize_additional_information(wire)
|
||||
assert runner_info["global_request_id"] == ["request-1"]
|
||||
request_metadata = runner_info["meta"][MINIMAX_H3_ENCODER_REQUEST_KEY]
|
||||
assert request_metadata["task"] == "ref2va"
|
||||
assert isinstance(
|
||||
runner_info["hidden_states"]["layers"][0],
|
||||
torch.Tensor,
|
||||
)
|
||||
assert isinstance(
|
||||
runner_info["hidden_states"]["layers"][1],
|
||||
torch.Tensor,
|
||||
)
|
||||
from vllm_omni.model_executor.models.minimax_h3.encoder import MiniMaxH3Encoder
|
||||
|
||||
media = MiniMaxH3Encoder._media_input(runner_info)
|
||||
assert media.task == "ref2va"
|
||||
assert len(media.images) == 1
|
||||
assert len(media.audios) == 1
|
||||
|
||||
|
||||
def test_text_encoder_prompt_rejects_injected_prepared_video_descriptor():
|
||||
prompt = {
|
||||
"prompt": "hello",
|
||||
"additional_information": {"meta": {MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY: '{"artifact_dir":"/tmp/user"}'}},
|
||||
}
|
||||
sampling = OmniDiffusionSamplingParams(height=256, width=448, extra_args={"task": "t2va"})
|
||||
def _mock_ref2va_video_with_audio(monkeypatch, *, duration_seconds: float) -> None:
|
||||
from vllm_omni.model_executor.models.minimax_h3 import encoder_processing
|
||||
|
||||
transformed = prepare_text_encoder_prompt(prompt, [sampling])
|
||||
|
||||
assert MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY not in transformed["additional_information"]["meta"]
|
||||
|
||||
|
||||
def test_prepare_ref2va_video_uses_shared_frame_sampler_once(monkeypatch):
|
||||
prepare_videos = Mock()
|
||||
|
||||
def prepare(_videos, *, target_frame_count, workdir, start_time_seconds):
|
||||
prepare_videos(_videos, target_frame_count, workdir, start_time_seconds)
|
||||
prepared_path = os.path.join(workdir, "prepared.mp4")
|
||||
open(prepared_path, "wb").close()
|
||||
return [
|
||||
{
|
||||
"original_path": "/tmp/input.mp4",
|
||||
"prepared_path": prepared_path,
|
||||
"input_has_audio": True,
|
||||
"width": 448,
|
||||
"height": 256,
|
||||
"start_time_seconds": 0.0,
|
||||
"duration_seconds": 5.2,
|
||||
"audio_duration_seconds": 5.2,
|
||||
}
|
||||
]
|
||||
|
||||
sample_frames = Mock(return_value={"frames": [np.zeros((4, 4, 3), dtype=np.uint8)]})
|
||||
monkeypatch.setattr("vllm_omni.model_executor.stage_input_processors.minimax_h3.prepare_reference_videos", prepare)
|
||||
sample_rate = 16_000
|
||||
frames = np.zeros((4, 32, 32, 3), dtype=np.uint8)
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.model_executor.stage_input_processors.minimax_h3.sample_reference_video_frames",
|
||||
sample_frames,
|
||||
encoder_processing,
|
||||
"prepare_reference_videos",
|
||||
lambda *args, **kwargs: [
|
||||
{
|
||||
"prepared_path": "prepared.mp4",
|
||||
"original_path": "original.mp4",
|
||||
"input_has_audio": True,
|
||||
"duration_seconds": duration_seconds,
|
||||
"audio_duration_seconds": duration_seconds,
|
||||
}
|
||||
],
|
||||
)
|
||||
monkeypatch.setattr(encoder_processing, "load_video_frames", lambda path: frames)
|
||||
monkeypatch.setattr(
|
||||
encoder_processing,
|
||||
"sample_reference_video_frames",
|
||||
lambda *args, **kwargs: {"block_timestamps": [0.0], "frames": frames[:1]},
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
encoder_processing,
|
||||
"load_video_audio",
|
||||
lambda *args, **kwargs: (torch.zeros(round(duration_seconds * sample_rate)), sample_rate),
|
||||
)
|
||||
|
||||
|
||||
def test_prepare_ref2va_rejects_short_embedded_video_audio(monkeypatch):
|
||||
_mock_ref2va_video_with_audio(monkeypatch, duration_seconds=1.0)
|
||||
sampling = OmniDiffusionSamplingParams(
|
||||
height=256,
|
||||
width=448,
|
||||
num_frames=96,
|
||||
extra_args={"task": "ref2va"},
|
||||
)
|
||||
prompt = {"prompt": "hello", "multi_modal_data": {"video": "original.mp4"}}
|
||||
|
||||
with pytest.raises(OmniClientError, match=r"duration must be in \[2, 15\] seconds, got 1\.000"):
|
||||
prepare_encoder_prompt(prompt, [sampling])
|
||||
|
||||
|
||||
def test_prepare_ref2va_keeps_embedded_and_standalone_audio_budgets_separate(monkeypatch):
|
||||
duration_seconds = 8.0
|
||||
sample_rate = 16_000
|
||||
_mock_ref2va_video_with_audio(monkeypatch, duration_seconds=duration_seconds)
|
||||
sampling = OmniDiffusionSamplingParams(
|
||||
height=256,
|
||||
width=448,
|
||||
num_frames=240,
|
||||
extra_args={"task": "ref2va"},
|
||||
)
|
||||
prompt = {
|
||||
"prompt": "hello",
|
||||
"multi_modal_data": {"video": "/tmp/input.mp4"},
|
||||
"multi_modal_data": {
|
||||
"video": "original.mp4",
|
||||
"audio": (torch.zeros(int(duration_seconds * sample_rate)), sample_rate),
|
||||
},
|
||||
}
|
||||
sampling = OmniDiffusionSamplingParams(height=256, width=448, num_frames=124, extra_args={"task": "ref2va"})
|
||||
|
||||
transformed = prepare_text_encoder_prompt(prompt, [sampling])
|
||||
|
||||
prepare_videos.assert_called_once()
|
||||
_, target_frame_count, artifact_dir, start_time_seconds = prepare_videos.call_args.args
|
||||
assert target_frame_count == 124
|
||||
assert start_time_seconds is None
|
||||
descriptor = transformed["additional_information"]["meta"][MINIMAX_H3_PREPARED_REFERENCE_VIDEOS_KEY]
|
||||
described_dir, described_videos = deserialize_prepared_reference_videos(descriptor)
|
||||
assert described_dir == artifact_dir
|
||||
assert os.path.isfile(described_videos[0]["prepared_path"])
|
||||
sample_frames.assert_called_once_with(described_videos[0]["prepared_path"])
|
||||
assert prompt["multi_modal_data"] == {"video": "/tmp/input.mp4"}
|
||||
transformed = prepare_encoder_prompt(prompt, [sampling])
|
||||
assert transformed["mm_processor_kwargs"][MINIMAX_H3_CONDITION_LABELS_KEY] == [
|
||||
("audio", 1),
|
||||
("video", 1),
|
||||
("audio", 2),
|
||||
]
|
||||
shutil.rmtree(artifact_dir)
|
||||
|
||||
|
||||
def test_stage_wire_rejects_multiple_text_encoder_sources():
|
||||
with pytest.raises(RuntimeError, match="exactly one text-encoder source"):
|
||||
text_encoder2diffusion([SimpleNamespace(), SimpleNamespace()], prompt={"prompt": "hello"})
|
||||
def _encoder_output() -> dict:
|
||||
return MiniMaxH3EncoderConditioning(
|
||||
hidden_states=torch.randn(3, 5120, dtype=torch.bfloat16),
|
||||
token_tags=torch.tensor([1, 0, 1], dtype=torch.int64),
|
||||
task="t2va",
|
||||
height=256,
|
||||
width=448,
|
||||
num_frames=17,
|
||||
latent_t=5,
|
||||
audio_t=10,
|
||||
).to_omni_payload()
|
||||
|
||||
|
||||
def test_stage_wire_rejects_request_id_mismatch():
|
||||
source = SimpleNamespace(request_id="other", outputs=[SimpleNamespace()])
|
||||
def test_encoder2diffusion_reuses_encoder_handoff() -> None:
|
||||
prompt = {
|
||||
"prompt": "hello",
|
||||
"additional_information": {"global_request_id": ["req-1"]},
|
||||
}
|
||||
|
||||
with pytest.raises(RuntimeError, match="request ID does not match"):
|
||||
text_encoder2diffusion([source], prompt=prompt)
|
||||
|
||||
|
||||
def test_stage_wire_rejects_multiple_completions():
|
||||
source = SimpleNamespace(request_id="req-1", outputs=[SimpleNamespace(), SimpleNamespace()])
|
||||
|
||||
with pytest.raises(RuntimeError, match="exactly one completion"):
|
||||
text_encoder2diffusion([source], prompt={"prompt": "hello"})
|
||||
|
||||
|
||||
def _source_output(payload: dict) -> SimpleNamespace:
|
||||
completion = SimpleNamespace(multimodal_output=payload)
|
||||
return SimpleNamespace(request_id="request-1", outputs=[completion])
|
||||
|
||||
|
||||
def _diffusion_prompt() -> dict:
|
||||
return {
|
||||
"request_id": "request-1",
|
||||
"prompt": "test prompt",
|
||||
"additional_information": {},
|
||||
}
|
||||
|
||||
|
||||
def _valid_text_conditioning_payload() -> OmniPayload:
|
||||
return {
|
||||
"hidden_states": {
|
||||
"output": torch.zeros(4, 5120, dtype=torch.bfloat16),
|
||||
},
|
||||
"meta": {
|
||||
"token_role_ids": torch.tensor([[1], [1], [0], [0]], dtype=torch.int64),
|
||||
"multi_modal_data": {"image": object()},
|
||||
"additional_information": {
|
||||
"private": "preserved",
|
||||
"meta": {MINIMAX_H3_ENCODER_REQUEST_KEY: {"task": "t2va"}},
|
||||
"hidden_states": {"layers": {0: torch.zeros(1)}},
|
||||
},
|
||||
"model_intermediate_buffer": {"private": "encoder-only"},
|
||||
}
|
||||
|
||||
|
||||
def test_text_encoder2diffusion_reads_hidden_states_output_and_token_role_ids() -> None:
|
||||
payload = _valid_text_conditioning_payload()
|
||||
hidden = payload["hidden_states"]["output"]
|
||||
token_role_ids = payload["meta"]["token_role_ids"]
|
||||
|
||||
result = text_encoder2diffusion(
|
||||
[_source_output(payload)],
|
||||
_diffusion_prompt(),
|
||||
source = SimpleNamespace(
|
||||
finished=True,
|
||||
request_id="request-1",
|
||||
outputs=[SimpleNamespace(multimodal_output=flatten_payload(_encoder_output()))],
|
||||
)
|
||||
|
||||
conditioning = result["additional_information"]["text_encoder_output"]
|
||||
torch.testing.assert_close(conditioning["hidden_states"], hidden)
|
||||
torch.testing.assert_close(conditioning["token_tags"], token_role_ids.squeeze(-1))
|
||||
result = encoder2diffusion([source], prompt)
|
||||
|
||||
assert result["prompt"] == "test prompt"
|
||||
assert result["multi_modal_data"] is None
|
||||
assert "model_intermediate_buffer" not in result
|
||||
additional_information = result["additional_information"]
|
||||
assert additional_information["private"] == "preserved"
|
||||
assert "hidden_states" not in additional_information
|
||||
assert "meta" not in additional_information
|
||||
parsed = MiniMaxH3EncoderConditioning.from_omni_payload(additional_information["encoder_output"])
|
||||
assert parsed.task == "t2va"
|
||||
|
||||
|
||||
def test_text_conditioning_artifact_round_trips_through_omni_payload_wire_format() -> None:
|
||||
payload = _valid_text_conditioning_payload()
|
||||
|
||||
wire = serialize_payload(payload)
|
||||
|
||||
assert wire is not None
|
||||
restored = deserialize_payload(wire)
|
||||
conditioning = MiniMaxH3TextConditioning.from_omni_payload(restored)
|
||||
assert MINIMAX_H3_TEXT_CONDITIONING_SCHEMA == "minimax_h3.text_conditioning/v1"
|
||||
torch.testing.assert_close(
|
||||
conditioning.hidden_states,
|
||||
payload["hidden_states"]["output"],
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
conditioning.token_tags,
|
||||
payload["meta"]["token_role_ids"].squeeze(-1),
|
||||
)
|
||||
assert conditioning.hidden_states.dtype == torch.bfloat16
|
||||
assert conditioning.hidden_states.is_contiguous()
|
||||
assert conditioning.token_tags.dtype == torch.int64
|
||||
assert conditioning.token_tags.is_contiguous()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("payload", "message"),
|
||||
[
|
||||
({"meta": {"token_role_ids": torch.ones(4, 1)}}, "no hidden_states payload"),
|
||||
({"hidden_states": {}, "meta": {"token_role_ids": torch.ones(4, 1)}}, "no hidden_states.output tensor"),
|
||||
({"hidden_states": {"output": torch.randn(4, 5120)}}, "no conditioning metadata"),
|
||||
({"hidden_states": {"output": torch.randn(4, 5120)}, "meta": {}}, "no token_role_ids tensor"),
|
||||
(
|
||||
{
|
||||
"hidden_states": {"output": torch.randn(4, 5120)},
|
||||
"meta": {"token_role_ids": torch.ones(4)},
|
||||
},
|
||||
r"must have shape \[tokens, 1\]",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_text_encoder2diffusion_rejects_invalid_conditioning_payload(payload: dict, message: str) -> None:
|
||||
with pytest.raises(RuntimeError, match=message):
|
||||
text_encoder2diffusion([_source_output(payload)], _diffusion_prompt())
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("case", "message"),
|
||||
[
|
||||
("hidden_width", r"hidden_states must have shape \[tokens, 5120\]"),
|
||||
("hidden_dtype", "hidden_states must have dtype torch.bfloat16"),
|
||||
("hidden_noncontiguous", "hidden_states must use contiguous strided layout"),
|
||||
("hidden_non_strided", "hidden_states must use contiguous strided layout"),
|
||||
("token_dtype", "stage-wire token_role_ids must have dtype torch.int64"),
|
||||
("token_layout", "stage-wire token_role_ids must use contiguous strided layout"),
|
||||
("token_count", "token_tags must align with hidden_states"),
|
||||
("token_value", "token_tags must contain only 0 and 1"),
|
||||
],
|
||||
)
|
||||
def test_text_encoder2diffusion_rejects_h3_v1_contract_mismatch(case: str, message: str) -> None:
|
||||
payload = _valid_text_conditioning_payload()
|
||||
|
||||
if case == "hidden_width":
|
||||
payload["hidden_states"]["output"] = torch.zeros(4, 5119, dtype=torch.bfloat16)
|
||||
elif case == "hidden_dtype":
|
||||
payload["hidden_states"]["output"] = torch.zeros(4, 5120, dtype=torch.float32)
|
||||
elif case == "hidden_noncontiguous":
|
||||
payload["hidden_states"]["output"] = torch.empty(5120, 4, dtype=torch.bfloat16).t()
|
||||
elif case == "hidden_non_strided":
|
||||
payload["hidden_states"]["output"] = torch.empty(
|
||||
(4, 5120),
|
||||
dtype=torch.bfloat16,
|
||||
layout=torch.sparse_coo,
|
||||
)
|
||||
elif case == "token_dtype":
|
||||
payload["meta"]["token_role_ids"] = torch.tensor([[1], [1], [0], [0]], dtype=torch.int32)
|
||||
elif case == "token_layout":
|
||||
payload["meta"]["token_role_ids"] = torch.tensor(
|
||||
[[1], [0], [1], [0], [1], [0], [1], [0]],
|
||||
dtype=torch.int64,
|
||||
)[::2]
|
||||
elif case == "token_count":
|
||||
payload["meta"]["token_role_ids"] = torch.tensor([[1], [0], [0]], dtype=torch.int64)
|
||||
elif case == "token_value":
|
||||
payload["meta"]["token_role_ids"] = torch.tensor([[1], [1], [2], [0]], dtype=torch.int64)
|
||||
else: # pragma: no cover - parameterization is exhaustive
|
||||
raise AssertionError(f"unknown test case: {case}")
|
||||
|
||||
with pytest.raises(RuntimeError, match=message) as exc_info:
|
||||
text_encoder2diffusion([_source_output(payload)], _diffusion_prompt())
|
||||
|
||||
assert MINIMAX_H3_TEXT_CONDITIONING_SCHEMA in str(exc_info.value)
|
||||
def test_encoder2diffusion_waits_for_one_finished_source() -> None:
|
||||
assert encoder2diffusion([SimpleNamespace(finished=False)], {"prompt": "hello"}) is None
|
||||
with pytest.raises(RuntimeError, match="exactly one encoder source"):
|
||||
encoder2diffusion([SimpleNamespace(), SimpleNamespace()], {"prompt": "hello"})
|
||||
|
||||
|
||||
def test_ref2va_one_image_tokens_and_tags_match_fused_presentation():
|
||||
@@ -421,15 +345,20 @@ def test_checkpoint_resolver_selects_local_partition(tmp_path):
|
||||
(root / "FL2VA" / "text_encoder").mkdir(parents=True)
|
||||
(root / "Ref2VA" / "text_encoder").mkdir(parents=True)
|
||||
|
||||
assert resolve_minimax_h3_model_root(str(root), None, "fl2va") == str(root / "FL2VA" / "text_encoder")
|
||||
assert resolve_minimax_h3_model_root(str(root), None, "ref2va") == str(root / "Ref2VA" / "text_encoder")
|
||||
assert resolve_minimax_h3_model_root(str(root / "Ref2VA"), None, None) == str(root / "Ref2VA" / "text_encoder")
|
||||
assert resolve_minimax_h3_model_root(str(root / "FL2VA"), None, "ref2va") == str(root / "Ref2VA" / "text_encoder")
|
||||
assert resolve_minimax_h3_encoder_model_root(str(root), None, "fl2va") == str(root / "FL2VA" / "text_encoder")
|
||||
assert resolve_minimax_h3_encoder_model_root(str(root), None, "ref2va") == str(root / "Ref2VA" / "text_encoder")
|
||||
assert resolve_minimax_h3_encoder_model_root(str(root), None, "combined") == str(root / "FL2VA" / "text_encoder")
|
||||
assert resolve_minimax_h3_encoder_model_root(str(root / "Ref2VA"), None, None) == str(
|
||||
root / "Ref2VA" / "text_encoder"
|
||||
)
|
||||
assert resolve_minimax_h3_encoder_model_root(str(root / "FL2VA"), None, "ref2va") == str(
|
||||
root / "Ref2VA" / "text_encoder"
|
||||
)
|
||||
|
||||
|
||||
def test_checkpoint_resolver_rejects_unknown_task(tmp_path):
|
||||
with pytest.raises(ValueError, match="task_type must be one of"):
|
||||
resolve_minimax_h3_model_root(str(tmp_path), None, "unknown")
|
||||
resolve_minimax_h3_encoder_model_root(str(tmp_path), None, "unknown")
|
||||
|
||||
|
||||
def test_partition_resolver_preserves_consumer_auto_default(tmp_path):
|
||||
|
||||
Reference in new issue
Block a user