diff --git a/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py index 9d5b11d3b..12c0ebf45 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py +++ b/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py @@ -8,15 +8,28 @@ import torch from PIL import Image from torch import nn -from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar +from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import ( Wan22I2VPipeline, + get_wan22_i2v_post_process_func, get_wan22_i2v_pre_process_func, ) +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] +def test_wan22_i2v_postprocess_honors_request_output_type() -> None: + video = torch.zeros(1, 4, 1, 2, 2) + + output = get_wan22_i2v_post_process_func(SimpleNamespace())( + video, + sampling_params=SimpleNamespace(output_type="latent"), + ) + + assert output is video + + def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline: pipeline = object.__new__(Wan22I2VPipeline) nn.Module.__init__(pipeline) @@ -31,6 +44,29 @@ def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline: return pipeline +def _make_i2v_sampling(**overrides): + values = { + "height": 16, + "width": 16, + "num_frames": 5, + "num_inference_steps": 1, + "guidance_scale_provided": True, + "guidance_scale": 1.0, + "guidance_scale_2": None, + "guidance_scale_2_provided": False, + "boundary_ratio": None, + "generator": None, + "seed": None, + "num_outputs_per_prompt": 1, + "max_sequence_length": 8, + "latents": None, + "output_type": "latent", + "extra_args": {}, + } + values.update(overrides) + return SimpleNamespace(**values) + + def test_i2v_preprocess_requires_image_and_resizes_to_480p_aspect() -> None: preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace()) request = SimpleNamespace( @@ -53,6 +89,88 @@ def test_i2v_preprocess_requires_image_and_resizes_to_480p_aspect() -> None: preprocess(missing_image) +def _make_i2v_preprocess_request(last_image): + return SimpleNamespace( + prompt={ + "prompt": "p", + "multi_modal_data": { + "image": Image.new("RGB", (320, 160), "red"), + "last_image": last_image, + }, + }, + sampling_params=SimpleNamespace(height=16, width=16), + ) + + +def test_i2v_preprocess_treats_empty_last_image_list_as_absent() -> None: + preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace()) + + result = preprocess(_make_i2v_preprocess_request([])) + + assert result.prompt["multi_modal_data"]["last_image"] is None + assert result.batch_compatibility_key == ("wan22_i2v_last_image", False) + + +def test_i2v_preprocess_unwraps_single_last_image_list() -> None: + preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace()) + last_image = Image.new("RGB", (16, 16), "blue") + + result = preprocess(_make_i2v_preprocess_request([last_image])) + + assert result.prompt["multi_modal_data"]["last_image"] is last_image + assert result.batch_compatibility_key == ("wan22_i2v_last_image", True) + + +@pytest.mark.parametrize( + "last_image", + [ + Image.new("RGB", (16, 16), "blue"), + torch.zeros(3, 16, 16), + ], +) +def test_i2v_preprocess_preserves_supported_last_image(last_image) -> None: + preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace()) + + result = preprocess(_make_i2v_preprocess_request(last_image)) + + assert result.prompt["multi_modal_data"]["last_image"] is last_image + assert result.batch_compatibility_key == ("wan22_i2v_last_image", True) + + +def test_i2v_preprocess_loads_last_image_path(tmp_path) -> None: + preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace()) + path = tmp_path / "last.png" + Image.new("RGB", (16, 16), "blue").save(path) + + result = preprocess(_make_i2v_preprocess_request(str(path))) + + last_image = result.prompt["multi_modal_data"]["last_image"] + assert isinstance(last_image, Image.Image) + assert last_image.mode == "RGB" + assert result.batch_compatibility_key == ("wan22_i2v_last_image", True) + + +def test_i2v_preprocess_rejects_multiple_last_images() -> None: + preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace()) + + with pytest.raises(ValueError, match="at most one last_image"): + preprocess( + _make_i2v_preprocess_request( + [ + Image.new("RGB", (16, 16), "blue"), + Image.new("RGB", (16, 16), "green"), + ] + ) + ) + + +def test_i2v_preprocess_rejects_unsupported_last_image_type() -> None: + preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace()) + + with pytest.raises(TypeError, match="Unsupported last_image format"): + preprocess(_make_i2v_preprocess_request({"not": "an image"})) + + def test_i2v_diffuse_selects_stage_guidance_and_expands_timesteps() -> None: pipeline = _make_i2v_pipeline(expand_timesteps=True) latents = torch.zeros(1, 4, 2, 4, 4) @@ -126,3 +244,169 @@ def test_i2v_prepare_latents_builds_expand_condition_and_first_frame_mask() -> N assert first_frame_mask.shape == (1, 1, 2, 2, 2) assert first_frame_mask[:, :, 0].sum() == 0 assert first_frame_mask[:, :, 1].sum() == 4 + + +def test_i2v_prepare_latents_preserves_batched_image_conditions() -> None: + pipeline = _make_i2v_pipeline(expand_timesteps=True) + generators = [ + torch.Generator(device="cpu").manual_seed(1), + torch.Generator(device="cpu").manual_seed(2), + ] + + latents, condition, first_frame_mask = pipeline.prepare_latents( + image=torch.zeros(2, 3, 16, 16), + batch_size=2, + num_channels_latents=4, + height=16, + width=16, + num_frames=5, + dtype=torch.float32, + device=torch.device("cpu"), + generator=generators, + ) + + assert latents.shape == (2, 4, 2, 2, 2) + assert condition.shape == (2, 4, 1, 2, 2) + assert first_frame_mask.shape == (2, 1, 2, 2, 2) + + +def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch) -> None: + monkeypatch.setattr( + "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v.current_omni_platform", + SimpleNamespace(is_available=lambda: False), + ) + pipeline = _make_i2v_pipeline(expand_timesteps=True) + pipeline.scheduler = StubScheduler([9]) + pipeline.od_config = SimpleNamespace(flow_shift=5.0) + pipeline._sample_solver = "unipc" + pipeline._flow_shift = 5.0 + pipeline.boundary_ratio = 0.875 + pipeline.has_image_encoder = False + pipeline._guidance_scale = None + pipeline._guidance_scale_2 = None + pipeline._num_timesteps = None + pipeline._current_timestep = None + pipeline.check_inputs = lambda **kwargs: None + prepare_call = {} + + def _fake_encode_prompt(**kwargs): + batch_size = len(kwargs["prompt"]) + n = batch_size * kwargs["num_videos_per_prompt"] + return torch.zeros(n, 2, 3), None + + def _fake_prepare_latents(**kwargs): + prepare_call.update(kwargs) + batch_size = kwargs["batch_size"] + return ( + kwargs["latents"], + torch.zeros(batch_size, 4, 1, 2, 2), + torch.ones(batch_size, 1, 2, 2, 2), + ) + + pipeline.encode_prompt = _fake_encode_prompt # type: ignore[method-assign] + pipeline.prepare_latents = _fake_prepare_latents # type: ignore[method-assign] + pipeline.diffuse = lambda **kwargs: kwargs["latents"] # type: ignore[method-assign] + + gen_a = torch.Generator(device="cpu").manual_seed(1) + gen_b = torch.Generator(device="cpu").manual_seed(2) + latents_a = torch.zeros(2, 4, 2, 2, 2) + latents_b = torch.ones(2, 4, 2, 2, 2) + image_a = torch.zeros(1, 3, 16, 16) + image_b = torch.ones(1, 3, 16, 16) + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="a", + prompt={"prompt": "first", "multi_modal_data": {"image": image_a}}, + sampling_params=_make_i2v_sampling( + generator=gen_a, + latents=latents_a, + num_outputs_per_prompt=2, + ), + ), + SimpleNamespace( + request_id="b", + prompt={"prompt": "second", "multi_modal_data": {"image": image_b}}, + sampling_params=_make_i2v_sampling( + generator=gen_b, + latents=latents_b, + num_outputs_per_prompt=2, + ), + ), + ] + ) + + outputs = pipeline.forward(batch) + + assert prepare_call["batch_size"] == 4 + assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b] + torch.testing.assert_close(prepare_call["latents"], torch.cat([latents_a, latents_b])) + torch.testing.assert_close( + prepare_call["image"], + torch.cat([image_a, image_a, image_b, image_b]), + ) + assert len(outputs) == 2 + torch.testing.assert_close(outputs[0].output, latents_a) + torch.testing.assert_close(outputs[1].output, latents_b) + + +def test_i2v_forward_rejects_mismatched_tensor_condition_shapes(monkeypatch) -> None: + monkeypatch.setattr( + "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v.current_omni_platform", + SimpleNamespace(is_available=lambda: False), + ) + pipeline = _make_i2v_pipeline(expand_timesteps=True) + pipeline.scheduler = StubScheduler([9]) + pipeline.od_config = SimpleNamespace(flow_shift=5.0) + pipeline._sample_solver = "unipc" + pipeline._flow_shift = 5.0 + pipeline.boundary_ratio = 0.875 + pipeline.has_image_encoder = False + pipeline._guidance_scale = None + pipeline._guidance_scale_2 = None + pipeline._num_timesteps = None + pipeline._current_timestep = None + pipeline.check_inputs = lambda **kwargs: None + pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None) # type: ignore[method-assign] + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="a", + prompt={"prompt": "first", "multi_modal_data": {"image": torch.zeros(1, 3, 16, 16)}}, + sampling_params=_make_i2v_sampling(), + ), + SimpleNamespace( + request_id="b", + prompt={"prompt": "second", "multi_modal_data": {"image": torch.zeros(1, 3, 16, 8)}}, + sampling_params=_make_i2v_sampling(), + ), + ] + ) + + with pytest.raises(ValueError, match="image condition"): + pipeline.forward(batch) + + +def test_i2v_forward_rejects_mixed_last_image_presence() -> None: + pipeline = _make_i2v_pipeline(expand_timesteps=True) + image = torch.zeros(1, 3, 16, 16) + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="a", + prompt={ + "prompt": "first", + "multi_modal_data": {"image": image, "last_image": image}, + }, + sampling_params=_make_i2v_sampling(), + ), + SimpleNamespace( + request_id="b", + prompt={"prompt": "second", "multi_modal_data": {"image": image}}, + sampling_params=_make_i2v_sampling(), + ), + ] + ) + + with pytest.raises(ValueError, match="mix of provided and missing last_image"): + pipeline.forward(batch) diff --git a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py index 8c06df034..325bd59da 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py +++ b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py @@ -93,6 +93,29 @@ def _make_pipeline() -> Wan22Pipeline: return pipeline +def _make_sampling(**overrides): + values = { + "height": None, + "width": None, + "num_frames": 1, + "num_inference_steps": 2, + "guidance_scale_provided": True, + "guidance_scale": 1.0, + "guidance_scale_2": None, + "guidance_scale_2_provided": False, + "boundary_ratio": None, + "generator": None, + "seed": None, + "num_outputs_per_prompt": 1, + "max_sequence_length": 32, + "latents": None, + "output_type": "latent", + "extra_args": {}, + } + values.update(overrides) + return SimpleNamespace(**values) + + @pytest.mark.parametrize( ("sampling_params_kwargs", "expected_low", "expected_high"), [ @@ -129,9 +152,10 @@ def test_forward_delegates_denoising_to_diffuse( ) batch = DiffusionRequestBatch(requests=[mock_req]) - output = pipeline.forward(batch) + outputs = pipeline.forward(batch) - assert torch.equal(output.output, torch.ones((1, 4, 1, 8, 8))) + assert len(outputs) == 1 + assert torch.equal(outputs[0].output, torch.ones((1, 4, 1, 8, 8))) assert torch.equal(captured["prompt_embeds"], torch.zeros(1, 32, 8)) assert torch.equal(captured["timesteps"], pipeline.scheduler.timesteps) assert captured["guidance_low"] == expected_low @@ -142,6 +166,134 @@ def test_forward_delegates_denoising_to_diffuse( assert pipeline.scheduler.set_timesteps_calls == [(2, torch.device("cpu"))] +def test_forward_batches_text_generators_latents_and_splits_outputs() -> None: + pipeline = _make_pipeline() + encode_call = {} + prepare_call = {} + + def _fake_encode_prompt(**kwargs): + encode_call.update(kwargs) + batch_size = len(kwargs["prompt"]) + n = batch_size * kwargs["num_videos_per_prompt"] + return torch.arange(n, dtype=torch.float32).view(n, 1, 1), torch.zeros(n, 1, 1) + + def _fake_prepare_latents(**kwargs): + prepare_call.update(kwargs) + return kwargs["latents"] + + pipeline.encode_prompt = _fake_encode_prompt # type: ignore[method-assign] + pipeline.prepare_latents = _fake_prepare_latents # type: ignore[method-assign] + pipeline.diffuse = lambda **kwargs: kwargs["latents"] # type: ignore[method-assign] + + gen_a = torch.Generator(device="cpu").manual_seed(1) + gen_b = torch.Generator(device="cpu").manual_seed(2) + latents_a = torch.zeros(2, 4, 1, 2, 2) + latents_b = torch.ones(2, 4, 1, 2, 2) + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="a", + prompt={"prompt": "first", "negative_prompt": "bad first"}, + sampling_params=_make_sampling( + generator=gen_a, + latents=latents_a, + num_outputs_per_prompt=2, + ), + ), + SimpleNamespace( + request_id="b", + prompt={"prompt": "second", "negative_prompt": "bad second"}, + sampling_params=_make_sampling( + generator=gen_b, + latents=latents_b, + num_outputs_per_prompt=2, + ), + ), + ] + ) + + outputs = pipeline.forward(batch) + + assert encode_call["prompt"] == ["first", "second"] + assert encode_call["negative_prompt"] == ["bad first", "bad second"] + assert prepare_call["batch_size"] == 4 + assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b] + torch.testing.assert_close(prepare_call["latents"], torch.cat([latents_a, latents_b])) + assert len(outputs) == 2 + torch.testing.assert_close(outputs[0].output, latents_a) + torch.testing.assert_close(outputs[1].output, latents_b) + + +def test_forward_batches_precomputed_prompt_embeddings() -> None: + pipeline = _make_pipeline() + diffuse_call = {} + pipeline.encode_prompt = lambda **kwargs: pytest.fail("text encoder must not run") # type: ignore[method-assign] + pipeline.prepare_latents = lambda **kwargs: torch.zeros(kwargs["batch_size"], 4, 1, 2, 2) # type: ignore[method-assign] + + def _fake_diffuse(**kwargs): + diffuse_call.update(kwargs) + return kwargs["latents"] + + pipeline.diffuse = _fake_diffuse # type: ignore[method-assign] + embeds_a = torch.zeros(3, 4) + embeds_b = torch.ones(3, 4) + negative_a = torch.full((3, 4), 2.0) + negative_b = torch.full((3, 4), 3.0) + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="a", + prompt={"prompt_embeds": embeds_a, "negative_prompt_embeds": negative_a}, + sampling_params=_make_sampling(guidance_scale=4.0), + ), + SimpleNamespace( + request_id="b", + prompt={"prompt_embeds": embeds_b, "negative_prompt_embeds": negative_b}, + sampling_params=_make_sampling(guidance_scale=4.0), + ), + ] + ) + + outputs = pipeline.forward(batch) + + torch.testing.assert_close(diffuse_call["prompt_embeds"], torch.stack([embeds_a, embeds_b])) + torch.testing.assert_close(diffuse_call["negative_prompt_embeds"], torch.stack([negative_a, negative_b])) + assert len(outputs) == 2 + + +def test_prepare_latents_with_request_generators_matches_single_generation() -> None: + pipeline = _make_pipeline() + kwargs = { + "num_channels_latents": 4, + "height": 16, + "width": 16, + "num_frames": 5, + "dtype": torch.float32, + "device": torch.device("cpu"), + } + + batched = Wan22Pipeline.prepare_latents( + pipeline, + batch_size=2, + generator=[torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)], + **kwargs, + ) + singles = torch.cat( + [ + Wan22Pipeline.prepare_latents( + pipeline, + batch_size=1, + generator=torch.Generator().manual_seed(seed), + **kwargs, + ) + for seed in (1, 2) + ] + ) + + torch.testing.assert_close(batched, singles) + assert not torch.equal(batched[0], batched[1]) + + def test_diffuse_runs_prediction_and_scheduler_for_each_timestep() -> None: pipeline = _make_pipeline() latents = torch.zeros((1, 1, 1, 2, 2), dtype=torch.float32) diff --git a/tests/diffusion/models/wan2_2/test_wan22_pipeline_helpers.py b/tests/diffusion/models/wan2_2/test_wan22_pipeline_helpers.py index 314717869..5a37e14f2 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_pipeline_helpers.py +++ b/tests/diffusion/models/wan2_2/test_wan22_pipeline_helpers.py @@ -10,6 +10,7 @@ import torch import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 as wan22_module from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( create_transformer_from_config, + get_wan22_post_process_func, load_transformer_config, retrieve_latents, ) @@ -26,6 +27,17 @@ class _LatentDist: return torch.tensor([2.0]) +def test_wan22_postprocess_honors_request_output_type() -> None: + video = torch.zeros(1, 4, 1, 2, 2) + + output = get_wan22_post_process_func(SimpleNamespace())( + video, + sampling_params=SimpleNamespace(output_type="latent"), + ) + + assert output is video + + def test_retrieve_latents_supports_sample_mode_argmax_and_direct_latents() -> None: generator = torch.Generator(device="cpu") diff --git a/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py index 4440a8f35..e5bae65df 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py +++ b/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py @@ -1,14 +1,307 @@ +from types import SimpleNamespace from unittest.mock import MagicMock, patch +import numpy as np +import PIL.Image import pytest import torch import torch.nn as nn +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import ( + Wan22S2VPipeline, + _make_clip_generators, +) from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import WanS2VTransformer3DModel +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +def _make_s2v_sampling(**overrides): + values = { + "height": 16, + "width": 16, + "num_frames": 8, + "num_inference_steps": 1, + "guidance_scale_provided": True, + "guidance_scale": 1.0, + "generator": None, + "seed": None, + "num_outputs_per_prompt": 1, + "max_sequence_length": 8, + "latents": None, + "output_type": "latent", + "extra_args": {}, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _make_s2v_validation_pipeline() -> Wan22S2VPipeline: + pipeline = object.__new__(Wan22S2VPipeline) + nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.transformer = SimpleNamespace(dtype=torch.float32) + 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.check_inputs = lambda *args, **kwargs: None + pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None) # type: ignore[method-assign] + return pipeline + + +def test_s2v_predict_noise_keeps_batch_dimension() -> None: + pipeline = object.__new__(Wan22S2VPipeline) + nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + transformer = MagicMock() + transformer.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) + transformer.return_value = (torch.zeros(1, 16, 2, 2, 2),) + pipeline.transformer = transformer + + result = pipeline.predict_noise(hidden_states=torch.zeros(1, 16, 2, 2, 2)) + + assert result.shape == (1, 16, 2, 2, 2) + + +def test_s2v_clip_generators_preserve_main_seed_per_clip_behavior() -> None: + explicit_generator = torch.Generator(device="cpu").manual_seed(999) + + generators = _make_clip_generators( + seeds=[1234, 5678, None, None], + request_generators=[ + torch.Generator(device="cpu").manual_seed(1234), + torch.Generator(device="cpu").manual_seed(5678), + explicit_generator, + None, + ], + clip_index=2, + device=torch.device("cpu"), + ) + + assert [generator.initial_seed() for generator in generators if generator is not None] == [1236, 5680, 999, 2] + assert generators[2] is explicit_generator + + +def test_s2v_forward_rejects_different_num_repeat() -> None: + pipeline = _make_s2v_validation_pipeline() + pipeline.encode_audio = lambda *args, **kwargs: (torch.zeros(1, 1, 2, 16), 2, 16) # type: ignore[method-assign] + image = PIL.Image.new("RGB", (16, 16)) + audio = np.zeros(16000, dtype=np.float32) + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="a", + prompt={ + "prompt": "first", + "multi_modal_data": {"image": image, "audio": audio}, + "additional_information": {"num_repeat": 1}, + }, + sampling_params=_make_s2v_sampling(), + ), + SimpleNamespace( + request_id="b", + prompt={ + "prompt": "second", + "multi_modal_data": {"image": image, "audio": audio}, + "additional_information": {"num_repeat": 2}, + }, + sampling_params=_make_s2v_sampling(), + ), + ] + ) + + with pytest.raises(ValueError, match="num_repeat"): + pipeline.forward(batch) + + +def test_s2v_forward_rejects_different_audio_time_lengths() -> None: + pipeline = _make_s2v_validation_pipeline() + pipeline.encode_audio = lambda audio, **kwargs: ( # type: ignore[method-assign] + torch.zeros(1, 1, 2, 16 if audio[0] == 0 else 24), + 1, + 16 if audio[0] == 0 else 24, + ) + image = PIL.Image.new("RGB", (16, 16)) + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="a", + prompt={ + "prompt": "first", + "multi_modal_data": {"image": image, "audio": np.zeros(16000, dtype=np.float32)}, + }, + sampling_params=_make_s2v_sampling(), + ), + SimpleNamespace( + request_id="b", + prompt={ + "prompt": "second", + "multi_modal_data": {"image": image, "audio": np.ones(16000, dtype=np.float32)}, + }, + sampling_params=_make_s2v_sampling(), + ), + ] + ) + + with pytest.raises(ValueError, match="audio embedding time lengths"): + pipeline.forward(batch) + + +def test_s2v_forward_rejects_different_init_first_frame() -> None: + pipeline = _make_s2v_validation_pipeline() + image = PIL.Image.new("RGB", (16, 16)) + audio = np.zeros(16000, dtype=np.float32) + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="a", + prompt={ + "prompt": "first", + "multi_modal_data": {"image": image, "audio": audio}, + "additional_information": {"init_first_frame": True}, + }, + sampling_params=_make_s2v_sampling(), + ), + SimpleNamespace( + request_id="b", + prompt={ + "prompt": "second", + "multi_modal_data": {"image": image, "audio": audio}, + "additional_information": {"init_first_frame": False}, + }, + sampling_params=_make_s2v_sampling(), + ), + ] + ) + + with pytest.raises(ValueError, match="init_first_frame"): + pipeline.forward(batch) + + +def test_s2v_forward_rejects_different_raw_audio_shapes() -> None: + pipeline = _make_s2v_validation_pipeline() + pipeline.encode_audio = lambda *args, **kwargs: (torch.zeros(1, 1, 2, 16), 1, 16) # type: ignore[method-assign] + image = PIL.Image.new("RGB", (16, 16)) + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="a", + prompt={ + "prompt": "first", + "multi_modal_data": {"image": image, "audio": np.zeros(16000, dtype=np.float32)}, + }, + sampling_params=_make_s2v_sampling(), + ), + SimpleNamespace( + request_id="b", + prompt={ + "prompt": "second", + "multi_modal_data": {"image": image, "audio": np.zeros(8000, dtype=np.float32)}, + }, + sampling_params=_make_s2v_sampling(), + ), + ] + ) + + with pytest.raises(ValueError, match="raw audio shapes and sample rates"): + pipeline.forward(batch) + + +def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None: + pipeline = object.__new__(Wan22S2VPipeline) + nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.transformer = MagicMock() + pipeline.transformer.dtype = torch.float32 + pipeline.transformer.parameters.return_value = 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 + 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.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"]) + + image = PIL.Image.new("RGB", (16, 16)) + audio_a = np.zeros(16000, dtype=np.float32) + audio_b = np.ones(16000, dtype=np.float32) + latents_a = torch.zeros(2, 16, 2, 2, 2) + latents_b = torch.ones(2, 16, 2, 2, 2) + generators_a = [torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)] + generators_b = [torch.Generator().manual_seed(3), torch.Generator().manual_seed(4)] + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="a", + prompt={ + "prompt": "first", + "negative_prompt": "negative first", + "multi_modal_data": {"image": image, "audio": audio_a}, + }, + sampling_params=_make_s2v_sampling( + num_outputs_per_prompt=2, + generator=generators_a, + latents=latents_a, + ), + ), + SimpleNamespace( + request_id="b", + prompt={ + "prompt": "second", + "negative_prompt": "negative second", + "multi_modal_data": {"image": image, "audio": audio_b}, + }, + sampling_params=_make_s2v_sampling( + num_outputs_per_prompt=2, + generator=generators_b, + latents=latents_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 len(outputs) == 2 + assert outputs[0].output[0].shape[0] == 2 + assert outputs[1].output[0].shape[0] == 2 + assert outputs[0].output[1].shape == (16000,) + assert outputs[1].output[1].shape == (16000,) + np.testing.assert_array_equal(outputs[0].output[1], audio_a) + np.testing.assert_array_equal(outputs[1].output[1], audio_b) + assert np.shares_memory(outputs[0].output[1], audio_a) + assert np.shares_memory(outputs[1].output[1], audio_b) + torch.testing.assert_close(pipeline.diffuse.call_args.kwargs["latents"], torch.cat([latents_a, latents_b])) + assert pipeline.diffuse.call_args.kwargs["clip_generator"] == generators_a + generators_b + assert pipeline.encode_prompt.call_args.kwargs["prompt"] == ["first", "second"] + assert pipeline.encode_prompt.call_args.kwargs["negative_prompt"] == ["negative first", "negative second"] + + def test_s2v_exposes_hsdp_shard_conditions_for_transformer_blocks(): model = object.__new__(WanS2VTransformer3DModel) nn.Module.__init__(model) @@ -260,31 +553,37 @@ def test_s2v_pipeline_hsdp_forward_complete_process(): pipeline._prompt_clean = Wan22S2VPipeline._prompt_clean # -- Build request -- - from vllm_omni.diffusion.request import OmniDiffusionRequest - - sampling_params = MagicMock() - sampling_params.height = 704 - sampling_params.width = 1024 - sampling_params.num_frames = 80 - sampling_params.num_inference_steps = 5 - sampling_params.guidance_scale = 4.5 - sampling_params.guidance_scale_provided = True - sampling_params.generator = None - sampling_params.seed = 42 + sampling_params = _make_s2v_sampling( + height=704, + width=1024, + num_frames=80, + num_inference_steps=5, + guidance_scale=4.5, + generator=torch.Generator(device="cpu").manual_seed(42), + max_sequence_length=512, + ) ref_image = PIL.Image.new("RGB", (1024, 704)) audio_data = np.zeros(16000, dtype=np.float32) - req = MagicMock(spec=OmniDiffusionRequest) - req.prompts = [ - { - "prompt": "test prompt", - "negative_prompt": "bad quality", - "multi_modal_data": {"image": ref_image, "audio": audio_data}, - "additional_information": {"audio_path": audio_data, "pose_video": None, "init_first_frame": False}, - } - ] - req.sampling_params = sampling_params + req = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="test", + prompt={ + "prompt": "test prompt", + "negative_prompt": "bad quality", + "multi_modal_data": {"image": ref_image, "audio": audio_data}, + "additional_information": { + "audio_path": audio_data, + "pose_video": None, + "init_first_frame": False, + }, + }, + sampling_params=sampling_params, + ) + ] + ) # -- Mock methods that use complex internal state -- pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 25, 64, 80), 1, 80)) @@ -295,7 +594,7 @@ def test_s2v_pipeline_hsdp_forward_complete_process(): pipeline.progress_bar.return_value.__exit__ = MagicMock(return_value=False) # Mock predict_noise_maybe_with_cfg from CFGParallelMixin - pipeline.predict_noise_maybe_with_cfg = MagicMock(return_value=torch.zeros(16, 20, 88, 128)) + pipeline.predict_noise_maybe_with_cfg = MagicMock(return_value=torch.zeros(1, 16, 20, 88, 128)) pipeline.predict_noise = Wan22S2VPipeline.predict_noise.__get__(pipeline) # Mock platform methods @@ -306,7 +605,7 @@ def test_s2v_pipeline_hsdp_forward_complete_process(): with patch( "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.load_audio", return_value=(audio_data, 16000) ): - result = Wan22S2VPipeline.forward(pipeline, req=req) + result = Wan22S2VPipeline.forward(pipeline, req=req)[0] # -- Assertions -- # Text encoder was called diff --git a/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py index 714b79d50..6534b7cff 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py +++ b/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py @@ -8,13 +8,14 @@ import torch from PIL import Image from torch import nn -from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar +from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace import ( Wan22VACEPipeline, create_vace_transformer_from_config, get_wan22_vace_pre_process_func, ) from vllm_omni.diffusion.models.wan2_2.wan2_2_vace_transformer import WanVACETransformer3DModel +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @@ -32,6 +33,27 @@ def _make_vace_pipeline() -> Wan22VACEPipeline: return pipeline +def _make_vace_sampling(**overrides): + values = { + "height": 16, + "width": 16, + "num_frames": 5, + "num_inference_steps": 1, + "guidance_scale_provided": True, + "guidance_scale": 1.0, + "boundary_ratio": None, + "generator": None, + "seed": None, + "num_outputs_per_prompt": 1, + "max_sequence_length": 8, + "latents": None, + "output_type": "latent", + "extra_args": {}, + } + values.update(overrides) + return SimpleNamespace(**values) + + def test_vace_preprocess_collects_reference_video_and_mask_inputs() -> None: preprocess = get_wan22_vace_pre_process_func(SimpleNamespace()) ref = Image.new("RGB", (320, 160), "green") @@ -138,3 +160,189 @@ def test_vace_diffuse_passes_context_and_scale_to_cfg_branches() -> None: assert calls[0]["positive_kwargs"]["vace_context"] is vace_context assert calls[0]["negative_kwargs"]["vace_context_scale"] == 0.75 torch.testing.assert_close(result, torch.ones_like(latents)) + + +def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> None: + monkeypatch.setattr( + "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform", + SimpleNamespace(is_available=lambda: False), + ) + pipeline = _make_vace_pipeline() + pipeline.transformer.vace_patch_embedding = None + pipeline.transformer_2 = None + pipeline.scheduler = StubScheduler([9]) + pipeline.od_config = SimpleNamespace(flow_shift=3.0) + pipeline._sample_solver = "unipc" + pipeline._flow_shift = 3.0 + pipeline.boundary_ratio = None + pipeline._guidance_scale = None + pipeline._num_timesteps = None + pipeline._current_timestep = None + pipeline.check_inputs = lambda **kwargs: None + prepare_call = {} + + def _fake_encode_prompt(**kwargs): + batch_size = len(kwargs["prompt"]) + n = batch_size * kwargs["num_videos_per_prompt"] + return torch.zeros(n, 2, 3), None + + def _fake_prepare_latents(**kwargs): + prepare_call.update(kwargs) + return kwargs["latents"] + + pipeline.encode_prompt = _fake_encode_prompt # type: ignore[method-assign] + pipeline.prepare_latents = _fake_prepare_latents # type: ignore[method-assign] + pipeline.diffuse = lambda **kwargs: kwargs["latents"] # type: ignore[method-assign] + gen_a = torch.Generator(device="cpu").manual_seed(1) + gen_b = torch.Generator(device="cpu").manual_seed(2) + latents_a = torch.zeros(2, 4, 2, 2, 2) + latents_b = torch.ones(2, 4, 2, 2, 2) + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="a", + prompt={"prompt": "first"}, + sampling_params=_make_vace_sampling( + generator=gen_a, + latents=latents_a, + num_outputs_per_prompt=2, + ), + ), + SimpleNamespace( + request_id="b", + prompt={"prompt": "second"}, + sampling_params=_make_vace_sampling( + generator=gen_b, + latents=latents_b, + num_outputs_per_prompt=2, + ), + ), + ] + ) + + outputs = pipeline.forward(batch) + + assert prepare_call["batch_size"] == 4 + assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b] + torch.testing.assert_close(prepare_call["latents"], torch.cat([latents_a, latents_b])) + assert len(outputs) == 2 + torch.testing.assert_close(outputs[0].output, latents_a) + torch.testing.assert_close(outputs[1].output, latents_b) + + +def test_vace_forward_rejects_mismatched_condition_shapes(monkeypatch) -> None: + monkeypatch.setattr( + "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform", + SimpleNamespace(is_available=lambda: False), + ) + pipeline = _make_vace_pipeline() + pipeline.transformer.vace_patch_embedding = object() + pipeline.transformer_2 = None + pipeline.scheduler = StubScheduler([9]) + pipeline.od_config = SimpleNamespace(flow_shift=3.0) + pipeline._sample_solver = "unipc" + pipeline._flow_shift = 3.0 + pipeline.boundary_ratio = None + pipeline._guidance_scale = None + pipeline._num_timesteps = None + pipeline._current_timestep = None + pipeline.check_inputs = lambda **kwargs: None + pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None) # type: ignore[method-assign] + + def _fake_preprocess(video, **kwargs): + del kwargs + mask = torch.ones_like(video) + return video, mask, [[]] + + pipeline.preprocess_conditions = _fake_preprocess # type: ignore[method-assign] + pipeline.prepare_video_latents = ( # type: ignore[method-assign] + lambda video, mask, refs, generator, device: torch.zeros(1, 8, video.shape[2], 2, 2) + ) + pipeline.prepare_masks = lambda mask, refs: torch.zeros(1, 4, mask.shape[2], 2, 2) # type: ignore[method-assign] + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="a", + prompt={ + "prompt": "first", + "additional_information": {"source_video": torch.zeros(1, 3, 5, 16, 16)}, + }, + sampling_params=_make_vace_sampling(), + ), + SimpleNamespace( + request_id="b", + prompt={ + "prompt": "second", + "additional_information": {"source_video": torch.zeros(1, 3, 3, 16, 16)}, + }, + sampling_params=_make_vace_sampling(), + ), + ] + ) + + with pytest.raises(ValueError, match="VACE condition"): + pipeline.forward(batch) + + +def test_vace_forward_rejects_mismatched_reference_image_counts(monkeypatch) -> None: + monkeypatch.setattr( + "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform", + SimpleNamespace(is_available=lambda: False), + ) + pipeline = _make_vace_pipeline() + pipeline.transformer.vace_patch_embedding = object() + pipeline.transformer_2 = None + pipeline.scheduler = StubScheduler([9]) + pipeline.od_config = SimpleNamespace(flow_shift=3.0) + pipeline._sample_solver = "unipc" + pipeline._flow_shift = 3.0 + pipeline.boundary_ratio = None + pipeline._guidance_scale = None + pipeline._num_timesteps = None + pipeline._current_timestep = None + pipeline.check_inputs = lambda **kwargs: None + pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None) # type: ignore[method-assign] + + def _fake_preprocess(video, mask, reference_images, **kwargs): + del kwargs + return video, mask, [reference_images] + + pipeline.preprocess_conditions = _fake_preprocess # type: ignore[method-assign] + pipeline.prepare_video_latents = ( # type: ignore[method-assign] + lambda video, mask, refs, generator, device: torch.zeros(1, 8, 5, 2, 2) + ) + pipeline.prepare_masks = lambda mask, refs: torch.zeros(1, 4, 5, 2, 2) # type: ignore[method-assign] + video = torch.zeros(1, 3, 5, 16, 16) + mask = torch.ones_like(video) + reference = torch.zeros(3, 16, 16) + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="a", + prompt={ + "prompt": "first", + "additional_information": { + "source_video": video, + "mask": mask, + "reference_images": [reference], + }, + }, + sampling_params=_make_vace_sampling(), + ), + SimpleNamespace( + request_id="b", + prompt={ + "prompt": "second", + "additional_information": { + "source_video": video, + "mask": mask, + "reference_images": [reference, reference], + }, + }, + sampling_params=_make_vace_sampling(), + ), + ] + ) + + with pytest.raises(ValueError, match="same number of reference images"): + pipeline.forward(batch) diff --git a/tests/diffusion/test_diffusion_engine_rpc_routing.py b/tests/diffusion/test_diffusion_engine_rpc_routing.py index 1d4c2d17e..5aef831aa 100644 --- a/tests/diffusion/test_diffusion_engine_rpc_routing.py +++ b/tests/diffusion/test_diffusion_engine_rpc_routing.py @@ -123,6 +123,7 @@ class _ConcurrencyTrackingExecutor: def _make_request(tag: str): sampling_params = dict(_SAMPLING_KEY_DEFAULTS) sampling_params["num_inference_steps"] = 1 + sampling_params["extra_args"] = {} return SimpleNamespace( request_id=tag, prompt=f"prompt_{tag}", diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py index ce9f0ebe1..46ec0cbe6 100644 --- a/tests/diffusion/test_diffusion_scheduler.py +++ b/tests/diffusion/test_diffusion_scheduler.py @@ -217,13 +217,21 @@ class TestGetRequestBatchSamplingParamsKey: num_inference_steps: int = 2, seed: int | None = 123, generator: torch.Generator | None = None, + extra_args: dict | None = None, + condition_key: tuple | None = None, ) -> OmniDiffusionRequest: sp = OmniDiffusionSamplingParams( num_inference_steps=num_inference_steps, seed=seed, generator=generator, + extra_args=extra_args or {}, + ) + return OmniDiffusionRequest( + prompt="prompt", + sampling_params=sp, + request_id=f"req-{num_inference_steps}", + batch_compatibility_key=condition_key, ) - return OmniDiffusionRequest(prompt="prompt", sampling_params=sp, request_id=f"req-{num_inference_steps}") def test_distinguishes_num_inference_steps(self) -> None: scheduler = RequestScheduler() @@ -240,6 +248,57 @@ class TestGetRequestBatchSamplingParamsKey: self._make(seed=1, generator=gen_a) ) == scheduler._build_sampling_params_key(self._make(seed=2, generator=gen_b)) + @pytest.mark.parametrize( + ("first_extra_args", "second_extra_args"), + [ + ({"sample_solver": "unipc"}, {"sample_solver": "euler"}), + ({"flow_shift": 3.0}, {"flow_shift": 5.0}), + ], + ) + def test_distinguishes_wan_scheduler_structure( + self, + first_extra_args: dict, + second_extra_args: dict, + ) -> None: + scheduler = RequestScheduler() + + assert scheduler._build_sampling_params_key( + self._make(extra_args=first_extra_args) + ) != scheduler._build_sampling_params_key(self._make(extra_args=second_extra_args)) + + @pytest.mark.parametrize( + ("first_extra_args", "second_extra_args"), + [ + ({"sample_solver": " Euler "}, {"sample_solver": "euler"}), + ({"flow_shift": "5.0"}, {"flow_shift": 5.0}), + ], + ) + def test_normalizes_equivalent_wan_scheduler_structure( + self, + first_extra_args: dict, + second_extra_args: dict, + ) -> None: + scheduler = RequestScheduler() + + assert scheduler._build_sampling_params_key( + self._make(extra_args=first_extra_args) + ) == scheduler._build_sampling_params_key(self._make(extra_args=second_extra_args)) + + def test_uses_none_for_unspecified_wan_scheduler_structure(self) -> None: + scheduler = RequestScheduler() + + key = scheduler._build_sampling_params_key(self._make()) + + assert key.sample_solver is None + assert key.flow_shift is None + + def test_distinguishes_pipeline_condition_structure(self) -> None: + scheduler = RequestScheduler() + + assert scheduler._build_sampling_params_key( + self._make(condition_key=("wan22_s2v_condition", True)) + ) != scheduler._build_sampling_params_key(self._make(condition_key=("wan22_s2v_condition", False))) + class TestRequestScheduler: def setup_method(self) -> None: @@ -602,6 +661,38 @@ class TestRequestScheduler: assert first.num_running_reqs == 1 assert first.num_waiting_reqs == 1 + def test_batches_incompatible_pipeline_conditions_separately(self) -> None: + scheduler = RequestScheduler() + scheduler.initialize(SimpleNamespace(max_num_seqs=2)) + + request_a = OmniDiffusionRequest( + prompt="a", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, seed=123), + request_id="a", + batch_compatibility_key=("wan22_s2v_condition", True), + ) + request_b = OmniDiffusionRequest( + prompt="b", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, seed=456), + request_id="b", + batch_compatibility_key=("wan22_s2v_condition", False), + ) + scheduler.add_request(request_a) + scheduler.add_request(request_b) + + first = scheduler.schedule() + + assert _new_ids(first) == ["a"] + assert first.num_running_reqs == 1 + assert first.num_waiting_reqs == 1 + + scheduler.update_from_output(first, _make_request_output("a")) + second = scheduler.schedule() + + assert _new_ids(second) == ["b"] + assert second.num_running_reqs == 1 + assert second.num_waiting_reqs == 0 + def test_batches_different_request_local_seed_together(self) -> None: scheduler = RequestScheduler() scheduler.initialize(SimpleNamespace(max_num_seqs=2)) @@ -625,6 +716,47 @@ class TestRequestScheduler: assert first.num_running_reqs == 2 assert first.num_waiting_reqs == 0 + @pytest.mark.parametrize( + ("first_extra_args", "second_extra_args"), + [ + ({"sample_solver": "unipc"}, {"sample_solver": "euler"}), + ({"flow_shift": 3.0}, {"flow_shift": 5.0}), + ({"sample_solver": "unipc"}, {}), + ({"flow_shift": 3.0}, {}), + ], + ) + def test_batches_incompatible_wan_scheduler_structure_separately( + self, + first_extra_args: dict, + second_extra_args: dict, + ) -> None: + scheduler = RequestScheduler() + scheduler.initialize(SimpleNamespace(max_num_seqs=2)) + first_id = scheduler.add_request( + _make_step_request( + "a", + sampling_params=OmniDiffusionSamplingParams( + num_inference_steps=2, + extra_args=first_extra_args, + ), + ) + ) + scheduler.add_request( + _make_step_request( + "b", + sampling_params=OmniDiffusionSamplingParams( + num_inference_steps=2, + extra_args=second_extra_args, + ), + ) + ) + + first = scheduler.schedule() + + assert _new_ids(first) == [first_id] + assert first.num_running_reqs == 1 + assert first.num_waiting_reqs == 1 + def test_incompatible_waiting_head_blocks_later_compatible_request(self) -> None: scheduler = RequestScheduler() scheduler.initialize(SimpleNamespace(max_num_seqs=3)) diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 9ee59f153..5fae5f1b2 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -399,6 +399,101 @@ def test_initialize_local_llm_replica_scopes_runtime_env(monkeypatch): assert os.environ[runtime_env_var] == "parent-value" +def test_initialize_diffusion_stage_applies_client_batch_size_to_engine(monkeypatch): + import vllm_omni.diffusion.stage_diffusion_client as client_mod + import vllm_omni.engine.stage_init_utils as init_mod + + od_config = types.SimpleNamespace(max_num_seqs=1) + captured: dict[str, object] = {} + monkeypatch.setattr(init_mod, "build_diffusion_config", lambda *args: od_config) + + def _capture_client(model, config, metadata, stage_init_timeout, batch_size, use_inline): + captured.update( + model=model, + config=config, + metadata=metadata, + stage_init_timeout=stage_init_timeout, + batch_size=batch_size, + use_inline=use_inline, + ) + return object() + + monkeypatch.setattr(client_mod, "create_diffusion_client", _capture_client) + metadata = _make_diffusion_metadata(0) + + init_mod.initialize_diffusion_stage( + 0, + "dummy-model", + types.SimpleNamespace(), + metadata, + stage_init_timeout=12, + batch_size=4, + use_inline=True, + ) + + assert od_config.max_num_seqs == 4 + assert captured == { + "model": "dummy-model", + "config": od_config, + "metadata": metadata, + "stage_init_timeout": 12, + "batch_size": 4, + "use_inline": True, + } + + +def test_launch_diffusion_stage_replica_applies_batch_size_to_config(monkeypatch): + import vllm_omni.diffusion.stage_diffusion_client as client_mod + import vllm_omni.diffusion.stage_diffusion_proc as proc_mod + import vllm_omni.engine.stage_engine_startup as startup_mod + + od_config = types.SimpleNamespace(max_num_seqs=1, parallel_config=types.SimpleNamespace(world_size=1)) + monkeypatch.setattr(startup_mod, "build_diffusion_config", lambda *args: od_config) + monkeypatch.setattr(startup_mod, "acquire_device_locks", lambda *args: []) + monkeypatch.setattr( + startup_mod, + "register_stage_with_omni_master", + lambda **kwargs: types.SimpleNamespace( + handshake_address="tcp://127.0.0.1:26001", + input_address="tcp://127.0.0.1:26002", + output_address="tcp://127.0.0.1:26003", + ), + ) + + omni_master_server = types.SimpleNamespace( + address="127.0.0.1", + port=25000, + release_route_port_reservations=lambda *args, **kwargs: None, + ) + proc_manager = types.SimpleNamespace( + addresses=types.SimpleNamespace( + inputs=["tcp://127.0.0.1:26002"], + outputs=["tcp://127.0.0.1:26003"], + ) + ) + monkeypatch.setattr(proc_mod, "StageDiffusionProcManager", lambda **kwargs: proc_manager) + sentinel_client = object() + monkeypatch.setattr( + client_mod.StageDiffusionClient, + "from_addresses", + lambda metadata, **kwargs: sentinel_client, + ) + + result, resources = startup_mod.launch_diffusion_stage_replica( + model="dummy-model", + stage_config=types.SimpleNamespace(), + metadata=types.SimpleNamespace(stage_id=0), + stage_init_timeout=12, + batch_size=4, + use_inline=False, + omni_master_server=omni_master_server, + ) + + assert result is sentinel_client + assert od_config.max_num_seqs == 4 + assert resources.manager is proc_manager + + def test_stage_runtime_initializes_stage_pools(monkeypatch): import vllm_omni.engine.stage_runtime as runtime_mod diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index 65ec024d5..64ae6627c 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/tests/engine/test_single_stage_mode.py @@ -1138,7 +1138,8 @@ class TestSingleStageReplicaInitialization: runtime._init_visible_devices_baseline = "0" mocker.patch.object(runtime_mod, "inject_kv_stage_info") - mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value="diffusion-config") + od_config = SimpleNamespace(max_num_seqs=None, parallel_config=SimpleNamespace(world_size=1)) + mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value=od_config) mock_register = mocker.patch( "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", return_value=StageRegistrationResponse( @@ -1174,6 +1175,7 @@ class TestSingleStageReplicaInitialization: os.environ[device_env_var] = prev_device_env assert result is sentinel_client + assert od_config.max_num_seqs == 4 mock_register.assert_called_once_with( omni_master_address="127.0.0.1", omni_master_port=25000, @@ -1183,7 +1185,7 @@ class TestSingleStageReplicaInitialization: ) mock_manager.assert_called_once_with( model="fake-model", - od_config="diffusion-config", + od_config=od_config, stage_init_timeout=60, handshake_address="tcp://127.0.0.1:26001", addresses=mocker.ANY, @@ -1230,7 +1232,8 @@ class TestSingleStageReplicaInitialization: runtime._init_visible_devices_baseline = "0" mocker.patch.object(runtime_mod, "inject_kv_stage_info") - mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value="diffusion-config") + od_config = SimpleNamespace(max_num_seqs=None, parallel_config=SimpleNamespace(world_size=1)) + mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value=od_config) mocker.patch( "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", return_value=StageRegistrationResponse( diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index f9bc9f04f..7ec6c3060 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -37,7 +37,7 @@ from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3 from vllm_omni.diffusion.postprocess import interpolate_video_tensor from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt from vllm_omni.platforms import current_omni_platform @@ -204,6 +204,8 @@ def get_wan22_post_process_func( output_type: str = "np", sampling_params=None, ): + if sampling_params is not None and sampling_params.output_type is not None: + output_type = sampling_params.output_type if output_type == "latent": return video video_metadata = {} @@ -233,6 +235,8 @@ def get_wan22_pre_process_func( prompt = request.prompt multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None + has_image = raw_image is not None and (not isinstance(raw_image, list) or bool(raw_image)) + request.batch_compatibility_key = ("wan22_image_condition", has_image) if isinstance(prompt, str): prompt = OmniTextPrompt(prompt=prompt) if "additional_information" not in prompt: @@ -288,6 +292,7 @@ class Wan22Pipeline( SupportsComponentDiscovery, WanLoraLoaderMixin, ): + supports_request_batch = True _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"] _encoder_modules: ClassVar[list[str]] = ["text_encoder"] _vae_modules: ClassVar[list[str]] = ["vae"] @@ -561,27 +566,33 @@ class Wan22Pipeline( return latents - def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: - prompt: str | None = None - negative_prompt: str | None = None - prompt_embeds: torch.Tensor | None = None - negative_prompt_embeds: torch.Tensor | None = None - if len(req.prompts) > 1: - raise ValueError( - """This model only supports a single prompt, not a batched request.""", - """Please pass in a single prompt object or string, or a single-item list.""", - ) - if len(req.prompts) == 1: - first_prompt = req.prompts[0] - prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "") - negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt") + def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]: + sampling_params_list = req.sampling_params_list + common = sampling_params_list[0] + prompt_texts = [prompt if isinstance(prompt, str) else (prompt.get("prompt") or "") for prompt in req.prompts] + negative_prompts = [ + None if isinstance(prompt, str) else prompt.get("negative_prompt") for prompt in req.prompts + ] + prompt_fields = DiffusionRequestBatch.collate_prompt_field_map( + req.prompts, + { + "prompt_embeds": None, + "negative_prompt_embeds": None, + }, + ) + prompt_embeds = prompt_fields["prompt_embeds"] + negative_prompt_embeds = prompt_fields["negative_prompt_embeds"] + prompt: list[str] | None = prompt_texts if prompt_embeds is None else None + negative_prompt: list[str] | None = None + if negative_prompt_embeds is None and any(value is not None for value in negative_prompts): + negative_prompt = [value or "" for value in negative_prompts] - if not prompt: - raise ValueError("Prompt is required for Wan2.2 generation.") + if prompt is not None and not all(prompt): + raise ValueError("Prompt is required for Wan2.2 generation when prompt_embeds are not provided.") - height = req.sampling_params.height or 480 - width = req.sampling_params.width or 832 - num_frames = req.sampling_params.num_frames or 81 + height = common.height or 480 + width = common.width or 832 + num_frames = common.num_frames or 81 # Ensure dimensions are compatible with VAE and patch size # For expand_timesteps mode, we need latent dims to be even (divisible by patch_size) @@ -589,19 +600,20 @@ class Wan22Pipeline( mod_value = self.vae_scale_factor_spatial * patch_size[1] # 16*2=32 for TI2V, 8*2=16 for I2V height = (height // mod_value) * mod_value width = (width // mod_value) * mod_value - num_steps = 40 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps + num_steps = 40 if common.num_inference_steps is None else common.num_inference_steps - output_type = req.sampling_params.output_type or "np" + output_type = common.output_type or "np" + num_outputs_per_prompt = common.num_outputs_per_prompt or 1 attention_kwargs: dict | None = None - guidance_low, guidance_high = resolve_wan_guidance_scales(req.sampling_params, default_guidance_scale=4.0) + guidance_low, guidance_high = resolve_wan_guidance_scales(common, default_guidance_scale=4.0) # record guidance for properties self._guidance_scale = guidance_low self._guidance_scale_2 = guidance_high # Prefer engine-configured boundary_ratio, but allow per-request fallback. - boundary_ratio = self.boundary_ratio if self.boundary_ratio is not None else req.sampling_params.boundary_ratio + boundary_ratio = self.boundary_ratio if self.boundary_ratio is not None else common.boundary_ratio if boundary_ratio is None: boundary_ratio = 0.875 @@ -633,32 +645,49 @@ class Wan22Pipeline( # Fallback to text_encoder dtype if no transformer loaded dtype = self.text_encoder.dtype - # Seed / generator - generator = req.sampling_params.generator - if generator is None and req.sampling_params.seed is not None: - generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed) + generator = req.collate_request_generators(num_outputs_per_prompt, None) + request_latents = req.collate_request_tensors("latents", None) if DEBUG_PERF: # Sync GPU before timing to ensure accurate measurements current_omni_platform.synchronize() _t_pipeline_start = time.perf_counter() _t_text_enc_start = _t_pipeline_start - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt=prompt, - negative_prompt=negative_prompt, - do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0, - num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1, - max_sequence_length=req.sampling_params.max_sequence_length or 512, - device=device, - dtype=dtype, - ) + do_classifier_free_guidance = guidance_low > 1.0 or guidance_high > 1.0 + if prompt_embeds is None: + prompt_embeds, negative_prompt_embeds = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=do_classifier_free_guidance, + num_videos_per_prompt=num_outputs_per_prompt, + max_sequence_length=common.max_sequence_length or 512, + device=device, + dtype=dtype, + ) + else: + prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) + prompt_embeds = prompt_embeds.repeat_interleave(num_outputs_per_prompt, dim=0) + if negative_prompt_embeds is not None: + negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype) + negative_prompt_embeds = negative_prompt_embeds.repeat_interleave(num_outputs_per_prompt, dim=0) + elif do_classifier_free_guidance: + _, negative_prompt_embeds = self.encode_prompt( + prompt=[""] * req.num_reqs, + negative_prompt=negative_prompt, + do_classifier_free_guidance=True, + num_videos_per_prompt=num_outputs_per_prompt, + max_sequence_length=common.max_sequence_length or 512, + device=device, + dtype=dtype, + ) if DEBUG_PERF: current_omni_platform.synchronize() _t_text_enc_ms = (time.perf_counter() - _t_text_enc_start) * 1000 - sample_solver = resolve_wan_sample_solver(req, default=self._sample_solver) - flow_shift = resolve_wan_flow_shift(req, self.od_config) + first_request = req.requests[0] + sample_solver = resolve_wan_sample_solver(first_request, default=self._sample_solver) + flow_shift = resolve_wan_flow_shift(first_request, self.od_config) if sample_solver != self._sample_solver or abs(flow_shift - self._flow_shift) > 1e-6: self.scheduler = build_wan_scheduler(sample_solver, flow_shift) self._sample_solver = sample_solver @@ -674,37 +703,41 @@ class Wan22Pipeline( if DEBUG_PERF: _t_latent_prep_start = time.perf_counter() - multi_modal_data = req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else None - raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None - if isinstance(raw_image, list): - if len(raw_image) > 1: - logger.warning( - """Received a list of image. Only a single image is supported by this model.""" - """Taking only the first image for now.""" - ) - raw_image = raw_image[0] - if raw_image is None: - image = None - elif isinstance(raw_image, str): - image = PIL.Image.open(raw_image) - else: - image = cast(PIL.Image.Image | torch.Tensor, raw_image) + images: list[PIL.Image.Image | torch.Tensor | None] = [] + for request_prompt in req.prompts: + multi_modal_data = request_prompt.get("multi_modal_data", {}) if not isinstance(request_prompt, str) else {} + raw_image = multi_modal_data.get("image") + if isinstance(raw_image, list): + if len(raw_image) > 1: + logger.warning("Received multiple images for one Wan request; using only the first image.") + raw_image = raw_image[0] if raw_image else None + if isinstance(raw_image, str): + raw_image = PIL.Image.open(raw_image) + images.append(cast(PIL.Image.Image | torch.Tensor | None, raw_image)) latent_condition = None first_frame_mask = None - if self.expand_timesteps and image is not None: + if self.expand_timesteps and any(image is not None for image in images): + if not all(image is not None for image in images): + raise ValueError("Cannot batch Wan requests with a mix of provided and missing image conditions.") # I2V mode: encode image and prepare condition from diffusers.video_processor import VideoProcessor video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) - # Preprocess image - if isinstance(image, PIL.Image.Image): - image = image.resize((width, height), PIL.Image.Resampling.LANCZOS) - image_tensor = video_processor.preprocess(image, height=height, width=width) - else: - image_tensor = image + image_tensors = [] + for image in images: + assert image is not None + if isinstance(image, PIL.Image.Image): + image = image.resize((width, height), PIL.Image.Resampling.LANCZOS) + image_tensor = video_processor.preprocess(image, height=height, width=width) + else: + image_tensor = image.unsqueeze(0) if image.ndim == 3 else image + image_tensors.append(image_tensor) + image_tensor = DiffusionRequestBatch.collate_tensors(image_tensors, "image condition", None) + assert image_tensor is not None + image_tensor = image_tensor.repeat_interleave(num_outputs_per_prompt, dim=0) # Use out_channels for noise latents (not in_channels which includes condition) num_channels_latents = self.transformer_config.out_channels @@ -720,7 +753,7 @@ class Wan22Pipeline( dtype=torch.float32, device=device, generator=generator, - latents=req.sampling_params.latents, + latents=request_latents, ) # Encode image condition @@ -731,7 +764,6 @@ class Wan22Pipeline( image_tensor = image_tensor.unsqueeze(2) # [B, C, 1, H, W] image_tensor = image_tensor.to(device=device, dtype=self.vae.dtype) latent_condition = retrieve_latents(self.vae.encode(image_tensor), sample_mode="argmax") - latent_condition = latent_condition.repeat(batch_size, 1, 1, 1, 1) # Normalize condition latents latents_mean = ( @@ -747,7 +779,7 @@ class Wan22Pipeline( # Create mask: 0 for first frame (condition), 1 for rest (to denoise) first_frame_mask = torch.ones( - 1, 1, num_latent_frames, latent_height, latent_width, dtype=torch.float32, device=device + batch_size, 1, num_latent_frames, latent_height, latent_width, dtype=torch.float32, device=device ) first_frame_mask[:, :, 0] = 0 else: @@ -762,7 +794,7 @@ class Wan22Pipeline( dtype=torch.float32, device=device, generator=generator, - latents=req.sampling_params.latents, + latents=request_latents, ) if DEBUG_PERF: current_omni_platform.synchronize() @@ -839,8 +871,13 @@ class Wan22Pipeline( _t_pipeline_wall_ms - _t_stages_sum, ) - return DiffusionOutput( - output=output, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + return split_diffusion_output_by_request( + DiffusionOutput( + output=output, + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + ), + req, + num_outputs_per_prompt=num_outputs_per_prompt, ) def predict_noise( diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index a2e444bb0..8e149d389 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py @@ -45,7 +45,7 @@ from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3 from vllm_omni.diffusion.postprocess import interpolate_video_tensor from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.platforms import current_omni_platform @@ -65,6 +65,8 @@ def get_wan22_i2v_post_process_func( output_type: str = "np", sampling_params=None, ): + if sampling_params is not None and sampling_params.output_type is not None: + output_type = sampling_params.output_type if output_type == "latent": return video video_metadata = {} @@ -84,6 +86,32 @@ def get_wan22_i2v_post_process_func( return post_process_func +def _normalize_i2v_last_image( + value: str | PIL.Image.Image | torch.Tensor | list[object] | None, +) -> PIL.Image.Image | torch.Tensor | None: + """Normalize the optional I2V last-frame condition for batching.""" + if value is None: + return None + + if isinstance(value, list): + if not value: + return None + if len(value) != 1: + raise ValueError("I2V accepts at most one last_image.") + value = value[0] + + if isinstance(value, str): + value = PIL.Image.open(value).convert("RGB") + + if not isinstance(value, (PIL.Image.Image, torch.Tensor)): + raise TypeError( + f"Unsupported last_image format {value.__class__}. " + "Expected a file path, PIL.Image.Image, torch.Tensor, or None." + ) + + return value + + def get_wan22_i2v_pre_process_func( od_config: OmniDiffusionConfig, ): @@ -109,6 +137,9 @@ def get_wan22_i2v_pre_process_func( """Please correctly set `"multi_modal_data": {"image": , …}`""", ) image = PIL.Image.open(raw_image).convert("RGB") if isinstance(raw_image, str) else raw_image + last_image = _normalize_i2v_last_image(multi_modal_data.get("last_image")) # type: ignore[union-attr] + prompt["multi_modal_data"]["last_image"] = last_image # type: ignore[index] + request.batch_compatibility_key = ("wan22_i2v_last_image", last_image is not None) # Calculate dimensions based on aspect ratio if not provided if request.sampling_params.height is None or request.sampling_params.width is None: @@ -157,6 +188,7 @@ class Wan22I2VPipeline( Wan2.2-style I2V (with expand_timesteps for TI2V-5B). """ + supports_request_batch = True _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"] _encoder_modules: ClassVar[list[str]] = ["text_encoder", "image_encoder"] _vae_modules: ClassVar[list[str]] = ["vae"] @@ -433,60 +465,67 @@ class Wan22I2VPipeline( quant_config = getattr(self.od_config, "quantization_config", None) return create_transformer_from_config(config, quant_config=quant_config) - def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: - prompt: str | None = None - negative_prompt: str | None = None - prompt_embeds: torch.Tensor | None = None - negative_prompt_embeds: torch.Tensor | None = None + def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]: + sampling_params_list = req.sampling_params_list + common = sampling_params_list[0] + prompt_texts = [prompt if isinstance(prompt, str) else (prompt.get("prompt") or "") for prompt in req.prompts] + negative_prompts = [ + None if isinstance(prompt, str) else prompt.get("negative_prompt") for prompt in req.prompts + ] + prompt_fields = DiffusionRequestBatch.collate_prompt_field_map( + req.prompts, + { + "prompt_embeds": None, + "negative_prompt_embeds": None, + }, + ) + prompt_embeds = prompt_fields["prompt_embeds"] + negative_prompt_embeds = prompt_fields["negative_prompt_embeds"] image_embeds: torch.Tensor | None = None - if len(req.prompts) > 1: - raise ValueError( - """This model only supports a single prompt, not a batched request.""", - """Please pass in a single prompt object or string, or a single-item list.""", - ) - if len(req.prompts) == 1: - first_prompt = req.prompts[0] - prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "") - negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt") + prompt: list[str] | None = prompt_texts if prompt_embeds is None else None + negative_prompt: list[str] | None = None + if negative_prompt_embeds is None and any(value is not None for value in negative_prompts): + negative_prompt = [value or "" for value in negative_prompts] + if prompt is not None and not all(prompt): + raise ValueError("Prompt is required for Wan2.2 I2V generation when prompt_embeds are not provided.") - if not prompt: - raise ValueError("Prompt is required for Wan2.2 generation.") + images: list[PIL.Image.Image | torch.Tensor] = [] + last_images: list[PIL.Image.Image | torch.Tensor | None] = [] + for request_prompt in req.prompts: + multi_modal_data = request_prompt.get("multi_modal_data", {}) if not isinstance(request_prompt, str) else {} + raw_image = multi_modal_data.get("image") + if raw_image is None: + raise ValueError("Image is required for I2V generation.") + if isinstance(raw_image, list): + if len(raw_image) > 1: + logger.warning("Received multiple images for one I2V request; using only the first image.") + raw_image = raw_image[0] + if isinstance(raw_image, str): + raw_image = PIL.Image.open(raw_image).convert("RGB") + images.append(cast(PIL.Image.Image | torch.Tensor, raw_image)) - # Get image from request - multi_modal_data = req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else None - raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None - if raw_image is None: - raise ValueError("Image is required for I2V generation.") - if isinstance(raw_image, list): - if len(raw_image) > 1: - logger.warning( - """Received a list of image. Only a single image is supported by this model.""" - """Taking only the first image for now.""" - ) - raw_image = raw_image[0] - if isinstance(raw_image, str): - image = PIL.Image.open(raw_image) - else: - image = cast(PIL.Image.Image | torch.Tensor, raw_image) + last_image = multi_modal_data.get("last_image") + if isinstance(last_image, str): + last_image = PIL.Image.open(last_image).convert("RGB") + last_images.append(cast(PIL.Image.Image | torch.Tensor | None, last_image)) + if any(image is not None for image in last_images) and not all(image is not None for image in last_images): + raise ValueError("Cannot batch I2V requests with a mix of provided and missing last_image conditions.") - last_image: PIL.Image.Image | torch.Tensor | None = None - if multi_modal_data is not None: - last_image = multi_modal_data.get("last_image", None) + height = common.height or 480 + width = common.width or 832 + num_frames = common.num_frames or 81 + num_steps = 40 if common.num_inference_steps is None else common.num_inference_steps - height = req.sampling_params.height or 480 - width = req.sampling_params.width or 832 - num_frames = req.sampling_params.num_frames or 81 - num_steps = 40 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps - - output_type = req.sampling_params.output_type or "np" + output_type = common.output_type or "np" + num_outputs_per_prompt = common.num_outputs_per_prompt or 1 attention_kwargs: dict | None = None - guidance_low, guidance_high = resolve_wan_guidance_scales(req.sampling_params, default_guidance_scale=5.0) + guidance_low, guidance_high = resolve_wan_guidance_scales(common, default_guidance_scale=5.0) self._guidance_scale = guidance_low self._guidance_scale_2 = guidance_high - boundary_ratio = self.boundary_ratio if self.boundary_ratio is not None else req.sampling_params.boundary_ratio + boundary_ratio = self.boundary_ratio if self.boundary_ratio is not None else common.boundary_ratio if boundary_ratio is None: boundary_ratio = 0.875 logger.warning("boundary_ratio is required for I2V generation. using default value 0.875") @@ -495,7 +534,7 @@ class Wan22I2VPipeline( self.check_inputs( prompt=prompt, negative_prompt=negative_prompt, - image=image, + image=images, height=height, width=width, prompt_embeds=prompt_embeds, @@ -513,10 +552,8 @@ class Wan22I2VPipeline( device = self.device dtype = self.transformer.dtype - # Generator setup - generator = req.sampling_params.generator - if generator is None and req.sampling_params.seed is not None: - generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed) + generator = req.collate_request_generators(num_outputs_per_prompt, None) + request_latents = req.collate_request_tensors("latents", None) if DEBUG_PERF: # Sync GPU before timing to ensure accurate measurements @@ -524,15 +561,33 @@ class Wan22I2VPipeline( _t_pipeline_start = time.perf_counter() _t_text_enc_start = _t_pipeline_start - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt=prompt, - negative_prompt=negative_prompt, - do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0, - num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1, - max_sequence_length=req.sampling_params.max_sequence_length or 512, - device=device, - dtype=dtype, - ) + do_classifier_free_guidance = guidance_low > 1.0 or guidance_high > 1.0 + if prompt_embeds is None: + prompt_embeds, negative_prompt_embeds = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=do_classifier_free_guidance, + num_videos_per_prompt=num_outputs_per_prompt, + max_sequence_length=common.max_sequence_length or 512, + device=device, + dtype=dtype, + ) + else: + prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) + prompt_embeds = prompt_embeds.repeat_interleave(num_outputs_per_prompt, dim=0) + if negative_prompt_embeds is not None: + negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype) + negative_prompt_embeds = negative_prompt_embeds.repeat_interleave(num_outputs_per_prompt, dim=0) + elif do_classifier_free_guidance: + _, negative_prompt_embeds = self.encode_prompt( + prompt=[""] * req.num_reqs, + negative_prompt=negative_prompt, + do_classifier_free_guidance=True, + num_videos_per_prompt=num_outputs_per_prompt, + max_sequence_length=common.max_sequence_length or 512, + device=device, + dtype=dtype, + ) if DEBUG_PERF: current_omni_platform.synchronize() @@ -543,12 +598,18 @@ class Wan22I2VPipeline( if DEBUG_PERF: _t_img_enc_start = time.perf_counter() if self.has_image_encoder and self.transformer.config.image_dim is not None: - if image_embeds is None: - if last_image is None: - image_embeds = self.encode_image(image, device) - else: - image_embeds = self.encode_image([image, last_image], device) - image_embeds = image_embeds.repeat(batch_size, 1, 1) + if all(last_image is None for last_image in last_images): + image_embeds = self.encode_image(images, device) + image_embeds = image_embeds.repeat_interleave(num_outputs_per_prompt, dim=0) + else: + image_pairs = [] + for first_image, last_image in zip(images, last_images): + assert last_image is not None + image_pairs.extend([first_image, last_image]) + image_embeds = self.encode_image(image_pairs, device) + image_embeds = image_embeds.view(req.num_reqs, 2, *image_embeds.shape[1:]) + image_embeds = image_embeds.repeat_interleave(num_outputs_per_prompt, dim=0) + image_embeds = image_embeds.flatten(0, 1) image_embeds = image_embeds.to(dtype) else: image_embeds = None @@ -557,8 +618,9 @@ class Wan22I2VPipeline( current_omni_platform.synchronize() _t_img_enc_ms = (time.perf_counter() - _t_img_enc_start) * 1000 - sample_solver = resolve_wan_sample_solver(req, default=self._sample_solver) - flow_shift = resolve_wan_flow_shift(req, self.od_config) + first_request = req.requests[0] + sample_solver = resolve_wan_sample_solver(first_request, default=self._sample_solver) + flow_shift = resolve_wan_flow_shift(first_request, self.od_config) if sample_solver != self._sample_solver or abs(flow_shift - self._flow_shift) > 1e-6: self.scheduler = build_wan_scheduler(sample_solver, flow_shift) self._sample_solver = sample_solver @@ -582,23 +644,36 @@ class Wan22I2VPipeline( video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) - if isinstance(image, PIL.Image.Image): - image = TF.to_tensor(image).to(device) - image_tensor = video_processor.preprocess(image, height=height, width=width) - else: - image_tensor = image - image_tensor = image_tensor.to(device=device, dtype=torch.float32) - - # Handle last_image if provided - if last_image is not None: - if isinstance(last_image, PIL.Image.Image): - image = TF.to_tensor(last_image).to(device) - last_image_tensor = video_processor.preprocess(last_image, height=height, width=width) + image_tensors = [] + for image in images: + if isinstance(image, PIL.Image.Image): + image = TF.to_tensor(image).to(device) + image_tensor = video_processor.preprocess(image, height=height, width=width) else: - last_image_tensor = last_image - last_image_tensor = last_image_tensor.to(device=device, dtype=torch.float32) - else: - last_image_tensor = None + image_tensor = image.unsqueeze(0) if image.ndim == 3 else image + image_tensors.append(image_tensor.to(device=device, dtype=torch.float32)) + image_tensor = DiffusionRequestBatch.collate_tensors(image_tensors, "image condition", None) + assert image_tensor is not None + image_tensor = image_tensor.repeat_interleave(num_outputs_per_prompt, dim=0) + + last_image_tensor = None + if all(last_image is not None for last_image in last_images): + last_image_tensors = [] + for last_image in last_images: + assert last_image is not None + if isinstance(last_image, PIL.Image.Image): + last_image = TF.to_tensor(last_image).to(device) + current_last_image = video_processor.preprocess(last_image, height=height, width=width) + else: + current_last_image = last_image.unsqueeze(0) if last_image.ndim == 3 else last_image + last_image_tensors.append(current_last_image.to(device=device, dtype=torch.float32)) + last_image_tensor = DiffusionRequestBatch.collate_tensors( + last_image_tensors, + "last image condition", + None, + ) + assert last_image_tensor is not None + last_image_tensor = last_image_tensor.repeat_interleave(num_outputs_per_prompt, dim=0) latents, condition, first_frame_mask = self.prepare_latents( image=image_tensor, @@ -610,7 +685,7 @@ class Wan22I2VPipeline( dtype=torch.float32, device=device, generator=generator, - latents=req.sampling_params.latents, + latents=request_latents, last_image=last_image_tensor, ) @@ -694,8 +769,13 @@ class Wan22I2VPipeline( _t_pipeline_wall_ms - _t_stages_sum, ) - return DiffusionOutput( - output=output, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + return split_diffusion_output_by_request( + DiffusionOutput( + output=output, + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + ), + req, + num_outputs_per_prompt=num_outputs_per_prompt, ) def predict_noise( @@ -849,7 +929,6 @@ class Wan22I2VPipeline( # Encode through VAE latent_condition = retrieve_latents(self.vae.encode(video_condition), sample_mode="argmax") - latent_condition = latent_condition.repeat(batch_size, 1, 1, 1, 1) # Normalize latents latents_mean = ( @@ -866,7 +945,7 @@ class Wan22I2VPipeline( if self.expand_timesteps: # TI2V-5B style: create mask where first frame is 0 (condition), rest is 1 (to denoise) first_frame_mask = torch.ones( - 1, 1, num_latent_frames, latent_height, latent_width, dtype=dtype, device=device + batch_size, 1, num_latent_frames, latent_height, latent_width, dtype=dtype, device=device ) first_frame_mask[:, :, 0] = 0 return latents, latent_condition, first_frame_mask @@ -892,7 +971,9 @@ class Wan22I2VPipeline( condition = torch.concat([mask_lat_size, latent_condition], dim=1) # For non-expand mode, first_frame_mask is not used in the same way - first_frame_mask = torch.ones(1, 1, num_latent_frames, latent_height, latent_width, dtype=dtype, device=device) + first_frame_mask = torch.ones( + batch_size, 1, num_latent_frames, latent_height, latent_width, dtype=dtype, device=device + ) return latents, condition, first_frame_mask diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py index 52a42cb79..bce71633c 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py @@ -43,7 +43,7 @@ from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import ( ) from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.platforms import current_omni_platform @@ -63,6 +63,30 @@ _S2V_DEFAULT_NEG_PROMPT = ( _AUDIO_VIDEO_RATE = 30 +def _s2v_audio_condition_key(audio: str | np.ndarray | None) -> tuple[Any, ...]: + if isinstance(audio, str): + # File paths are deliberately conservative: two different files may + # have different durations even when their request metadata matches. + return ("path", audio) + if isinstance(audio, np.ndarray): + return ("array", tuple(audio.shape), str(audio.dtype)) + return (type(audio).__name__,) + + +def _make_clip_generators( + seeds: list[int | None], + request_generators: list[torch.Generator | None], + clip_index: int, + device: torch.device, +) -> list[torch.Generator | None]: + return [ + torch.Generator(device=device).manual_seed((seed if seed is not None else 0) + clip_index) + if seed is not None or request_generator is None + else request_generator + for seed, request_generator in zip(seeds, request_generators) + ] + + def _linear_interpolation( features: torch.Tensor, input_fps: float, @@ -337,6 +361,12 @@ def get_wan22_s2v_pre_process_func( prompt["additional_information"]["init_first_frame"] = ( multi_modal_data.get("init_first_frame", False) if multi_modal_data is not None else False ) + request.batch_compatibility_key = ( + "wan22_s2v_condition", + bool(prompt["additional_information"]["init_first_frame"]), + prompt["additional_information"].get("num_repeat"), + _s2v_audio_condition_key(raw_audio), + ) request.prompt = prompt return request @@ -486,6 +516,7 @@ class Wan22S2VPipeline( """ # Default config values from Wan2.2/wan/configs/wan_s2v_14B.py + supports_request_batch = True _DEFAULT_MOTION_FRAMES = 73 _DEFAULT_INFER_FRAMES = 80 _DEFAULT_FPS = 16 @@ -942,7 +973,6 @@ class Wan22S2VPipeline( """Forward pass through the S2V transformer to predict noise. With return_dict=False, the model returns (output,) where output is [B, C, T, H, W]. - We squeeze batch dim since S2V processes one sample at a time. """ if current_model is None: current_model = self.transformer @@ -952,8 +982,8 @@ class Wan22S2VPipeline( if isinstance(result, list): return result[0] if isinstance(result, tuple): - return result[0].squeeze(0) - return result.sample.squeeze(0) + return result[0] + return result.sample def diffuse( self, @@ -962,7 +992,7 @@ class Wan22S2VPipeline( prompt_embeds: torch.Tensor, negative_prompt_embeds: torch.Tensor | None, guidance_scale: float, - clip_generator: torch.Generator, + clip_generator: torch.Generator | list[torch.Generator] | None, dtype: torch.dtype, device: torch.device, max_seq_len: int, @@ -971,13 +1001,13 @@ class Wan22S2VPipeline( ref_latents: torch.Tensor, motion_frames: list[int], drop_first_motion: bool, - positive_audio_emb: torch.Tensor, - negative_audio_emb: torch.Tensor | None, + positive_audio_emb: dict[str, torch.Tensor], + negative_audio_emb: dict[str, torch.Tensor] | None, ) -> torch.Tensor: - """Denoising diffusion loop for one S2V clip. + """Denoising diffusion loop for a batch of compatible S2V clips. Args: - latents: Initial noise tensor [C, T, H, W] + latents: Initial noise tensor [B, C, T, H, W] timesteps: Denoising timesteps prompt_embeds: Text embeddings [1, seq_len, dim] negative_prompt_embeds: Negative text embeddings (optional) @@ -995,7 +1025,7 @@ class Wan22S2VPipeline( negative_audio_emb: Precomputed audio embeddings for negative prompt (optional) Returns: - Denoised latents [C, T, H, W] + Denoised latents [B, C, T, H, W] """ do_true_cfg = self.do_classifier_free_guidance and negative_prompt_embeds is not None @@ -1004,13 +1034,13 @@ class Wan22S2VPipeline( self._current_timestep = t self.record_denoise_step(step_idx, t) - latent_model_input = [latents.to(device)] - timestep = t.unsqueeze(0).to(device) if t.dim() == 0 else t.to(device) + latent_model_input = latents.to(device) + timestep = t.expand(latents.shape[0]).to(device) positive_kwargs = { "hidden_states": latent_model_input, "timestep": timestep, - "encoder_hidden_states": prompt_embeds[0:1], + "encoder_hidden_states": prompt_embeds, "cond_states": cond_latents, "motion_latents": input_motion_latents, "ref_latents": ref_latents, @@ -1023,7 +1053,7 @@ class Wan22S2VPipeline( negative_kwargs = { "hidden_states": latent_model_input, "timestep": timestep, - "encoder_hidden_states": negative_prompt_embeds[0:1], + "encoder_hidden_states": negative_prompt_embeds, "cond_states": cond_latents, "motion_latents": input_motion_latents, "ref_latents": ref_latents, @@ -1044,12 +1074,12 @@ class Wan22S2VPipeline( # Scheduler step latents = self.scheduler.step( - noise_pred.unsqueeze(0), + noise_pred, t, - latents.unsqueeze(0), + latents, return_dict=False, - generator=clip_generator, - )[0].squeeze(0) + generator=clip_generator if isinstance(clip_generator, torch.Generator) else None, + )[0] pbar.update() @@ -1080,111 +1110,178 @@ class Wan22S2VPipeline( prompt_embeds: torch.Tensor | None = None, negative_prompt_embeds: torch.Tensor | None = None, **kwargs, - ) -> DiffusionOutput: + ) -> list[DiffusionOutput]: """Run S2V generation — may produce multiple autoregressive clips. This method mirrors ``WanS2V.generate()``, reorganized into the vLLM-Omni pipeline structure. """ - # ---- Extract params from request ---- - if len(req.prompts) > 1: - raise ValueError("S2V only supports a single prompt per request.") - - if len(req.prompts) == 1: - prompt_data = req.prompts[0] - prompt = prompt_data if isinstance(prompt_data, str) else prompt_data.get("prompt") - negative_prompt = None if isinstance(prompt_data, str) else prompt_data.get("negative_prompt") + sampling_params_list = req.sampling_params_list + common = sampling_params_list[0] + prompt_fields = DiffusionRequestBatch.collate_prompt_field_map( + req.prompts, + { + "prompt_embeds": prompt_embeds, + "negative_prompt_embeds": negative_prompt_embeds, + }, + ) + prompt_embeds = prompt_fields["prompt_embeds"] + negative_prompt_embeds = prompt_fields["negative_prompt_embeds"] + prompt_texts: list[str] = [] + negative_prompts: list[str] = [] + images: list[PIL.Image.Image | None] = [] + audio_paths: list[str | np.ndarray | None] = [] + requested_repeats: list[int | None] = [] + init_first_frames: list[bool] = [] + for prompt_data in req.prompts: multi_modal_data = prompt_data.get("multi_modal_data", {}) if not isinstance(prompt_data, str) else {} additional_info = prompt_data.get("additional_information", {}) if not isinstance(prompt_data, str) else {} - if image is None: - raw_image = multi_modal_data.get("image", None) - if isinstance(raw_image, str): - image = PIL.Image.open(raw_image).convert("RGB") - elif isinstance(raw_image, PIL.Image.Image): - image = raw_image + prompt_texts.append( + prompt_data if isinstance(prompt_data, str) else (prompt_data.get("prompt") or prompt or "") + ) + request_negative_prompt = ( + negative_prompt if isinstance(prompt_data, str) else prompt_data.get("negative_prompt", negative_prompt) + ) + negative_prompts.append(request_negative_prompt or _S2V_DEFAULT_NEG_PROMPT) - if audio_path is None: - audio_path = additional_info.get("audio_path") - if audio_path is None: - audio_path = multi_modal_data.get("audio") + raw_image = multi_modal_data.get("image", image) + if isinstance(raw_image, list): + raw_image = raw_image[0] if raw_image else None + if isinstance(raw_image, str): + raw_image = PIL.Image.open(raw_image).convert("RGB") + images.append(raw_image if isinstance(raw_image, PIL.Image.Image) else None) - if pose_video is None: - pose_video = additional_info.get("pose_video") + request_audio_path = additional_info.get("audio_path", audio_path) + if request_audio_path is None: + request_audio_path = multi_modal_data.get("audio") + audio_paths.append(request_audio_path) + requested_repeats.append(additional_info.get("num_repeat", num_repeat)) + init_first_frames.append(additional_info.get("init_first_frame", init_first_frame)) - init_first_frame = additional_info.get("init_first_frame", init_first_frame) - - if negative_prompt is None: - negative_prompt = _S2V_DEFAULT_NEG_PROMPT - - height = req.sampling_params.height or height - width = req.sampling_params.width or width + height = common.height or height + width = common.width or width # num_frames defaults to 1 (image-model sentinel); S2V needs many frames. - req_num_frames = req.sampling_params.num_frames + req_num_frames = common.num_frames infer_frames = ( infer_frames or (req_num_frames if req_num_frames and req_num_frames > 1 else None) or self._DEFAULT_INFER_FRAMES ) - num_steps = req.sampling_params.num_inference_steps or num_inference_steps + num_steps = common.num_inference_steps or num_inference_steps + num_outputs_per_prompt = common.num_outputs_per_prompt or 1 - if req.sampling_params.guidance_scale_provided: - guidance_scale = req.sampling_params.guidance_scale + if common.guidance_scale_provided: + guidance_scale = common.guidance_scale self._guidance_scale = guidance_scale # ---- Validate ---- - self.check_inputs(prompt, image, audio_path, height, width, prompt_embeds) + if len(set(init_first_frames)) != 1: + raise ValueError("Batched S2V requests must have matching init_first_frame values.") + init_first_frame = init_first_frames[0] + for request_prompt, request_image, request_audio_path in zip(prompt_texts, images, audio_paths): + self.check_inputs( + request_prompt if prompt_embeds is None else None, + request_image, + request_audio_path, + height, + width, + prompt_embeds, + ) device = self.device dtype = self.transformer.dtype if hasattr(self.transformer, "dtype") else torch.bfloat16 - - if generator is None: - generator = req.sampling_params.generator - seed = req.sampling_params.seed - if generator is None and seed is not None: - generator = torch.Generator(device=device).manual_seed(seed) - if seed is None: - seed = 0 + request_generators = req.collate_request_generators(num_outputs_per_prompt, None) + batch_size = req.num_reqs * num_outputs_per_prompt + if isinstance(request_generators, list): + generators = request_generators + else: + generators = [request_generators] * batch_size + seeds = [sampling.seed for sampling in sampling_params_list for _ in range(num_outputs_per_prompt)] + request_latents = req.collate_request_tensors("latents", None) # ---- 1. Text encoding ---- if prompt_embeds is None: prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt=prompt, - negative_prompt=negative_prompt, + prompt=prompt_texts, + negative_prompt=negative_prompts, do_classifier_free_guidance=self.do_classifier_free_guidance, - num_videos_per_prompt=1, - max_sequence_length=512, + num_videos_per_prompt=num_outputs_per_prompt, + max_sequence_length=common.max_sequence_length or 512, device=device, dtype=dtype, ) else: prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) + prompt_embeds = prompt_embeds.repeat_interleave(num_outputs_per_prompt, dim=0) if negative_prompt_embeds is not None: negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype) + negative_prompt_embeds = negative_prompt_embeds.repeat_interleave(num_outputs_per_prompt, dim=0) # ---- 2. Audio encoding ---- - audio_emb, audio_num_repeat, target_video_frames = self.encode_audio( - audio_path, infer_frames=infer_frames, device=device, dtype=dtype - ) - if num_repeat is None or num_repeat > audio_num_repeat: - num_repeat = audio_num_repeat + audio_embeddings: list[torch.Tensor] = [] + effective_repeats: list[int] = [] + target_video_frames: list[int] = [] + for request_audio_path, requested_repeat in zip(audio_paths, requested_repeats): + assert request_audio_path is not None + audio_emb, audio_num_repeat, request_target_video_frames = self.encode_audio( + request_audio_path, infer_frames=infer_frames, device=device, dtype=dtype + ) + audio_embeddings.append(audio_emb.repeat_interleave(num_outputs_per_prompt, dim=0)) + effective_repeats.append( + audio_num_repeat + if requested_repeat is None or requested_repeat > audio_num_repeat + else requested_repeat + ) + target_video_frames.append(request_target_video_frames) + + if len(set(effective_repeats)) != 1: + raise ValueError(f"Batched S2V requests must have matching num_repeat values, got {effective_repeats}.") + if len(set(target_video_frames)) != 1: + raise ValueError( + "Batched S2V requests must have matching audio embedding time lengths; " + f"got target video frames {target_video_frames}." + ) + num_repeat = effective_repeats[0] + audio_emb = DiffusionRequestBatch.collate_tensors(audio_embeddings, "S2V audio embeddings", None) + assert audio_emb is not None # Load raw audio waveform for muxing into output video - if isinstance(audio_path, np.ndarray): - raw_audio_waveform = audio_path.astype(np.float32) - raw_audio_sr = 16000 - else: - raw_audio_waveform, raw_audio_sr = load_audio(audio_path, sr=None, mono=True) + raw_audio_waveforms: list[np.ndarray] = [] + raw_audio_sample_rates: list[int] = [] + for request_audio_path in audio_paths: + assert request_audio_path is not None + if isinstance(request_audio_path, np.ndarray): + raw_audio_waveforms.append(request_audio_path.astype(np.float32, copy=False)) + raw_audio_sample_rates.append(16000) + else: + raw_audio_waveform, raw_audio_sr = load_audio(request_audio_path, sr=None, mono=True) + raw_audio_waveforms.append(raw_audio_waveform) + raw_audio_sample_rates.append(raw_audio_sr) + if len(set(raw_audio_sample_rates)) != 1 or len({waveform.shape for waveform in raw_audio_waveforms}) != 1: + raise ValueError("Batched S2V requests must have matching raw audio shapes and sample rates.") + raw_audio_sr = raw_audio_sample_rates[0] # ---- 3. Reference image encoding ---- - ref_latents = self.encode_ref_image(image, height, width, device=device).to(dtype=dtype) + ref_latents = DiffusionRequestBatch.collate_tensors( + [ + self.encode_ref_image(request_image, height, width, device=device) + .to(dtype=dtype) + .repeat_interleave(num_outputs_per_prompt, dim=0) + for request_image in images + if request_image is not None + ], + "S2V reference image latents", + None, + ) + assert ref_latents is not None # ---- 4. Initial motion latents (zeros) ---- motion_frames = min(self.motion_frames, infer_frames - 1) motion_pixels = torch.zeros( - [1, 3, motion_frames, height, width], + [batch_size, 3, motion_frames, height, width], dtype=dtype, device=device, ) @@ -1193,8 +1290,11 @@ class Wan22S2VPipeline( drop_first_motion = False # Place reference image in the last 6 frames of motion tensor_trans = transforms.ToTensor() - ref_tensor = tensor_trans(image).to(device=device, dtype=dtype) * 2 - 1.0 # [C, H, W] - motion_pixels[:, :, -6:] = ref_tensor.unsqueeze(0).unsqueeze(2).expand(-1, -1, 6, -1, -1) + ref_tensors = torch.stack( + [tensor_trans(request_image) for request_image in images if request_image is not None] + ).to(device=device, dtype=dtype) + ref_tensors = ref_tensors.repeat_interleave(num_outputs_per_prompt, dim=0) * 2 - 1.0 + motion_pixels[:, :, -6:] = ref_tensors.unsqueeze(2).expand(-1, -1, 6, -1, -1) motion_latents = self.prepare_motion_latents(motion_pixels, device=device).to(dtype=dtype) @@ -1211,24 +1311,29 @@ class Wan22S2VPipeline( clips: list[torch.Tensor] = [] # Keep a pixel-space buffer of the trailing motion_frames for the # autoregressive connection between clips. - videos_last_frames = torch.zeros([1, 3, motion_frames, height, width], dtype=dtype, device=device) + videos_last_frames = torch.zeros([batch_size, 3, motion_frames, height, width], dtype=dtype, device=device) for r in range(num_repeat): - # Per-clip seed - clip_seed = seed + r - clip_generator = torch.Generator(device=device).manual_seed(clip_seed) - + clip_generators = _make_clip_generators(seeds, generators, r, device) # -- Noise -- - latents = self.prepare_latents( - infer_frames=infer_frames, - height=height, - width=width, - motion_frames=motion_frames, - lat_motion_frames=lat_motion_frames, - dtype=dtype, - device=device, - generator=clip_generator, - ) + if r == 0 and request_latents is not None: + latents = request_latents.to(device=device, dtype=dtype) + else: + latents = torch.stack( + [ + self.prepare_latents( + infer_frames=infer_frames, + height=height, + width=width, + motion_frames=motion_frames, + lat_motion_frames=lat_motion_frames, + dtype=dtype, + device=device, + generator=clip_generator, + ) + for clip_generator in clip_generators + ] + ) # -- Scheduler -- self.scheduler.set_timesteps(num_steps, device=device) @@ -1242,11 +1347,11 @@ class Wan22S2VPipeline( # -- Pose condition for this clip -- # Default: zero condition (no pose driving) - # Shape: [B=1, C=16, T, H, W] — passed to transformer which iterates over batch dim + # Shape: [B, C=16, T, H, W] lat_h = height // self.vae_scale_factor_spatial lat_w = width // self.vae_scale_factor_spatial lat_target_frames = (infer_frames + 3 + motion_frames) // 4 - lat_motion_frames - cond_latents = torch.zeros([1, 16, lat_target_frames, lat_h, lat_w], dtype=dtype, device=device) + cond_latents = torch.zeros([batch_size, 16, lat_target_frames, lat_h, lat_w], dtype=dtype, device=device) # -- Clone motion latents for this clip -- input_motion_latents = motion_latents.clone() @@ -1282,7 +1387,7 @@ class Wan22S2VPipeline( prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, guidance_scale=guidance_scale, - clip_generator=clip_generator, + clip_generator=clip_generators, dtype=dtype, device=device, max_seq_len=max_seq_len, @@ -1300,11 +1405,10 @@ class Wan22S2VPipeline( self.transformer.to("cpu") current_omni_platform.empty_cache() - latents_for_decode = latents.unsqueeze(0) # [1, C, T, H, W] if not (drop_first_motion and r == 0): - decode_latents = torch.cat([motion_latents, latents_for_decode], dim=2) + decode_latents = torch.cat([motion_latents, latents], dim=2) else: - decode_latents = torch.cat([ref_latents, latents_for_decode], dim=2) + decode_latents = torch.cat([ref_latents, latents], dim=2) decode_latents = self._denormalize_latents(decode_latents) decode_latents = decode_latents.to(self.vae.dtype) @@ -1321,7 +1425,7 @@ class Wan22S2VPipeline( # Create buffer for broadcast clip_video = torch.empty( - (1, 3, total_frames, height, width), + (batch_size, 3, total_frames, height, width), device=decode_latents.device, dtype=decode_latents.dtype, ) @@ -1356,12 +1460,19 @@ class Wan22S2VPipeline( current_omni_platform.empty_cache() # ---- Concatenate all clips ---- - output = torch.cat(clips, dim=2) # [1, C, T_total, H, W] + output = torch.cat(clips, dim=2) # [B, C, T_total, H, W] - return DiffusionOutput( - output=(output, raw_audio_waveform, raw_audio_sr), - stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + outputs = split_diffusion_output_by_request( + DiffusionOutput( + output=output, + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + ), + req, + num_outputs_per_prompt=num_outputs_per_prompt, ) + for request_output, raw_audio_waveform in zip(outputs, raw_audio_waveforms): + request_output.output = (request_output.output, raw_audio_waveform, raw_audio_sr) + return outputs def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights using AutoWeightsLoader for vLLM integration.""" diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py index d0cdb1250..ae869fd57 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py @@ -35,13 +35,27 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( ) from vllm_omni.diffusion.models.wan2_2.wan2_2_vace_transformer import WanVACETransformer3DModel from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) +def _vace_condition_structure(value: object) -> tuple: + if value is None: + return ("none",) + if isinstance(value, str): + return ("path", value) + if isinstance(value, PIL.Image.Image): + return ("image", value.size) + if isinstance(value, torch.Tensor): + return ("tensor", tuple(value.shape), str(value.dtype)) + if isinstance(value, list): + return ("list", tuple(_vace_condition_structure(item) for item in value)) + return (type(value).__name__,) + + def create_vace_transformer_from_config( config: dict, quant_config: QuantizationConfig | None = None, @@ -104,6 +118,7 @@ def get_wan22_vace_pre_process_func(od_config: OmniDiffusionConfig): prompt["additional_information"] = {} if not multi_modal_data: + request.batch_compatibility_key = ("wan22_vace_condition", ("none",), ("none",), ("none",)) request.prompt = prompt return request @@ -155,6 +170,13 @@ def get_wan22_vace_pre_process_func(od_config: OmniDiffusionConfig): mask = [mask] prompt["additional_information"]["mask"] = mask + request.batch_compatibility_key = ( + "wan22_vace_condition", + _vace_condition_structure(ref_images), + _vace_condition_structure(source_video), + _vace_condition_structure(mask), + ) + request.prompt = prompt return request @@ -482,7 +504,7 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput): attention_kwargs: dict | None = None, vace_context_scale: float | list[float] = 1.0, **kwargs, - ) -> DiffusionOutput: + ) -> list[DiffusionOutput]: """Generate or edit video using VACE. The mode is determined by which inputs are provided in the request: @@ -507,47 +529,48 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput): attention_kwargs: Additional kwargs for attention layers. vace_context_scale: VACE conditioning strength. """ - # Get parameters from request or arguments - if len(req.prompts) > 1: - raise ValueError( - "This model only supports a single prompt, not a batched request. " - "Please pass in a single prompt object or string, or a single-item list." - ) - - reference_images = None - source_video = None - source_mask = None - - if len(req.prompts) == 1: - first_prompt = req.prompts[0] - if isinstance(first_prompt, str): - prompt = first_prompt - else: - prompt = first_prompt.get("prompt") - negative_prompt = negative_prompt or first_prompt.get("negative_prompt") - prompt_embeds = prompt_embeds if prompt_embeds is not None else first_prompt.get("prompt_embeds") - negative_prompt_embeds = ( - negative_prompt_embeds - if negative_prompt_embeds is not None - else first_prompt.get("negative_prompt_embeds") - ) - - additional_info = first_prompt.get("additional_information", {}) - reference_images = additional_info.get("reference_images") - source_video = additional_info.get("source_video") - source_mask = additional_info.get("mask") - - if prompt is None and prompt_embeds is None: + sampling_params_list = req.sampling_params_list + common = sampling_params_list[0] + prompt_texts = [item if isinstance(item, str) else (item.get("prompt") or "") for item in req.prompts] + negative_prompts = [None if isinstance(item, str) else item.get("negative_prompt") for item in req.prompts] + prompt_fields = DiffusionRequestBatch.collate_prompt_field_map( + req.prompts, + { + "prompt_embeds": prompt_embeds, + "negative_prompt_embeds": negative_prompt_embeds, + }, + ) + prompt_embeds = prompt_fields["prompt_embeds"] + negative_prompt_embeds = prompt_fields["negative_prompt_embeds"] + prompt = prompt_texts if prompt_embeds is None else None + negative_prompt = None + if negative_prompt_embeds is None and any(value is not None for value in negative_prompts): + negative_prompt = [value or "" for value in negative_prompts] + if prompt is not None and not all(prompt): raise ValueError("Prompt or prompt_embeds is required for VACE generation.") - height = req.sampling_params.height or height - width = req.sampling_params.width or width - num_frames = req.sampling_params.num_frames or frame_num - num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps - generator = req.sampling_params.generator or generator + reference_images_list = [] + source_videos = [] + source_masks = [] + for request_prompt in req.prompts: + additional_info = ( + request_prompt.get("additional_information", {}) if not isinstance(request_prompt, str) else {} + ) + reference_images_list.append(additional_info.get("reference_images")) + source_videos.append(additional_info.get("source_video")) + source_masks.append(additional_info.get("mask")) - if req.sampling_params.guidance_scale_provided: - guidance_scale = req.sampling_params.guidance_scale + height = common.height or height + width = common.width or width + num_frames = common.num_frames or frame_num + num_inference_steps = common.num_inference_steps or num_inference_steps + num_outputs_per_prompt = common.num_outputs_per_prompt or 1 + generator = req.collate_request_generators(num_outputs_per_prompt, None) + request_latents = req.collate_request_tensors("latents", None) + output_type = common.output_type or output_type + + if common.guidance_scale_provided: + guidance_scale = common.guidance_scale # Ensure dimensions are compatible with VAE and patch size mod_value = self.vae_scale_factor_spatial * 2 # 8 * 2 = 16 @@ -558,47 +581,47 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput): num_frames = num_frames // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1 num_frames = max(num_frames, 1) - self.check_inputs( - prompt=prompt, - negative_prompt=negative_prompt, - height=height, - width=width, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - video=source_video, - mask=source_mask, - reference_images=reference_images, - ) + for index in range(req.num_reqs): + self.check_inputs( + prompt=None if prompt is None else prompt[index], + negative_prompt=None if negative_prompt is None else negative_prompt[index], + height=height, + width=width, + prompt_embeds=None if prompt_embeds is None else prompt_embeds[index], + negative_prompt_embeds=None if negative_prompt_embeds is None else negative_prompt_embeds[index], + video=source_videos[index], + mask=source_masks[index], + reference_images=reference_images_list[index], + ) device = self.device self._guidance_scale = guidance_scale dtype = self.transformer.dtype if self.transformer is not None else torch.bfloat16 - if generator is None and req.sampling_params.seed is not None: - generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed) - # Encode prompts if prompt_embeds is None: - if prompt is None: - raise ValueError("Either prompt or prompt_embeds must be provided.") prompt_embeds, negative_prompt_embeds = self.encode_prompt( prompt=prompt, negative_prompt=negative_prompt, do_classifier_free_guidance=guidance_scale > 1.0, - num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1, - max_sequence_length=req.sampling_params.max_sequence_length or 512, + num_videos_per_prompt=num_outputs_per_prompt, + max_sequence_length=common.max_sequence_length or 512, device=device, dtype=dtype, ) else: prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) + prompt_embeds = prompt_embeds.repeat_interleave(num_outputs_per_prompt, dim=0) if negative_prompt_embeds is not None: negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype) + negative_prompt_embeds = negative_prompt_embeds.repeat_interleave(num_outputs_per_prompt, dim=0) elif guidance_scale > 1.0: _, negative_prompt_embeds = self.encode_prompt( - prompt="", - negative_prompt=None, + prompt=[""] * req.num_reqs, + negative_prompt=negative_prompt, do_classifier_free_guidance=True, + num_videos_per_prompt=num_outputs_per_prompt, + max_sequence_length=common.max_sequence_length or 512, device=device, dtype=dtype, ) @@ -608,24 +631,43 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput): # so read it from whichever transformer exists. active_transformer = self.transformer if self.transformer is not None else self.transformer_2 if active_transformer.vace_patch_embedding is not None: - video, mask, ref_images_processed = self.preprocess_conditions( - video=source_video, - mask=source_mask, - reference_images=reference_images, - height=height, - width=width, - num_frames=num_frames, - dtype=dtype, - device=device, - ) + contexts = [] + reference_counts = [] + for index, (source_video, source_mask, reference_images) in enumerate( + zip(source_videos, source_masks, reference_images_list) + ): + video, mask, ref_images_processed = self.preprocess_conditions( + video=source_video, + mask=source_mask, + reference_images=reference_images, + height=height, + width=width, + num_frames=num_frames, + dtype=dtype, + device=device, + ) + request_generator = sampling_params_list[index].generator + if isinstance(request_generator, list): + request_generator = request_generator[0] + conditioning_latents = self.prepare_video_latents( + video, + mask, + ref_images_processed, + request_generator, + device, + ) + mask_encoded = self.prepare_masks(mask, ref_images_processed) + context = torch.cat([conditioning_latents, mask_encoded], dim=1) + contexts.append(context.repeat_interleave(num_outputs_per_prompt, dim=0)) + reference_counts.append(len(ref_images_processed[0]) if ref_images_processed else 0) - conditioning_latents = self.prepare_video_latents(video, mask, ref_images_processed, generator, device) - mask_encoded = self.prepare_masks(mask, ref_images_processed) - - # Unified VACE context: [video_latents, mask] along channels -> [B, C, T, H, W] - vace_context = torch.cat([conditioning_latents, mask_encoded], dim=1) - - num_reference_images = len(ref_images_processed[0]) if ref_images_processed else 0 + if len(set(reference_counts)) != 1: + raise ValueError( + f"Batched VACE conditions must have the same number of reference images; got {reference_counts}." + ) + vace_context = DiffusionRequestBatch.collate_tensors(contexts, "VACE condition", None) + assert vace_context is not None + num_reference_images = reference_counts[0] else: vace_context = None @@ -641,7 +683,7 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput): dtype=torch.float32, device=device, generator=generator, - latents=req.sampling_params.latents, + latents=request_latents, ) # Set up scheduler @@ -654,9 +696,7 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput): # stays None and diffuse() uses the same transformer for every step. boundary_timestep = None if self.transformer_2 is not None: - boundary_ratio = ( - self.boundary_ratio if self.boundary_ratio is not None else req.sampling_params.boundary_ratio - ) + boundary_ratio = self.boundary_ratio if self.boundary_ratio is not None else common.boundary_ratio if boundary_ratio is None: boundary_ratio = 0.875 logger.warning_once("boundary_ratio is required for Wan2.2 VACE generation. using default value 0.875") @@ -700,7 +740,11 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput): latents = latents / latents_std + latents_mean output = self.vae.decode(latents, return_dict=False)[0] - return DiffusionOutput(output=output) + return split_diffusion_output_by_request( + DiffusionOutput(output=output), + req, + num_outputs_per_prompt=num_outputs_per_prompt, + ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights using AutoWeightsLoader for vLLM integration.""" diff --git a/vllm_omni/diffusion/request.py b/vllm_omni/diffusion/request.py index 8834e408e..ce097a2d9 100644 --- a/vllm_omni/diffusion/request.py +++ b/vllm_omni/diffusion/request.py @@ -43,6 +43,10 @@ class OmniDiffusionRequest: # BaseScheduler takes ownership and clears this field before the request is # sent to a Worker. diffusion_kv_requests: tuple[DiffusionKVRequest, ...] | None = None + # Optional model-specific condition structure used by request-batch admission. + # This is populated by a pipeline preprocessor before the request reaches + # the scheduler; ``None`` keeps the default behavior for other pipelines. + batch_compatibility_key: tuple[Any, ...] | None = None def __post_init__(self): """Initialize dependent fields after dataclass initialization.""" diff --git a/vllm_omni/diffusion/sched/interface.py b/vllm_omni/diffusion/sched/interface.py index 8fb36f405..e543ed251 100644 --- a/vllm_omni/diffusion/sched/interface.py +++ b/vllm_omni/diffusion/sched/interface.py @@ -88,9 +88,10 @@ class RequestBatchSamplingParamsKey: """Request level Batch-compatibility key derived from ``OmniDiffusionSamplingParams``. Only request-batch-wide fields belong here. Request-local values such as - seeds, generators, latent tensors, timesteps, and pipeline-specific + seeds, generators, latent tensors, timesteps, and request-local ``extra_args`` are read per request from - ``DiffusionRequestBatch.sampling_params_list``. + ``DiffusionRequestBatch.sampling_params_list``. Structural ``extra_args`` + that affect batching are listed explicitly below. """ # Spatial / temporal shape. @@ -131,6 +132,16 @@ class RequestBatchSamplingParamsKey: # the model and must remain distinct from explicit ``lossless``. quality: str | None = None + # Wan scheduler structure is carried through extra_args. Requests using + # different solvers or flow shifts must not share a request batch. + sample_solver: str | None = None + flow_shift: float | None = None + + # Pipeline-specific condition structure populated during preprocessing. + # It prevents independently valid requests with incompatible conditions + # from being admitted to the same request batch. + condition_key: tuple[Any, ...] | None = None + # LoRA identity. lora_int_id: int | None = None lora_scale: float = 1.0 diff --git a/vllm_omni/diffusion/sched/request_scheduler.py b/vllm_omni/diffusion/sched/request_scheduler.py index 968360183..4168ae974 100644 --- a/vllm_omni/diffusion/sched/request_scheduler.py +++ b/vllm_omni/diffusion/sched/request_scheduler.py @@ -22,7 +22,21 @@ if TYPE_CHECKING: # on sampling params, so it must be resolved separately from the bulk lookup. _REQUEST_BATCH_SAMPLING_PARAMS_KEY_FIELD_NAMES = frozenset( field.name for field in fields(RequestBatchSamplingParamsKey) -) - {"lora_int_id"} +) - {"condition_key", "flow_shift", "lora_int_id", "sample_solver"} + + +def _normalize_explicit_sample_solver(value: object | None) -> str | None: + """Normalize an explicitly provided solver without selecting a default.""" + if value is None: + return None + return str(value).strip().lower() + + +def _normalize_explicit_flow_shift(value: object | None) -> float | None: + """Normalize an explicitly provided flow shift without selecting a default.""" + if value is None: + return None + return float(value) def build_request_batch_sampling_params_key(request: OmniDiffusionRequest) -> RequestBatchSamplingParamsKey: @@ -31,6 +45,13 @@ def build_request_batch_sampling_params_key(request: OmniDiffusionRequest) -> Re # LoRA identity is optional on sampling params (and on test stubs). lora_request = getattr(sampling, "lora_request", None) key_kwargs = {name: getattr(sampling, name) for name in _REQUEST_BATCH_SAMPLING_PARAMS_KEY_FIELD_NAMES} + extra_args = sampling.extra_args or {} + # Match pipeline resolution for explicit overrides, but preserve None: + # pipeline/engine defaults are configuration-dependent and must not be + # inferred while building the request-batch key. + key_kwargs["sample_solver"] = _normalize_explicit_sample_solver(extra_args.get("sample_solver")) + key_kwargs["flow_shift"] = _normalize_explicit_flow_shift(extra_args.get("flow_shift")) + key_kwargs["condition_key"] = getattr(request, "batch_compatibility_key", None) key_kwargs["lora_int_id"] = lora_request.lora_int_id if lora_request is not None else None return RequestBatchSamplingParamsKey(**key_kwargs) diff --git a/vllm_omni/engine/stage_engine_startup.py b/vllm_omni/engine/stage_engine_startup.py index 5dacc31df..081b14703 100644 --- a/vllm_omni/engine/stage_engine_startup.py +++ b/vllm_omni/engine/stage_engine_startup.py @@ -1558,6 +1558,7 @@ def launch_diffusion_stage_replica( from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient od_config = build_diffusion_config(model, stage_config, metadata) + od_config.max_num_seqs = batch_size parallel_config = getattr(od_config, "parallel_config", None) world_size = getattr(parallel_config, "world_size", 1) try: diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index c2e1ab5bc..054a2fb22 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -1441,6 +1441,7 @@ def initialize_diffusion_stage( from vllm_omni.diffusion.stage_diffusion_client import create_diffusion_client od_config = build_diffusion_config(model, stage_cfg, metadata) + od_config.max_num_seqs = batch_size return create_diffusion_client(model, od_config, metadata, stage_init_timeout, batch_size, use_inline)