[4/N] Unify worker-side MP4 encoding and wire Wan into it (#7048)
Signed-off-by: specture724 <specture724@gmail.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
12 files changed
+1177
-232
No files matched your search
@@ -2852,7 +2852,7 @@ def test_decode_to_mp4_batches_consumer_transfers(monkeypatch):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
|
||||
|
||||
class FakeEncoder:
|
||||
instances = []
|
||||
instances: list["FakeEncoder"] = []
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
self.pushes = []
|
||||
@@ -2873,6 +2873,8 @@ def test_decode_to_mp4_batches_consumer_transfers(monkeypatch):
|
||||
return torch.zeros(1, 1, 2)
|
||||
|
||||
class FakeVideoVAE:
|
||||
chunk_value_range = (0.0, 1.0)
|
||||
|
||||
def decode_with_chunks(self, latent, *, on_chunk):
|
||||
for value in (0.0, 0.25, 0.5):
|
||||
on_chunk(torch.full((1, 3, 3, 2, 2), value))
|
||||
@@ -2884,7 +2886,7 @@ def test_decode_to_mp4_batches_consumer_transfers(monkeypatch):
|
||||
pipeline.device = torch.device("cpu")
|
||||
monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False)
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.utils.media_utils.ChunkedMP4Encoder",
|
||||
"vllm_omni.diffusion.utils.chunked_video.ChunkedMP4Encoder",
|
||||
FakeEncoder,
|
||||
)
|
||||
|
||||
@@ -2904,12 +2906,81 @@ def test_decode_to_mp4_batches_consumer_transfers(monkeypatch):
|
||||
assert encoder.pushes[1].shape == (3, 2, 2, 3)
|
||||
|
||||
|
||||
def _peer_rank_preencode_pipeline(monkeypatch):
|
||||
"""An H3 pipeline whose video VAE is a distributed rank that owns no output.
|
||||
|
||||
Every rank of the VAE group must call ``decode_with_chunks`` to keep the
|
||||
temporal collectives in lockstep, but only the owner receives chunks, so a
|
||||
peer's callback is never invoked.
|
||||
"""
|
||||
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
|
||||
|
||||
class FakeAudioVAE:
|
||||
def decode_latent(self, latent):
|
||||
return torch.zeros(1, 1, 2)
|
||||
|
||||
class PeerRankVideoVAE:
|
||||
chunk_value_range = (0.0, 1.0)
|
||||
|
||||
def decode_with_chunks(self, latent, *, on_chunk):
|
||||
del latent, on_chunk
|
||||
|
||||
pipeline = object.__new__(mod.MiniMaxH3Pipeline)
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
pipeline.audio_vae = FakeAudioVAE()
|
||||
pipeline.video_vae = PeerRankVideoVAE()
|
||||
pipeline.device = torch.device("cpu")
|
||||
pipeline.od_config = SimpleNamespace()
|
||||
monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False)
|
||||
return pipeline
|
||||
|
||||
|
||||
def test_peer_vae_rank_returns_no_preencoded_output_instead_of_failing(monkeypatch):
|
||||
"""A rank that owns no output must not turn its silence into a request error."""
|
||||
pipeline = _peer_rank_preencode_pipeline(monkeypatch)
|
||||
|
||||
output = pipeline.decode_to_mp4(torch.zeros(1), torch.zeros(1), height=2, width=2)
|
||||
|
||||
assert output == b""
|
||||
|
||||
|
||||
def test_peer_vae_rank_post_decode_reaches_post_processing(monkeypatch):
|
||||
"""The real H3 post_decode chain, not just the shared consumer, tolerates a peer rank."""
|
||||
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
|
||||
from vllm_omni.diffusion.worker.utils import StepRequestState
|
||||
|
||||
pipeline = _peer_rank_preencode_pipeline(monkeypatch)
|
||||
monkeypatch.setattr(pipeline, "_unpack_denoised_rows", lambda *args, **kwargs: (torch.zeros(1), torch.zeros(1)))
|
||||
state = StepRequestState(
|
||||
request_id="peer-rank",
|
||||
sampling=SimpleNamespace(num_outputs_per_prompt=1),
|
||||
prompt="a prompt",
|
||||
)
|
||||
state.latents = torch.zeros(1)
|
||||
state.extra[mod._STEP_BRANCH] = object()
|
||||
state.extra[mod._STEP_AUDIO_ROWS] = torch.zeros(1)
|
||||
state.extra[mod._STEP_SHAPE] = {
|
||||
"latent_t": 1,
|
||||
"latent_h": 1,
|
||||
"latent_w": 1,
|
||||
"audio_t": 1,
|
||||
"height": 2,
|
||||
"width": 2,
|
||||
"preencode_mp4": True,
|
||||
}
|
||||
|
||||
output = pipeline.post_decode(state)
|
||||
|
||||
assert output.output == (b"", None)
|
||||
assert mod._minimax_h3_post_process(output.output)["video"] == [b""]
|
||||
|
||||
|
||||
def test_request_video_codec_options_reach_the_preencoded_mp4_encoder(monkeypatch):
|
||||
"""A client's encoder options must survive the worker-side pre-encode path."""
|
||||
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
|
||||
|
||||
class FakeEncoder:
|
||||
instances = []
|
||||
instances: list["FakeEncoder"] = []
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
self.kwargs = kwargs
|
||||
@@ -2929,6 +3000,8 @@ def test_request_video_codec_options_reach_the_preencoded_mp4_encoder(monkeypatc
|
||||
return torch.zeros(1, 1, 2)
|
||||
|
||||
class FakeVideoVAE:
|
||||
chunk_value_range = (0.0, 1.0)
|
||||
|
||||
def decode_with_chunks(self, latent, *, on_chunk):
|
||||
on_chunk(torch.zeros(1, 3, 1, 2, 2))
|
||||
|
||||
@@ -2938,7 +3011,7 @@ def test_request_video_codec_options_reach_the_preencoded_mp4_encoder(monkeypatc
|
||||
pipeline.video_vae = FakeVideoVAE()
|
||||
pipeline.device = torch.device("cpu")
|
||||
monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False)
|
||||
monkeypatch.setattr("vllm_omni.diffusion.utils.media_utils.ChunkedMP4Encoder", FakeEncoder)
|
||||
monkeypatch.setattr("vllm_omni.diffusion.utils.chunked_video.ChunkedMP4Encoder", FakeEncoder)
|
||||
|
||||
pipeline.decode_to_mp4(
|
||||
torch.zeros(1),
|
||||
|
||||
@@ -0,0 +1,165 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import io
|
||||
from types import SimpleNamespace
|
||||
|
||||
import av
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.models.wan2_2.chunked_mp4 import (
|
||||
WAN_DEFAULT_OUTPUT_FPS,
|
||||
resolve_wan_output_fps,
|
||||
resolve_wan_preencode_batch_frames,
|
||||
resolve_wan_preencode_mp4,
|
||||
resolve_wan_video_codec_options,
|
||||
wan_preencoded_mp4_payload,
|
||||
)
|
||||
from vllm_omni.diffusion.utils.chunked_video import decode_to_mp4
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
class _FakeWanVAE:
|
||||
"""Publish temporal chunks the way the Wan VAE seam does."""
|
||||
|
||||
chunk_value_range = (-1.0, 1.0)
|
||||
|
||||
def __init__(self, *, batch: int, chunks: int, frames_per_chunk: int, height: int, width: int):
|
||||
self.batch = batch
|
||||
self.chunks = chunks
|
||||
self.frames_per_chunk = frames_per_chunk
|
||||
self.height = height
|
||||
self.width = width
|
||||
self.chunk_calls = 0
|
||||
|
||||
def decode_with_chunks(self, latents, *, on_chunk):
|
||||
del latents
|
||||
assert on_chunk is not None
|
||||
for index in range(self.chunks):
|
||||
self.chunk_calls += 1
|
||||
value = -1.0 + 2.0 * index / max(self.chunks - 1, 1)
|
||||
on_chunk(torch.full((self.batch, 3, self.frames_per_chunk, self.height, self.width), value))
|
||||
|
||||
|
||||
def _decoded_frame_count(data: bytes) -> int:
|
||||
with av.open(io.BytesIO(data)) as container:
|
||||
return sum(1 for _ in container.decode(video=0))
|
||||
|
||||
|
||||
def test_decode_wan_latents_to_mp4_emits_one_playable_mp4_per_batch_entry():
|
||||
vae = _FakeWanVAE(batch=2, chunks=6, frames_per_chunk=2, height=16, width=32)
|
||||
|
||||
videos = decode_to_mp4(vae, torch.zeros(1), fps=24, batch_frames=4)
|
||||
|
||||
assert len(videos) == 2
|
||||
for data in videos:
|
||||
assert data[:4] and _decoded_frame_count(data) == 12
|
||||
|
||||
|
||||
def test_decode_wan_latents_to_mp4_coalesces_transfers_without_losing_frames():
|
||||
"""batch_frames only changes how many transfers carry the same frames."""
|
||||
unbatched = _FakeWanVAE(batch=1, chunks=8, frames_per_chunk=1, height=16, width=32)
|
||||
batched = _FakeWanVAE(batch=1, chunks=8, frames_per_chunk=1, height=16, width=32)
|
||||
|
||||
per_chunk = decode_to_mp4(unbatched, torch.zeros(1), fps=24, batch_frames=1)
|
||||
coalesced = decode_to_mp4(batched, torch.zeros(1), fps=24, batch_frames=4)
|
||||
|
||||
assert _decoded_frame_count(per_chunk[0]) == _decoded_frame_count(coalesced[0]) == 8
|
||||
|
||||
|
||||
def test_decode_wan_latents_to_mp4_returns_nothing_for_a_rank_without_output():
|
||||
class SilentVAE:
|
||||
chunk_value_range = (-1.0, 1.0)
|
||||
|
||||
def decode_with_chunks(self, latents, *, on_chunk):
|
||||
del latents, on_chunk
|
||||
|
||||
assert decode_to_mp4(SilentVAE(), torch.zeros(1), fps=24) == []
|
||||
|
||||
|
||||
def test_decode_wan_latents_to_mp4_rejects_a_non_positive_batch():
|
||||
with pytest.raises(ValueError, match="batch_frames"):
|
||||
decode_to_mp4(
|
||||
_FakeWanVAE(batch=1, chunks=1, frames_per_chunk=1, height=16, width=16), None, fps=24, batch_frames=0
|
||||
)
|
||||
|
||||
|
||||
def test_decode_wan_latents_to_mp4_rejects_a_vae_without_the_capability():
|
||||
class PlainVAE:
|
||||
chunk_value_range = (-1.0, 1.0)
|
||||
|
||||
def decode(self, z, return_dict=True):
|
||||
del z, return_dict
|
||||
|
||||
with pytest.raises(TypeError, match="chunked VAE decode capability"):
|
||||
decode_to_mp4(PlainVAE(), torch.zeros(1), fps=24)
|
||||
|
||||
|
||||
def test_preencode_flag_is_off_unless_the_request_asks_for_it():
|
||||
assert resolve_wan_preencode_mp4(SimpleNamespace(extra_args=None), output_type="np") is False
|
||||
assert resolve_wan_preencode_mp4(SimpleNamespace(extra_args={}), output_type="np") is False
|
||||
assert resolve_wan_preencode_mp4(SimpleNamespace(extra_args={"preencode_mp4": True}), output_type="np") is True
|
||||
|
||||
|
||||
def test_preencode_rejects_combinations_it_cannot_serve():
|
||||
interpolating = SimpleNamespace(extra_args={"preencode_mp4": True}, enable_frame_interpolation=True)
|
||||
with pytest.raises(ValueError, match="enable_frame_interpolation"):
|
||||
resolve_wan_preencode_mp4(interpolating, output_type="np")
|
||||
|
||||
with pytest.raises(ValueError, match="output_type"):
|
||||
resolve_wan_preencode_mp4(SimpleNamespace(extra_args={"preencode_mp4": True}), output_type="pil")
|
||||
|
||||
|
||||
def test_output_fps_mirrors_the_serving_fallback():
|
||||
assert resolve_wan_output_fps(SimpleNamespace(fps=16)) == 16
|
||||
assert resolve_wan_output_fps(SimpleNamespace(fps=[16])) == 16
|
||||
assert resolve_wan_output_fps(SimpleNamespace(fps=None)) == WAN_DEFAULT_OUTPUT_FPS
|
||||
assert resolve_wan_output_fps(SimpleNamespace()) == WAN_DEFAULT_OUTPUT_FPS
|
||||
|
||||
|
||||
def test_codec_options_reach_the_encoder_as_strings():
|
||||
params = SimpleNamespace(extra_args={"video_codec_options": {"preset": "ultrafast", "threads": 0}})
|
||||
assert resolve_wan_video_codec_options(params) == {"preset": "ultrafast", "threads": "0"}
|
||||
assert resolve_wan_video_codec_options(SimpleNamespace(extra_args={})) == {"preset": "ultrafast", "threads": "0"}
|
||||
assert resolve_wan_video_codec_options(SimpleNamespace(extra_args={"video_codec_options": None})) is None
|
||||
assert resolve_wan_video_codec_options(SimpleNamespace(extra_args={"video_codec_options": {}})) == {}
|
||||
|
||||
|
||||
def test_preencoded_payload_passes_bytes_through_and_ignores_tensors():
|
||||
assert wan_preencoded_mp4_payload(torch.zeros(1, 3, 2, 4, 4)) is None
|
||||
assert wan_preencoded_mp4_payload([b"one", b"two"]) == {"payload": {"video": [b"one", b"two"]}, "metadata": {}}
|
||||
assert wan_preencoded_mp4_payload(b"one") == {"payload": {"video": [b"one"]}, "metadata": {}}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", [1, 5, 33])
|
||||
def test_request_batch_threshold_controls_transfers_without_losing_frames(monkeypatch, value):
|
||||
from vllm_omni.diffusion.utils import chunked_video
|
||||
|
||||
transfers = []
|
||||
quantize = chunked_video.chunk_to_uint8_frames
|
||||
|
||||
def capture(chunk, value_range):
|
||||
transfers.append(chunk.shape[2])
|
||||
return quantize(chunk, value_range)
|
||||
|
||||
monkeypatch.setattr(chunked_video, "chunk_to_uint8_frames", capture)
|
||||
params = SimpleNamespace(extra_args={"preencode_mp4": True, "preencode_batch_frames": value})
|
||||
vae = _FakeWanVAE(batch=1, chunks=4, frames_per_chunk=3, height=16, width=32)
|
||||
videos = decode_to_mp4(vae, torch.zeros(1), fps=24, batch_frames=resolve_wan_preencode_batch_frames(params))
|
||||
|
||||
assert transfers == {1: [3, 3, 3, 3], 5: [6, 6], 33: [12]}[value]
|
||||
assert _decoded_frame_count(videos[0]) == 12
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", [0, -1, True, 1.5, "17", None])
|
||||
def test_request_batch_threshold_rejects_invalid_values(value):
|
||||
with pytest.raises(ValueError, match="preencode_batch_frames"):
|
||||
resolve_wan_preencode_batch_frames(SimpleNamespace(extra_args={"preencode_batch_frames": value}))
|
||||
|
||||
|
||||
def test_request_batch_threshold_preserves_model_defaults():
|
||||
params = SimpleNamespace(extra_args=None)
|
||||
assert resolve_wan_preencode_batch_frames(params) == 17
|
||||
assert resolve_wan_preencode_batch_frames(params, default=1) == 1
|
||||
@@ -647,3 +647,245 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
|
||||
assert video.shape[1] == 3 # channels
|
||||
assert audio_waveform is not None
|
||||
assert audio_sr == 16000
|
||||
|
||||
|
||||
def _make_s2v_preencode_pipeline() -> Wan22S2VPipeline:
|
||||
"""Build the same stub pipeline the batching test drives, for preencode runs."""
|
||||
pipeline = object.__new__(Wan22S2VPipeline)
|
||||
nn.Module.__init__(pipeline)
|
||||
pipeline.device = torch.device("cpu")
|
||||
pipeline.transformer = MagicMock()
|
||||
pipeline.transformer.dtype = torch.float32
|
||||
# A fresh iterator per call: the clip loop reads this once per clip.
|
||||
pipeline.transformer.parameters.side_effect = lambda: 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
|
||||
# The shared consumer reads the published pixel range off the VAE.
|
||||
pipeline.vae.chunk_value_range = (-1.0, 1.0)
|
||||
# One decoded clip: [B, C, T, H, W] for a batch of four (2 requests x 2 outputs).
|
||||
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._flow_shift = 3.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"])
|
||||
return pipeline
|
||||
|
||||
|
||||
def _make_s2v_preencode_batch(audio_a, audio_b, **sampling_overrides) -> DiffusionRequestBatch:
|
||||
image = PIL.Image.new("RGB", (16, 16))
|
||||
return DiffusionRequestBatch(
|
||||
requests=[
|
||||
SimpleNamespace(
|
||||
request_id="a",
|
||||
prompt={"prompt": "first", "multi_modal_data": {"image": image, "audio": audio_a}},
|
||||
sampling_params=_make_s2v_sampling(
|
||||
num_outputs_per_prompt=2,
|
||||
latents=torch.zeros(2, 16, 2, 2, 2),
|
||||
output_type="np",
|
||||
**sampling_overrides,
|
||||
),
|
||||
),
|
||||
SimpleNamespace(
|
||||
request_id="b",
|
||||
prompt={"prompt": "second", "multi_modal_data": {"image": image, "audio": audio_b}},
|
||||
sampling_params=_make_s2v_sampling(
|
||||
num_outputs_per_prompt=2,
|
||||
latents=torch.ones(2, 16, 2, 2, 2),
|
||||
output_type="np",
|
||||
**sampling_overrides,
|
||||
),
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def _decode_mp4(data: bytes):
|
||||
"""Return (video frame count, first audio samples) for one encoded container."""
|
||||
import io
|
||||
|
||||
import av
|
||||
|
||||
with av.open(io.BytesIO(data)) as container:
|
||||
frames = sum(1 for _ in container.decode(video=0))
|
||||
with av.open(io.BytesIO(data)) as container:
|
||||
assert container.streams.audio, "the worker must mux the waveform into the container"
|
||||
samples = np.concatenate(
|
||||
[frame.to_ndarray().reshape(-1) for frame in container.decode(audio=0)][:4],
|
||||
)
|
||||
return frames, samples
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_frames", [1, 1000])
|
||||
def test_s2v_preencode_returns_playable_mp4_bytes_per_request(monkeypatch, batch_frames) -> None:
|
||||
"""The clip loop hands finished clips to the encoder instead of concatenating."""
|
||||
pipeline = _make_s2v_preencode_pipeline()
|
||||
# Two clips, so the autoregressive motion feedback runs between pushes.
|
||||
pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 1, 2, 8), 2, 16))
|
||||
audio_a = np.zeros(16000, dtype=np.float32)
|
||||
audio_b = np.full(16000, 0.5, dtype=np.float32)
|
||||
batch = _make_s2v_preencode_batch(
|
||||
audio_a, audio_b, extra_args={"preencode_mp4": True, "preencode_batch_frames": batch_frames}
|
||||
)
|
||||
from vllm_omni.diffusion.utils import chunked_video
|
||||
|
||||
transfers = []
|
||||
quantize = chunked_video.chunk_to_uint8_frames
|
||||
|
||||
def capture(chunk, value_range):
|
||||
transfers.append(chunk.shape[2])
|
||||
return quantize(chunk, value_range)
|
||||
|
||||
monkeypatch.setattr(chunked_video, "chunk_to_uint8_frames", capture)
|
||||
|
||||
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 pipeline.vae.decode.call_count == 2, "both clips must reach the encoder"
|
||||
assert len(transfers) == (2 if batch_frames == 1 else 1)
|
||||
|
||||
decoded = []
|
||||
for request_output in outputs:
|
||||
# Two outputs per prompt, each already a complete container.
|
||||
assert isinstance(request_output.output, list)
|
||||
assert len(request_output.output) == 2
|
||||
for data in request_output.output:
|
||||
assert isinstance(data, bytes)
|
||||
frames, samples = _decode_mp4(data)
|
||||
assert frames > 0
|
||||
decoded.append(samples)
|
||||
|
||||
# Four entries: request a's waveform for its two outputs, then request b's.
|
||||
# A silent and a non-silent waveform must not land on the same request.
|
||||
assert np.abs(decoded[0]).max() == pytest.approx(0.0, abs=1e-3)
|
||||
assert np.abs(decoded[1]).max() == pytest.approx(0.0, abs=1e-3)
|
||||
assert np.abs(decoded[2]).max() > 0.1
|
||||
assert np.abs(decoded[3]).max() > 0.1
|
||||
|
||||
|
||||
def test_s2v_preencode_aborts_encoders_when_a_later_clip_fails(monkeypatch) -> None:
|
||||
"""A failure after the first push must not strand encoder worker threads."""
|
||||
pipeline = _make_s2v_preencode_pipeline()
|
||||
pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 1, 2, 8), 2, 16))
|
||||
batch = _make_s2v_preencode_batch(
|
||||
np.zeros(16000, dtype=np.float32),
|
||||
np.ones(16000, dtype=np.float32),
|
||||
extra_args={"preencode_mp4": True, "preencode_batch_frames": 1},
|
||||
)
|
||||
sessions = []
|
||||
|
||||
from vllm_omni.diffusion.utils.chunked_video import ChunkedVideoMP4Session
|
||||
|
||||
def capture_session(**kwargs):
|
||||
session = ChunkedVideoMP4Session(**kwargs)
|
||||
sessions.append(session)
|
||||
return session
|
||||
|
||||
def fail_on_second_clip(**kwargs):
|
||||
if pipeline.vae.decode.call_count:
|
||||
raise RuntimeError("second clip failed")
|
||||
return kwargs["latents"]
|
||||
|
||||
pipeline.diffuse = fail_on_second_clip
|
||||
monkeypatch.setattr("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.ChunkedVideoMP4Session", capture_session)
|
||||
|
||||
with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as platform:
|
||||
platform.is_available.return_value = False
|
||||
with pytest.raises(RuntimeError, match="second clip failed"):
|
||||
pipeline.forward(batch)
|
||||
|
||||
assert len(sessions) == 1
|
||||
assert len(sessions[0]._encoders) == 4
|
||||
assert all(not encoder._thread.is_alive() for encoder in sessions[0]._encoders)
|
||||
|
||||
|
||||
def test_s2v_preencode_skips_encoding_on_a_vae_patch_parallel_peer(monkeypatch) -> None:
|
||||
"""A peer still receives every clip for the motion loop but encodes none of them."""
|
||||
pipeline = _make_s2v_preencode_pipeline()
|
||||
pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 1, 2, 8), 2, 16))
|
||||
# A peer's patch-parallel decode returns an empty placeholder.
|
||||
pipeline.vae.decode.return_value = (torch.empty(0),)
|
||||
pipeline.vae._vae_pp_group = object()
|
||||
broadcasts = []
|
||||
monkeypatch.setattr(
|
||||
"torch.distributed.broadcast",
|
||||
lambda tensor, src, group: broadcasts.append(tensor.zero_()),
|
||||
)
|
||||
batch = _make_s2v_preencode_batch(
|
||||
np.zeros(16000, dtype=np.float32),
|
||||
np.ones(16000, dtype=np.float32),
|
||||
extra_args={"preencode_mp4": True, "preencode_batch_frames": 1},
|
||||
)
|
||||
sessions = []
|
||||
|
||||
from vllm_omni.diffusion.utils.chunked_video import ChunkedVideoMP4Session
|
||||
|
||||
def capture_session(**kwargs):
|
||||
session = ChunkedVideoMP4Session(**kwargs)
|
||||
sessions.append(session)
|
||||
return session
|
||||
|
||||
monkeypatch.setattr("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.ChunkedVideoMP4Session", capture_session)
|
||||
|
||||
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(broadcasts) == 2, "the peer must still join the broadcast for both clips"
|
||||
# Once for the initial motion latents, once more from the peer's broadcast clip.
|
||||
assert pipeline.prepare_motion_latents.call_count == 2, "the motion loop still consumes the peer's frames"
|
||||
assert sessions[0]._encoders == []
|
||||
assert [request_output.output for request_output in outputs] == [[], []]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("output_type", ["pil", "pt", "latent"])
|
||||
def test_s2v_preencode_rejects_request_output_types_it_cannot_serve(output_type) -> None:
|
||||
"""Pre-encoding returns MP4 bytes, so a request asking for frames must fail."""
|
||||
pipeline = _make_s2v_preencode_pipeline()
|
||||
batch = _make_s2v_preencode_batch(
|
||||
np.zeros(16000, dtype=np.float32),
|
||||
np.ones(16000, dtype=np.float32),
|
||||
extra_args={"preencode_mp4": True},
|
||||
)
|
||||
for request in batch.requests:
|
||||
request.sampling_params.output_type = output_type
|
||||
|
||||
with pytest.raises(ValueError, match="output_type"):
|
||||
pipeline.forward(batch)
|
||||
|
||||
|
||||
def test_s2v_preencode_keeps_the_full_decode_path_untouched() -> None:
|
||||
"""Without the flag the loop still returns the (video, audio, rate) tuple."""
|
||||
pipeline = _make_s2v_preencode_pipeline()
|
||||
audio_a = np.zeros(16000, dtype=np.float32)
|
||||
audio_b = np.ones(16000, dtype=np.float32)
|
||||
batch = _make_s2v_preencode_batch(audio_a, audio_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 outputs[0].output[0].shape[0] == 2
|
||||
np.testing.assert_array_equal(outputs[0].output[1], audio_a)
|
||||
np.testing.assert_array_equal(outputs[1].output[1], audio_b)
|
||||
Reference in new issue
Block a user