232 lines
10 KiB
Python
232 lines
10 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.models.minimax_h3.continuation import (
|
|
diffuse_continuation,
|
|
plan_continuation_windows,
|
|
resolve_continuation,
|
|
)
|
|
from vllm_omni.diffusion.models.minimax_h3.packed_sequence import minimax_h3_packed_sequence_ref2va_blocks
|
|
from vllm_omni.errors import OmniClientError
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
|
|
|
|
|
@pytest.mark.parametrize("vary_prompts", [False, True])
|
|
@pytest.mark.parametrize("total,window", [(1807, 277), (1807, 243), (124, 277)])
|
|
def test_continuation_retains_prefix_and_global_stereo_timeline(total, window, vary_prompts):
|
|
audio_t = round(total * 40 / 24)
|
|
source = torch.arange(2 * audio_t * 32).reshape(2, audio_t, 32).float()
|
|
calls = []
|
|
plan = plan_continuation_windows(total, window, 22)
|
|
texts = [(torch.full((3 + i, 8), float(i)), torch.ones(3 + i, dtype=torch.long)) for i in range(len(plan))]
|
|
|
|
def sample(**kw):
|
|
calls.append(kw)
|
|
t = kw["latent_t"]
|
|
if vary_prompts:
|
|
assert kw["text_embeddings"] is texts[len(calls) - 1][0]
|
|
assert kw["text_tags"] is texts[len(calls) - 1][1]
|
|
assert kw["media_time_origin"] == len(texts[-1][0])
|
|
# Different constants expose replacement of an old prefix or an
|
|
# accidental append of the hidden overlap.
|
|
video = torch.full((1, 24, t, 2, 2), float(len(calls)))
|
|
audio = kw["locked_audio_rows"].reshape(2, kw["audio_t"], 32).permute(0, 2, 1)
|
|
if len(calls) > 1:
|
|
assert torch.all(kw["visual_condition"] == len(calls) - 1)
|
|
assert kw["ref_blocks"][-1]["kind"] == "latent_guide"
|
|
return video, audio
|
|
|
|
kwargs = dict(
|
|
num_frames=total,
|
|
latent_t=(total - 5) // 17 * 5 + 2,
|
|
latent_h=2,
|
|
latent_w=2,
|
|
audio_t=audio_t,
|
|
locked_audio_rows=source.reshape(-1, 32),
|
|
)
|
|
video, audio = diffuse_continuation(
|
|
sample, kwargs, window_frames=window, overlap_frames=22, text_conditioning=texts if vary_prompts else None
|
|
)
|
|
torch.testing.assert_close(audio, source.permute(0, 2, 1))
|
|
plan = plan_continuation_windows(total, window, 22)
|
|
assert [call["temporal_offset"] for call in calls] == pytest.approx([part.start * 40 / 24 for part in plan])
|
|
offset = 0
|
|
for i, part in enumerate(plan):
|
|
new_frames = part.end - part.start - part.overlap
|
|
count = (new_frames - 5) // 17 * 5 + 2 if i == 0 else new_frames // 17 * 5
|
|
assert torch.all(video[:, :, offset : offset + count] == i + 1)
|
|
offset += count
|
|
assert offset == video.shape[2] == kwargs["latent_t"]
|
|
|
|
|
|
@pytest.mark.parametrize("temporal_offset", [0.0, 425.0, 221 * 40 / 24])
|
|
def test_latent_guide_shares_target_clock_and_remains_condition_only(temporal_offset):
|
|
packed = minimax_h3_packed_sequence_ref2va_blocks(
|
|
temporal_offset=temporal_offset,
|
|
text_len=3,
|
|
latent_t=12,
|
|
latent_h=4,
|
|
latent_w=4,
|
|
audio_t=68,
|
|
ref_blocks=[
|
|
{"kind": "image", "latent_h": 4, "latent_w": 4},
|
|
{"kind": "latent_guide", "latent_t": 7, "latent_h": 4, "latent_w": 4, "ref_audio_t": 37},
|
|
],
|
|
)
|
|
pos = packed["img_position_ids"]
|
|
visual = packed["img_pos"]
|
|
target = visual[packed["update_mask"]]
|
|
guide = visual[~packed["update_mask"]][4:]
|
|
torch.testing.assert_close(pos[guide], pos[target[:28]])
|
|
audio = packed["audio_pos"]
|
|
guide_audio = audio[~packed["audio_update_mask"]].reshape(2, 37)
|
|
target_audio = audio[packed["audio_update_mask"]].reshape(2, 68)
|
|
torch.testing.assert_close(pos[guide_audio], pos[target_audio[:, :37]])
|
|
assert pos[target[0], 0] == pytest.approx(4 + temporal_offset)
|
|
torch.testing.assert_close(pos[:3, 0], torch.arange(3, dtype=torch.float64))
|
|
assert torch.all(pos[3:7, 0] == 3) # static reference image stays fixed
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"extra,step",
|
|
[
|
|
({"long_video_mode": "bad"}, False),
|
|
({"long_video_mode": "continuation"}, True),
|
|
({"long_video_mode": "continuation", "continuation_overlap_frames": 23}, False),
|
|
({"long_video_mode": "continuation", "continuation_window_frames": True}, False),
|
|
({"long_video_mode": "continuation", "continuation_overlap_frames": 277}, False),
|
|
],
|
|
)
|
|
def test_invalid_continuation_options(extra, step):
|
|
with pytest.raises(OmniClientError):
|
|
resolve_continuation(extra, task="ref2va", step_execution=step)
|
|
|
|
|
|
def test_global_offset_moves_temporal_references_but_not_images_or_padding():
|
|
kwargs = dict(
|
|
text_len=3,
|
|
latent_t=12,
|
|
latent_h=4,
|
|
latent_w=4,
|
|
audio_t=68,
|
|
ref_blocks=[
|
|
{"kind": "image", "latent_h": 4, "latent_w": 4},
|
|
{"kind": "audio", "ref_audio_t": 3},
|
|
{"kind": "video_audio", "ref_audio_t": 2, "latent_t": 2, "latent_h": 4, "latent_w": 4},
|
|
{"kind": "image", "latent_h": 4, "latent_w": 4},
|
|
{"kind": "latent_guide", "ref_audio_t": 37, "latent_t": 7, "latent_h": 4, "latent_w": 4},
|
|
],
|
|
)
|
|
base = minimax_h3_packed_sequence_ref2va_blocks(**kwargs)
|
|
shifted = minimax_h3_packed_sequence_ref2va_blocks(**kwargs, temporal_offset=425.0)
|
|
moving = base["image_mask"] | base["audio_mask"]
|
|
moving[3:7] = False # first static image
|
|
moving[25:29] = False # static image after temporal references
|
|
delta = shifted["img_position_ids"] - base["img_position_ids"]
|
|
torch.testing.assert_close(delta[:, 0], moving.double() * 425)
|
|
assert torch.count_nonzero(delta[:, 1:]) == 0
|
|
for key, value in base.items():
|
|
if isinstance(value, torch.Tensor) and key != "img_position_ids":
|
|
torch.testing.assert_close(shifted[key], value)
|
|
|
|
|
|
def test_long_ref2va_defaults_to_continuation_with_explicit_full_opt_out():
|
|
assert resolve_continuation({"long_video": True}, task="ref2va", step_execution=False) == (277, 22)
|
|
assert resolve_continuation({}, task="ref2va", step_execution=False) is None
|
|
assert (
|
|
resolve_continuation({"long_video": True, "long_video_mode": "full"}, task="ref2va", step_execution=False)
|
|
is None
|
|
)
|
|
|
|
|
|
def test_different_chunk_prompt_lengths_keep_the_same_media_clock():
|
|
kwargs = dict(
|
|
latent_t=7,
|
|
latent_h=4,
|
|
latent_w=4,
|
|
audio_t=37,
|
|
ref_blocks=[{"kind": "image", "latent_h": 4, "latent_w": 4}],
|
|
media_time_origin=20,
|
|
)
|
|
first = minimax_h3_packed_sequence_ref2va_blocks(text_len=3, **kwargs)
|
|
second = minimax_h3_packed_sequence_ref2va_blocks(text_len=15, temporal_offset=425, **kwargs)
|
|
for key, mask in (("img_pos", "update_mask"), ("audio_pos", "audio_update_mask")):
|
|
a = first["img_position_ids"][first[key][first[mask]]]
|
|
b = second["img_position_ids"][second[key][second[mask]]]
|
|
torch.testing.assert_close(b[:, 0] - a[:, 0], torch.full_like(a[:, 0], 425))
|
|
torch.testing.assert_close(b[:, 1:], a[:, 1:])
|
|
|
|
|
|
@pytest.mark.parametrize("prompt_count", [None, 7, 6])
|
|
def test_request_encodes_only_window_prompts_and_prepares_references_once(monkeypatch, prompt_count):
|
|
from PIL import Image
|
|
|
|
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderMediaConditioning
|
|
|
|
pipeline = object.__new__(module.MiniMaxH3Pipeline)
|
|
torch.nn.Module.__init__(pipeline)
|
|
pipeline.load_text_encoder = True
|
|
pipeline.device = torch.device("cpu")
|
|
pipeline.od_config = None
|
|
pipeline._active_turbo_spec = lambda sampling: None
|
|
pipeline._has_active_native_lora = lambda sampling: False
|
|
pipeline._resolve_task = lambda *args, **kwargs: "ref2va"
|
|
encoded, prepared_inputs, encoded_media = [], [], []
|
|
prepare = module.prepare_encoder_inputs
|
|
|
|
def prepare_once(*args, **kwargs):
|
|
result = prepare(*args, **kwargs)
|
|
prepared_inputs.append(result)
|
|
return result
|
|
|
|
def encode_text(prepared):
|
|
encoded.append(prepared)
|
|
return torch.full((3, 5120), len(encoded), dtype=torch.bfloat16), torch.zeros(3, dtype=torch.long)
|
|
|
|
def encode_media(media):
|
|
encoded_media.append(media)
|
|
return MiniMaxH3EncoderMediaConditioning(
|
|
task=media.task,
|
|
height=media.height,
|
|
width=media.width,
|
|
num_frames=media.num_frames,
|
|
latent_t=media.latent_t,
|
|
audio_t=media.audio_t,
|
|
)
|
|
|
|
monkeypatch.setattr(module, "prepare_encoder_inputs", prepare_once)
|
|
pipeline.encode_prompt = encode_text
|
|
pipeline._encode_local_media = encode_media
|
|
pipeline._prepare_encoder_conditioning_inputs = lambda conditioning, sampling: {
|
|
"text_embeddings": conditioning.hidden_states,
|
|
"continuation": (277, 22),
|
|
"num_frames": conditioning.num_frames,
|
|
"task": conditioning.task,
|
|
}
|
|
extra = {"task": "ref2va", "long_video": True, "duration": 75}
|
|
if prompt_count is not None:
|
|
extra["continuation_prompts"] = [f"window {i}" for i in range(prompt_count)]
|
|
sampling = OmniDiffusionSamplingParams(width=64, height=64, fps=24, extra_args=extra)
|
|
request = {"prompt": "main prompt", "multi_modal_data": {"image": Image.new("RGB", (256, 256))}}
|
|
if prompt_count == 6:
|
|
with pytest.raises(OmniClientError, match="exactly 7"):
|
|
pipeline._prepare_request_inputs(request, sampling)
|
|
assert not encoded and not encoded_media
|
|
return
|
|
context = pipeline._prepare_request_inputs(request, sampling)
|
|
assert [item.prompt for item in encoded] == extra.get("continuation_prompts", ["main prompt"])
|
|
assert len(prepared_inputs) == len(encoded_media) == 1
|
|
assert all(item.media is prepared_inputs[0].media for item in encoded)
|
|
if prompt_count is not None:
|
|
texts = context["continuation_text_conditioning"]
|
|
assert len(texts) == 7
|
|
assert texts[0][0] is context["text_embeddings"]
|
|
assert [int(hidden[0, 0]) for hidden, _ in texts] == list(range(1, 8))
|
|
else:
|
|
assert "continuation_text_conditioning" not in context
|