280 lines
9.0 KiB
Python
280 lines
9.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
|
|
|
|
|
def test_video_patchify_round_trip_preserves_values():
|
|
from vllm_omni.diffusion.models.minimax_h3.packed_tokens import (
|
|
minimax_h3_patchify_video_latent,
|
|
minimax_h3_unpatchify_video_tokens,
|
|
)
|
|
|
|
latent = torch.arange(2 * 3 * 2 * 4 * 6).reshape(2, 3, 2, 4, 6)
|
|
rows = minimax_h3_patchify_video_latent(
|
|
latent,
|
|
patch_size=(1, 2, 2),
|
|
)
|
|
restored = minimax_h3_unpatchify_video_tokens(
|
|
rows,
|
|
latent_shape=(2, 2, 3, 3),
|
|
patch_size=(1, 2, 2),
|
|
)
|
|
|
|
torch.testing.assert_close(restored, latent)
|
|
|
|
|
|
def test_audio_pack_round_trip_preserves_channel_major_order():
|
|
from vllm_omni.diffusion.models.minimax_h3.packed_tokens import (
|
|
minimax_h3_pack_audio_latent,
|
|
minimax_h3_unpack_audio_tokens,
|
|
)
|
|
|
|
latent = torch.arange(2 * 4 * 5).reshape(2, 4, 5)
|
|
rows = minimax_h3_pack_audio_latent(latent)
|
|
restored = minimax_h3_unpack_audio_tokens(
|
|
rows,
|
|
audio_t=10,
|
|
audio_channel=2,
|
|
)
|
|
|
|
torch.testing.assert_close(restored, latent)
|
|
|
|
|
|
def test_t2va_and_fl2va_packing_keep_update_rows_separate():
|
|
from vllm_omni.diffusion.models.minimax_h3.packed_sequence import (
|
|
minimax_h3_packed_sequence,
|
|
)
|
|
|
|
common = dict(
|
|
text_len=4,
|
|
latent_t=2,
|
|
latent_h=4,
|
|
latent_w=6,
|
|
audio_t=3,
|
|
)
|
|
t2va = minimax_h3_packed_sequence(
|
|
**common,
|
|
include_keyframe_cond=False,
|
|
)
|
|
fl2va = minimax_h3_packed_sequence(
|
|
**common,
|
|
include_keyframe_cond=True,
|
|
keyframe_frame_indices=[0],
|
|
frame_count=5,
|
|
)
|
|
|
|
assert int(t2va["seq_len"]) == 64
|
|
assert t2va["img_pos"].numel() == 12
|
|
assert t2va["update_mask"].all()
|
|
assert fl2va["img_pos"].numel() == 18
|
|
assert fl2va["update_mask"].sum().item() == 12
|
|
assert (~fl2va["update_mask"]).sum().item() == 6
|
|
assert t2va["audio_pos"].numel() == fl2va["audio_pos"].numel() == 6
|
|
assert t2va["video_spans"] == ({"start": 10, "latent_grid": (2, 2, 3), "role": "target"},)
|
|
assert fl2va["video_spans"] == ({"start": 16, "latent_grid": (2, 2, 3), "role": "target"},)
|
|
|
|
|
|
def test_ref2va_packing_tracks_video_and_audio_update_masks():
|
|
from vllm_omni.diffusion.models.minimax_h3.packed_sequence import (
|
|
minimax_h3_packed_sequence_ref2va_blocks,
|
|
)
|
|
|
|
packed = minimax_h3_packed_sequence_ref2va_blocks(
|
|
text_len=4,
|
|
latent_t=2,
|
|
latent_h=4,
|
|
latent_w=6,
|
|
audio_t=3,
|
|
ref_blocks=[
|
|
{"kind": "image", "latent_h": 4, "latent_w": 4},
|
|
{"kind": "audio", "ref_audio_t": 2},
|
|
],
|
|
)
|
|
|
|
assert int(packed["seq_len"]) == 64
|
|
assert packed["img_pos"].numel() == 16
|
|
assert packed["update_mask"].sum().item() == 12
|
|
assert packed["audio_pos"].numel() == 10
|
|
assert packed["audio_update_mask"].sum().item() == 6
|
|
assert packed["cu_seqlens"].tolist() == [0, 30, 64]
|
|
|
|
|
|
def test_ref2va_packing_publishes_physical_video_spans_only():
|
|
from vllm_omni.diffusion.models.minimax_h3.packed_sequence import (
|
|
minimax_h3_packed_sequence_ref2va_blocks,
|
|
)
|
|
|
|
packed = minimax_h3_packed_sequence_ref2va_blocks(
|
|
text_len=4,
|
|
latent_t=2,
|
|
latent_h=4,
|
|
latent_w=6,
|
|
audio_t=3,
|
|
ref_blocks=[
|
|
{"kind": "image", "latent_h": 4, "latent_w": 4},
|
|
{"kind": "video_audio", "ref_audio_t": 2, "latent_t": 2, "latent_h": 4, "latent_w": 6},
|
|
{"kind": "audio", "ref_audio_t": 1},
|
|
{"kind": "video", "ref_audio_t": 0, "latent_t": 1, "latent_h": 4, "latent_w": 4},
|
|
],
|
|
)
|
|
|
|
# Image and all audio rows intentionally stay out of the sparse spans.
|
|
assert packed["video_spans"] == (
|
|
{"start": 12, "latent_grid": (2, 2, 3), "role": "reference"},
|
|
{"start": 26, "latent_grid": (1, 2, 2), "role": "reference"},
|
|
{"start": 36, "latent_grid": (2, 2, 3), "role": "target"},
|
|
)
|
|
|
|
|
|
def test_condition_noise_is_seeded_and_keeps_clean_anchor_at_timestep_one():
|
|
from vllm_omni.diffusion.models.minimax_h3.condition_noise import (
|
|
minimax_h3_audio_cond_noise_aug_rows,
|
|
minimax_h3_imgvid_cond_noise_aug_rows,
|
|
)
|
|
|
|
image_rows = torch.randn(4, 96)
|
|
audio_rows = torch.randn(6, 32)
|
|
|
|
torch.testing.assert_close(
|
|
minimax_h3_imgvid_cond_noise_aug_rows(
|
|
image_rows,
|
|
condition_shapes=[(1, 4, 4)],
|
|
target_latent_t=2,
|
|
imgvid_cond_num_frames=1,
|
|
seed=42,
|
|
noise_aug=1.0,
|
|
),
|
|
image_rows,
|
|
)
|
|
torch.testing.assert_close(
|
|
minimax_h3_audio_cond_noise_aug_rows(
|
|
audio_rows,
|
|
condition_audio_t=[3],
|
|
seed=42,
|
|
noise_aug=1.0,
|
|
),
|
|
audio_rows,
|
|
)
|
|
|
|
first = minimax_h3_audio_cond_noise_aug_rows(
|
|
audio_rows,
|
|
condition_audio_t=[3],
|
|
seed=42,
|
|
noise_aug=0.25,
|
|
)
|
|
second = minimax_h3_audio_cond_noise_aug_rows(
|
|
audio_rows,
|
|
condition_audio_t=[3],
|
|
seed=42,
|
|
noise_aug=0.25,
|
|
)
|
|
torch.testing.assert_close(first, second)
|
|
|
|
|
|
def test_condition_noise_accepts_a_reference_video_longer_than_the_target():
|
|
from vllm_omni.diffusion.models.minimax_h3.condition_noise import (
|
|
minimax_h3_imgvid_cond_noise_aug_rows,
|
|
)
|
|
|
|
rows = torch.zeros(4 * 2 * 2, 96)
|
|
result = minimax_h3_imgvid_cond_noise_aug_rows(
|
|
rows,
|
|
condition_shapes=[(4, 4, 4)],
|
|
target_latent_t=2,
|
|
imgvid_cond_num_frames=1,
|
|
seed=7,
|
|
noise_aug=0.5,
|
|
)
|
|
assert result.shape == rows.shape
|
|
|
|
|
|
def test_explicit_seq_len_pins_one_shape_across_prompt_lengths():
|
|
"""Prompts of different token counts must land on the same packed length.
|
|
|
|
Without the pin they fall into different 64-row buckets, and each bucket is
|
|
a packed shape the compiled transformer has not seen before.
|
|
"""
|
|
from vllm_omni.diffusion.models.minimax_h3.packed_sequence import (
|
|
minimax_h3_packed_sequence,
|
|
)
|
|
|
|
common = dict(
|
|
latent_t=2,
|
|
latent_h=4,
|
|
latent_w=6,
|
|
audio_t=3,
|
|
include_keyframe_cond=False,
|
|
)
|
|
# used = text_len + 18 rows: 22 -> bucket 64; 118 -> bucket 128.
|
|
short = minimax_h3_packed_sequence(text_len=4, **common)
|
|
long = minimax_h3_packed_sequence(text_len=100, **common)
|
|
assert int(short["seq_len"]) == 64
|
|
assert int(long["seq_len"]) == 128
|
|
|
|
short_pinned = minimax_h3_packed_sequence(text_len=4, seq_len=192, **common)
|
|
long_pinned = minimax_h3_packed_sequence(text_len=100, seq_len=192, **common)
|
|
assert int(short_pinned["seq_len"]) == int(long_pinned["seq_len"]) == 192
|
|
# The pin only adds padding rows; the used prefix is untouched.
|
|
assert int(short_pinned["cu_seqlens"][1]) == int(short["cu_seqlens"][1]) == 22
|
|
assert int(long_pinned["cu_seqlens"][1]) == int(long["cu_seqlens"][1]) == 118
|
|
assert short_pinned["img_pos"].tolist() == short["img_pos"].tolist()
|
|
|
|
|
|
def test_the_ref2va_packer_pins_one_shape_across_prompt_lengths_too():
|
|
"""Both packers share one resolver, so both honour the pin."""
|
|
from vllm_omni.diffusion.models.minimax_h3.packed_sequence import (
|
|
minimax_h3_packed_sequence_ref2va_blocks,
|
|
)
|
|
|
|
common = dict(
|
|
latent_t=2,
|
|
latent_h=4,
|
|
latent_w=6,
|
|
audio_t=3,
|
|
ref_blocks=[{"kind": "audio", "ref_audio_t": 1}],
|
|
)
|
|
short = minimax_h3_packed_sequence_ref2va_blocks(text_len=4, **common)
|
|
long = minimax_h3_packed_sequence_ref2va_blocks(text_len=100, **common)
|
|
assert int(short["seq_len"]) != int(long["seq_len"])
|
|
|
|
short_pinned = minimax_h3_packed_sequence_ref2va_blocks(text_len=4, seq_len=192, **common)
|
|
long_pinned = minimax_h3_packed_sequence_ref2va_blocks(text_len=100, seq_len=192, **common)
|
|
assert int(short_pinned["seq_len"]) == int(long_pinned["seq_len"]) == 192
|
|
assert int(short_pinned["cu_seqlens"][1]) == int(short["cu_seqlens"][1])
|
|
assert int(long_pinned["cu_seqlens"][1]) == int(long["cu_seqlens"][1])
|
|
|
|
|
|
def test_explicit_seq_len_below_used_rows_is_a_client_error():
|
|
"""Both packers reject an under-sized pin as a 4xx, not a server error."""
|
|
from vllm_omni.diffusion.models.minimax_h3.packed_sequence import (
|
|
minimax_h3_packed_sequence,
|
|
minimax_h3_packed_sequence_ref2va_blocks,
|
|
)
|
|
from vllm_omni.errors import OmniClientError
|
|
|
|
with pytest.raises(OmniClientError, match="used rows"):
|
|
minimax_h3_packed_sequence(
|
|
text_len=4,
|
|
latent_t=2,
|
|
latent_h=4,
|
|
latent_w=6,
|
|
audio_t=3,
|
|
include_keyframe_cond=False,
|
|
seq_len=8,
|
|
)
|
|
|
|
with pytest.raises(OmniClientError, match="used rows"):
|
|
minimax_h3_packed_sequence_ref2va_blocks(
|
|
text_len=4,
|
|
latent_t=2,
|
|
latent_h=4,
|
|
latent_w=6,
|
|
audio_t=3,
|
|
ref_blocks=[{"kind": "audio", "ref_audio_t": 1}],
|
|
seq_len=8,
|
|
)
|