From f58ef57ac4594b966441ee57cd36e2c331a6a056 Mon Sep 17 00:00:00 2001 From: Aman <67479108+Bezdarnost@users.noreply.github.com> Date: Thu, 24 Sep 2026 22:22:17 +0600 Subject: [PATCH] [Bugfix] Enable pipeline profiling for LingBot Video (#7744) Signed-off-by: bezdarnost --- .../test_pipeline_lingbot_video.py | 70 +++++++++++++++++++ .../lingbot_video/pipeline_lingbot_video.py | 13 +++- 2 files changed, 81 insertions(+), 2 deletions(-) diff --git a/tests/diffusion/models/lingbot_video/test_pipeline_lingbot_video.py b/tests/diffusion/models/lingbot_video/test_pipeline_lingbot_video.py index 3775ef665..f3e0b1ce3 100644 --- a/tests/diffusion/models/lingbot_video/test_pipeline_lingbot_video.py +++ b/tests/diffusion/models/lingbot_video/test_pipeline_lingbot_video.py @@ -26,6 +26,7 @@ def _make_pipeline(): pipeline.default_image_negative_prompt = "default image negative" pipeline.img_prompt_template = "" pipeline.od_config = SimpleNamespace(flow_shift=None) + pipeline.enable_diffusion_pipeline_profiler = False return pipeline @@ -572,3 +573,72 @@ def test_load_weights_rejects_external_weight_stream(): assert pipeline.load_weights([]) == set() with pytest.raises(RuntimeError, match="components are loaded directly"): pipeline.load_weights([("transformer.weight", torch.zeros(1))]) + + +@pytest.mark.parametrize("enabled", [False, True], ids=["disabled", "enabled"]) +def test_pipeline_profiler_records_executed_stages_per_request(monkeypatch, enabled): + from vllm_omni.diffusion.models.lingbot_video import pipeline_lingbot_video as module + from vllm_omni.diffusion.profiler import diffusion_pipeline_profiler as profiler_module + from vllm_omni.diffusion.worker.utils import consume_pipeline_stage_durations + + class TextEncoder(nn.Module): + def forward(self): + return torch.ones(1, 2, 4), torch.ones(1, 2, dtype=torch.long) + + class VAE(nn.Module): + def __init__(self): + super().__init__() + self.anchor = nn.Parameter(torch.zeros(())) + self.config = SimpleNamespace(latents_mean=[0.0], latents_std=[1.0]) + + def encode(self, value): + return value + + def decode(self, latents): + return (latents.repeat(1, 3, 1, 1, 1),) + + components = [ + (module.LingBotVideoTransformer3DModel, _RecordingTransformer()), + (module.Qwen3VLForConditionalGeneration, TextEncoder()), + (module.Qwen3VLProcessor, SimpleNamespace()), + (module.AutoencoderKLWan, VAE()), + (module.FlowUniPCMultistepScheduler, _CorruptingScheduler()), + ] + for component_class, component in components: + monkeypatch.setattr(component_class, "from_pretrained", lambda *args, value=component, **kwargs: value) + monkeypatch.setattr(module, "get_local_device", lambda: torch.device("cpu")) + monkeypatch.setattr(module.LingBotVideoPipeline, "encode_prompt", lambda self, *args, **kwargs: self.text_encoder()) + monkeypatch.setattr(profiler_module.current_omni_platform, "is_available", lambda: False) + ticks = iter(range(100)) + monkeypatch.setattr(profiler_module.time, "perf_counter", lambda: float(next(ticks))) + pipeline = module.LingBotVideoPipeline( + od_config=SimpleNamespace( + model="unused", dtype=torch.float32, enable_diffusion_pipeline_profiler=enabled, flow_shift=None + ) + ) + request = _make_request_batch( + {"prompt": "a robot", "modalities": ["video"]}, + height=16, + width=16, + num_frames=9, + num_inference_steps=2, + guidance_scale=1.0, + seed=42, + ) + prefix = "LingBotVideoPipeline." + expected = {prefix + "text_encoder.forward": 1.0, prefix + "transformer.forward": 2.0, prefix + "vae.decode": 1.0} + first = pipeline(request) + assert first.stage_durations == (expected if enabled else {}) + assert first.output["video"].shape == (3, 2, 2, 3) + # A latent-only request must not inherit the previous request's VAE timing. + request.requests[0].sampling_params.output_type = "latent" + second = pipeline(request) + latent_expected = {key: value for key, value in expected.items() if not key.endswith("vae.decode")} + assert second.stage_durations == (latent_expected if enabled else {}) + assert first.stage_durations == (expected if enabled else {}) + consumed = consume_pipeline_stage_durations(pipeline) + if enabled: + assert consumed == {**latent_expected, prefix + "forward": 7.0} + else: + assert consumed == {} + assert consume_pipeline_stage_durations(pipeline) == {} diff --git a/vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py b/vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py index 1c94fecc6..51687a241 100644 --- a/vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py +++ b/vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Adapted from LingBot-Video (https://github.com/Robbyant/lingbot-video). from __future__ import annotations @@ -36,6 +36,7 @@ from vllm_omni.diffusion.models.lingbot_video.request_utils import ( ) from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler +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.errors import OmniClientError @@ -306,6 +307,7 @@ class LingBotVideoPipeline( SupportImageInput, ProgressBarMixin, SupportsComponentDiscovery, + DiffusionPipelineProfilerMixin, ): """Native vLLM-Omni entry for LingBot-Video checkpoints. @@ -382,6 +384,10 @@ class LingBotVideoPipeline( self.set_progress_bar_config(disable=bool(model_config.get("quiet_progress", True))) self.default_negative_prompt = DEFAULT_NEGATIVE_PROMPT self.default_image_negative_prompt = DEFAULT_NEGATIVE_PROMPT_IMAGE + self.setup_diffusion_pipeline_profiler( + profiler_targets=["text_encoder.forward", "transformer.forward", "vae.encode", "vae.decode"], + enable_diffusion_pipeline_profiler=od_config.enable_diffusion_pipeline_profiler, + ) def to(self, *args, **kwargs): device, dtype, non_blocking, _ = torch._C._nn._parse_to(*args, **kwargs) @@ -871,4 +877,7 @@ class LingBotVideoPipeline( ), ) output_key = "image" if request_config.mode is LingBotGenerationMode.T2I else "video" - return DiffusionOutput(output={output_key: frames}) + return DiffusionOutput( + output={output_key: frames}, + stage_durations=self.stage_durations if self.enable_diffusion_pipeline_profiler else {}, + )