From fde4808723709c3404073dde87dab9bf011f9c20 Mon Sep 17 00:00:00 2001 From: dengyunyang <584797741@qq.com> Date: Thu, 19 Mar 2026 12:26:12 +0800 Subject: [PATCH] [Performance] diffusion timing (#1757) Signed-off-by: dengyunyang <584797741@qq.com> Signed-off-by: bjf-frz Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com> Co-authored-by: bjf-frz Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com> --- benchmarks/diffusion/backends.py | 14 +++ .../model/adding_diffusion_model.md | 101 ++++++++++++++++++ examples/offline_inference/bagel/end2end.py | 6 ++ .../offline_inference/glm_image/end2end.py | 6 ++ examples/offline_inference/helios/end2end.py | 6 ++ .../hunyuan_image3/image_to_text.py | 10 +- .../image_to_image/image_edit.py | 6 ++ .../image_to_video/image_to_video.py | 6 ++ .../run_mammothmoda2_image_summarize.py | 6 ++ .../offline_inference/mimo_audio/end2end.py | 6 ++ .../offline_inference/qwen2_5_omni/end2end.py | 6 ++ .../offline_inference/qwen3_omni/end2end.py | 6 ++ .../qwen3_omni/end2end_async_chunk.py | 6 ++ .../offline_inference/qwen3_tts/end2end.py | 7 ++ .../text_to_audio/text_to_audio.py | 10 +- .../text_to_image/text_to_image.py | 6 ++ .../text_to_video/text_to_video.py | 6 ++ vllm_omni/diffusion/data.py | 6 ++ vllm_omni/diffusion/diffusion_engine.py | 4 + .../diffusion/models/bagel/pipeline_bagel.py | 10 +- .../diffusion/models/flux/pipeline_flux.py | 11 +- .../flux2_klein/pipeline_flux2_klein.py | 10 +- .../models/glm_image/pipeline_glm_image.py | 10 +- .../models/helios/pipeline_helios.py | 10 +- .../longcat_image/pipeline_longcat_image.py | 10 +- .../pipeline_longcat_image_edit.py | 10 +- .../nextstep_1_1/pipeline_nextstep_1_1.py | 10 +- .../models/ovis_image/pipeline_ovis_image.py | 10 +- .../models/qwen_image/pipeline_qwen_image.py | 11 +- .../qwen_image/pipeline_qwen_image_edit.py | 10 +- .../pipeline_qwen_image_edit_plus.py | 12 ++- .../qwen_image/pipeline_qwen_image_layered.py | 10 +- .../diffusion/models/sd3/pipeline_sd3.py | 10 +- .../stable_audio/pipeline_stable_audio.py | 10 +- .../models/wan2_2/pipeline_wan2_2.py | 11 +- .../models/wan2_2/pipeline_wan2_2_i2v.py | 12 ++- .../models/z_image/pipeline_z_image.py | 10 +- .../profiler/diffusion_pipeline_profiler.py | 101 ++++++++++++++++++ vllm_omni/engine/async_omni_engine.py | 1 + vllm_omni/entrypoints/cli/serve.py | 7 ++ vllm_omni/entrypoints/omni_base.py | 2 + vllm_omni/entrypoints/openai/serving_chat.py | 6 ++ vllm_omni/outputs.py | 6 ++ 43 files changed, 496 insertions(+), 38 deletions(-) create mode 100644 vllm_omni/diffusion/profiler/diffusion_pipeline_profiler.py diff --git a/benchmarks/diffusion/backends.py b/benchmarks/diffusion/backends.py index 98d3a30fe..6097cdccc 100644 --- a/benchmarks/diffusion/backends.py +++ b/benchmarks/diffusion/backends.py @@ -35,6 +35,7 @@ class RequestFuncOutput: error: str = "" start_time: float = 0.0 response_body: dict[str, Any] = field(default_factory=dict) + stage_durations: dict[str, float] = field(default_factory=dict) peak_memory_mb: float = 0.0 slo_achieved: bool | None = None @@ -55,6 +56,7 @@ async def async_request_chat_completions( input: RequestFuncInput, session: aiohttp.ClientSession, pbar: tqdm | None = None, + enable_diffusion_pipeline_profiler: bool = False, ) -> RequestFuncOutput: output = RequestFuncOutput() output.start_time = time.perf_counter() @@ -108,6 +110,18 @@ async def async_request_chat_completions( output.success = True if "peak_memory_mb" in resp_json: output.peak_memory_mb = resp_json["peak_memory_mb"] + try: + choices = resp_json.get("choices", []) + if choices and isinstance(choices, list): + msg = choices[0].get("message", {}) + if isinstance(msg, dict): + content = msg.get("content", []) + if content and isinstance(content, list) and len(content) > 0: + first_item = content[0] + if isinstance(first_item, dict): + output.stage_durations = first_item.get("stage_durations") + except (IndexError, TypeError, AttributeError): + pass else: output.error = f"HTTP {response.status}: {await response.text()}" output.success = False diff --git a/docs/contributing/model/adding_diffusion_model.md b/docs/contributing/model/adding_diffusion_model.md index 818bd15b7..3d04bf942 100644 --- a/docs/contributing/model/adding_diffusion_model.md +++ b/docs/contributing/model/adding_diffusion_model.md @@ -815,6 +815,107 @@ class WanTransformer3DModel(nn.Module): --- +### Diffusion Timing (Performance Profiling) +When adapting a new diffusion model, it is often useful to analyze the latency of key components such as text encoding, diffusion denoising, and VAE decoding. +vLLM-Omni provides a timing utility via `DiffusionPipelineProfilerMixin` to help developers quickly identify performance bottlenecks. + +!!! info + `DiffusionPipelineProfilerMixin` is different from using `torch.profiler` for diffusion models, as introduced in this [tutorial](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/profiling.md#3-profiling-diffusion-models). `DiffusionPipelineProfilerMixin` only prints the timing information of multiple functions (such as `vae.decode`), while `torch.profiler` saves detailed GPU/CPU computation time, call/execution steps. + +This tool automatically measures the execution time of selected pipeline modules and prints the results in the logs. + +**Enabling Diffusion Timing** + + +Enable timing by setting: +``` +vllm serve Qwen/Qwen-Image --omni --port 8091 --enable-diffusion-pipeline-profiler +``` +You can optionally specify which modules to profile: +``` +class YourPipeline(xxx, DiffusionPipelineProfilerMixin): + def __init__(self, xxx): + ... + self.setup_diffusion_pipeline_profiler(profiler_targets=["diffuse"], enable_diffusion_pipeline_profiler) +``` +If not specified, the default targets are used: +``` +["vae.encode", "vae.decode", "diffuse", "text_encoder.forward", "tokenizer.forward"] +``` +**Adding DiffusionPipelineProfilerMixin to a Pipeline** +To enable timing support in your pipeline, inherit from DiffusionPipelineProfilerMixin. +```python +from vllm_omni.diffusion.utils.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin + +class YourModelPipeline(nn.Module, DiffusionPipelineProfilerMixin): + # Optional: Specify custom timing targets + _PROFILER_TARGETS = ["vae.encode", "vae.decode", "diffuse", "text_encoder.forward", "tokenizer.forward"] + + def __init__( + self, + *, + od_config: OmniDiffusionConfig, + prefix: str = "", + ): + super().__init__() + self.od_config = od_config + self.parallel_config = od_config.parallel_config + # initialize pipeline components + ... + + # initialize timing profiler + self.setup_diffusion_pipeline_profiler(enable_diffusion_pipeline_profiler) +``` +The mixin dynamically wraps selected methods and records their execution time during inference. + +If you need to fetch the execution time of different modules, you will need to pass `self.stage_durations` to `DiffusionOutput`, as shown below: + +```diff +- return DiffusionOutput(output=img) ++ return DiffusionOutput( + output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + ) +``` + +**Pipeline Design for Timing** +The current diffusion timing utility is function-based, meaning it measures the execution time of individual methods. + +When implementing a new pipeline, avoid putting all logic inside a single function (e.g., forward). Instead, structure the pipeline in a modular way by separating key stages into independent methods, such as the diffusion loop. + +For example: +``` +def forward(self, req: OmniDiffusionRequest): + prompt_embeds = self.encode_prompt(req) + latents = self.diffuse(prompt_embeds, req) + images = self.vae.decode(latents) + return DiffusionOutput(output=images) +``` +This allows the timing utility to measure each stage (e.g., encode_prompt, diffuse, vae.decode) separately and helps identify performance bottlenecks more easily. + + +**Default Profiled Modules** + +By default, the following pipeline modules are timed: +``` +vae.encode +vae.decode +diffuse +text_encoder.forward +tokenizer.forward +``` + +**Example Output** + +When enabled, timing logs appear like this: +``` +[DiffusionTiming] text_encoder.forward took 0.018s +[DiffusionTiming] diffuse took 2.412s +[DiffusionTiming] vae.decode took 0.063s +``` +These measurements help identify bottlenecks during model adaptation and optimization + + + ## Troubleshooting diff --git a/examples/offline_inference/bagel/end2end.py b/examples/offline_inference/bagel/end2end.py index dca0907a1..b1f308709 100644 --- a/examples/offline_inference/bagel/end2end.py +++ b/examples/offline_inference/bagel/end2end.py @@ -58,6 +58,11 @@ def parse_args(): help="CFG parallel size: 1=batched (single GPU), 2=parallel with 2 branches (text CFG only), 3=parallel (3 GPUs).", ) parser.add_argument("--seed", type=int, default=None, help="Random seed for generation.") + parser.add_argument( + "--enable-diffusion-pipeline-profiler", + action="store_true", + help="Enable diffusion pipeline profiler to display stage durations.", + ) args = parser.parse_args() return args @@ -103,6 +108,7 @@ def main(): "shm_threshold_bytes": args.shm_threshold_bytes, "worker_backend": args.worker_backend, "ray_address": args.ray_address, + "enable_diffusion_pipeline_profiler": args.enable_diffusion_pipeline_profiler, } ) diff --git a/examples/offline_inference/glm_image/end2end.py b/examples/offline_inference/glm_image/end2end.py index a3c1b7dc2..3cb74f82c 100644 --- a/examples/offline_inference/glm_image/end2end.py +++ b/examples/offline_inference/glm_image/end2end.py @@ -247,6 +247,7 @@ def main(args: argparse.Namespace) -> None: stage_configs_path=config_path, log_stats=args.enable_stats, stage_init_timeout=args.stage_init_timeout, + enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, ) init_time = time.time() - start_time @@ -463,6 +464,11 @@ def parse_args() -> argparse.Namespace: default=False, help="Enable verbose output", ) + parser.add_argument( + "--enable-diffusion-pipeline-profiler", + action="store_true", + help="Enable diffusion pipeline profiler to display stage durations.", + ) return parser.parse_args() diff --git a/examples/offline_inference/helios/end2end.py b/examples/offline_inference/helios/end2end.py index 36f5bff16..dae707c2c 100644 --- a/examples/offline_inference/helios/end2end.py +++ b/examples/offline_inference/helios/end2end.py @@ -173,6 +173,11 @@ def parse_args() -> argparse.Namespace: default=1, help="Number of initial denoising steps using zero prediction (default: 1).", ) + parser.add_argument( + "--enable-diffusion-pipeline-profiler", + action="store_true", + help="Enable diffusion pipeline profiler to display stage durations.", + ) # Memory & parallelism parser.add_argument("--vae-use-slicing", action="store_true", help="Enable VAE slicing.") @@ -207,6 +212,7 @@ def main(): enable_cpu_offload=args.enable_cpu_offload, parallel_config=parallel_config, enforce_eager=args.enforce_eager, + enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, ) # Validate I2V / V2V arguments diff --git a/examples/offline_inference/hunyuan_image3/image_to_text.py b/examples/offline_inference/hunyuan_image3/image_to_text.py index dbe6cdea5..f5182c770 100644 --- a/examples/offline_inference/hunyuan_image3/image_to_text.py +++ b/examples/offline_inference/hunyuan_image3/image_to_text.py @@ -41,6 +41,11 @@ def parse_args() -> argparse.Namespace: required=True, help="Pretrain template prompt: <|startoftext|>{system}{question}", ) + parser.add_argument( + "--enable-diffusion-pipeline-profiler", + action="store_true", + help="Enable diffusion pipeline profiler to display stage durations.", + ) return parser.parse_args() @@ -52,7 +57,10 @@ def load_image(image_path: str) -> Image.Image: def main(args: argparse.Namespace) -> None: - omni = Omni(model=args.model) + omni = Omni( + model=args.model, + enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, + ) prompt_dict = { "prompt": args.prompt, diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py index c15478b71..a8035a3fd 100644 --- a/examples/offline_inference/image_to_image/image_edit.py +++ b/examples/offline_inference/image_to_image/image_edit.py @@ -325,6 +325,11 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable layerwise (blockwise) offloading on DiT modules.", ) + parser.add_argument( + "--enable-diffusion-pipeline-profiler", + action="store_true", + help="Enable diffusion pipeline profiler to display stage durations.", + ) return parser.parse_args() @@ -389,6 +394,7 @@ def main(): parallel_config=parallel_config, enforce_eager=args.enforce_eager, enable_cpu_offload=args.enable_cpu_offload, + enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, ) print("Pipeline loaded") diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index 6d4dc06ac..e312c6e72 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -176,6 +176,11 @@ def parse_args() -> argparse.Namespace: "Default 1 means pure sharding (no replication). " ), ) + parser.add_argument( + "--enable-diffusion-pipeline-profiler", + action="store_true", + help="Enable diffusion pipeline profiler to display stage durations.", + ) return parser.parse_args() @@ -281,6 +286,7 @@ def main(): model_class_name=model_class_name, cache_backend=args.cache_backend, cache_config=cache_config, + enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, ) if profiler_enabled: diff --git a/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py b/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py index 830557aa0..ca87b9e9a 100644 --- a/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py +++ b/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py @@ -43,6 +43,11 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--top-p", type=float, default=0.9) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--trust-remote-code", action="store_true") + parser.add_argument( + "--enable-diffusion-pipeline-profiler", + action="store_true", + help="Enable diffusion pipeline profiler to display stage durations.", + ) return parser.parse_args() @@ -70,6 +75,7 @@ def main() -> None: model=args.model, stage_configs_path=args.stage_config, trust_remote_code=args.trust_remote_code, + enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, ) try: sp = SamplingParams( diff --git a/examples/offline_inference/mimo_audio/end2end.py b/examples/offline_inference/mimo_audio/end2end.py index ae044d2e8..bdcfcd52b 100644 --- a/examples/offline_inference/mimo_audio/end2end.py +++ b/examples/offline_inference/mimo_audio/end2end.py @@ -189,6 +189,7 @@ def main(args): batch_timeout=args.batch_timeout, init_timeout=args.init_timeout, shm_threshold_bytes=args.shm_threshold_bytes, + enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, ) thinker_sampling_params = SamplingParams( @@ -433,6 +434,11 @@ def parse_args(): default="../../../model_executor/stage_configs/mimo_audio.yaml", help="Path to a stage configs file.", ) + parser.add_argument( + "--enable-diffusion-pipeline-profiler", + action="store_true", + help="Enable diffusion pipeline profiler to display stage durations.", + ) return parser.parse_args() diff --git a/examples/offline_inference/qwen2_5_omni/end2end.py b/examples/offline_inference/qwen2_5_omni/end2end.py index 7bba59983..9ae99793c 100644 --- a/examples/offline_inference/qwen2_5_omni/end2end.py +++ b/examples/offline_inference/qwen2_5_omni/end2end.py @@ -327,6 +327,7 @@ def main(args): batch_timeout=args.batch_timeout, init_timeout=args.init_timeout, shm_threshold_bytes=args.shm_threshold_bytes, + enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, ) thinker_sampling_params = SamplingParams( temperature=0.0, # Deterministic - no randomness @@ -540,6 +541,11 @@ def parse_args(): default=False, help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.", ) + parser.add_argument( + "--enable-diffusion-pipeline-profiler", + action="store_true", + help="Enable diffusion pipeline profiler to display stage durations.", + ) return parser.parse_args() diff --git a/examples/offline_inference/qwen3_omni/end2end.py b/examples/offline_inference/qwen3_omni/end2end.py index 0bf143e0c..639a6187c 100644 --- a/examples/offline_inference/qwen3_omni/end2end.py +++ b/examples/offline_inference/qwen3_omni/end2end.py @@ -299,6 +299,7 @@ def main(args): stage_configs_path=args.stage_configs_path, log_stats=args.log_stats, stage_init_timeout=args.stage_init_timeout, + enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, ) thinker_sampling_params = SamplingParams( @@ -526,6 +527,11 @@ def parse_args(): default=False, help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.", ) + parser.add_argument( + "--enable-diffusion-pipeline-profiler", + action="store_true", + help="Enable diffusion pipeline profiler to display stage durations.", + ) return parser.parse_args() diff --git a/examples/offline_inference/qwen3_omni/end2end_async_chunk.py b/examples/offline_inference/qwen3_omni/end2end_async_chunk.py index 8adbae9eb..3eaf30a1c 100644 --- a/examples/offline_inference/qwen3_omni/end2end_async_chunk.py +++ b/examples/offline_inference/qwen3_omni/end2end_async_chunk.py @@ -382,6 +382,7 @@ async def run_all(args): stage_configs_path=args.stage_configs_path, log_stats=args.log_stats, stage_init_timeout=args.stage_init_timeout, + enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, ) # Use default sampling params from stage config (they are pre-configured @@ -584,6 +585,11 @@ def parse_args(): default=16000, help="Sampling rate for audio loading.", ) + parser.add_argument( + "--enable-diffusion-pipeline-profiler", + action="store_true", + help="Enable diffusion pipeline profiler to display stage durations.", + ) return parser.parse_args() diff --git a/examples/offline_inference/qwen3_tts/end2end.py b/examples/offline_inference/qwen3_tts/end2end.py index 901418c39..72ae67b48 100644 --- a/examples/offline_inference/qwen3_tts/end2end.py +++ b/examples/offline_inference/qwen3_tts/end2end.py @@ -371,6 +371,7 @@ def main(args): stage_configs_path=args.stage_configs_path, log_stats=args.log_stats, stage_init_timeout=args.stage_init_timeout, + enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, ) batch_size = args.batch_size @@ -392,6 +393,7 @@ async def main_streaming(args): stage_configs_path=args.stage_configs_path, log_stats=args.log_stats, stage_init_timeout=args.stage_init_timeout, + enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, ) for i, prompt in enumerate(inputs): @@ -533,6 +535,11 @@ def parse_args(): default=1, help="Number of prompts per batch (default: 1, sequential).", ) + parser.add_argument( + "--enable-diffusion-pipeline-profiler", + action="store_true", + help="Enable diffusion pipeline profiler to display stage durations.", + ) return parser.parse_args() diff --git a/examples/offline_inference/text_to_audio/text_to_audio.py b/examples/offline_inference/text_to_audio/text_to_audio.py index 68aeee265..a6968c419 100644 --- a/examples/offline_inference/text_to_audio/text_to_audio.py +++ b/examples/offline_inference/text_to_audio/text_to_audio.py @@ -90,6 +90,11 @@ def parse_args() -> argparse.Namespace: default=44100, help="Sample rate for output audio (Stable Audio uses 44100 Hz).", ) + parser.add_argument( + "--enable-diffusion-pipeline-profiler", + action="store_true", + help="Enable diffusion pipeline profiler to display stage durations.", + ) return parser.parse_args() @@ -133,7 +138,10 @@ def main(): print(f"{'=' * 60}\n") # Initialize Omni with Stable Audio model - omni = Omni(model=args.model) + omni = Omni( + model=args.model, + enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, + ) # Calculate audio end time audio_end_in_s = args.audio_start + args.audio_length diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index 37c296285..e2463619a 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -217,6 +217,11 @@ def parse_args() -> argparse.Namespace: action="store_true", help="[NextStep-1.1 only] Apply layer normalization to sampled tokens.", ) + parser.add_argument( + "--enable-diffusion-pipeline-profiler", + action="store_true", + help="Enable diffusion pipeline profiler to display stage durations.", + ) return parser.parse_args() @@ -304,6 +309,7 @@ def main(): "parallel_config": parallel_config, "enforce_eager": args.enforce_eager, "enable_cpu_offload": args.enable_cpu_offload, + "enable_diffusion_pipeline_profiler": args.enable_diffusion_pipeline_profiler, **lora_args, **quant_kwargs, } diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index 148652a13..e0493d4e6 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -132,6 +132,11 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable expert parallelism for MoE layers.", ) + parser.add_argument( + "--enable-diffusion-pipeline-profiler", + action="store_true", + help="Enable diffusion pipeline profiler to display stage durations.", + ) return parser.parse_args() @@ -184,6 +189,7 @@ def main(): enforce_eager=args.enforce_eager, cache_backend=args.cache_backend, cache_config=cache_config, + enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, ) if profiler_enabled: diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index b67736c81..c22b163f6 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -461,6 +461,9 @@ class OmniDiffusionConfig: quantization: str | None = None quantization_config: "DiffusionQuantizationConfig | dict[str, Any] | None" = None + # Diffusion pipeline Profiling config + enable_diffusion_pipeline_profiler: bool = False + @property def is_moe(self) -> bool: num_experts = self.tf_model_config.get("num_experts", None) @@ -642,6 +645,9 @@ class DiffusionOutput: # logged timings info, directly from Req.timings # timings: Optional["RequestTimings"] = None + # logged duration of stages + stage_durations: dict[str, float] = field(default_factory=dict) + class AttentionBackendEnum(enum.Enum): FA = enum.auto() diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 34f18f4df..b80d96bfe 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -170,6 +170,7 @@ class DiffusionEngine: latents=output.trajectory_latents, multimodal_output={"audio": request_audio_payload}, final_output_type="audio", + stage_durations=output.stage_durations, ), ] else: @@ -185,6 +186,7 @@ class DiffusionEngine: latents=output.trajectory_latents, custom_output=output.custom_output or {}, multimodal_output=mm_output, + stage_durations=output.stage_durations, ), ] else: @@ -214,6 +216,7 @@ class DiffusionEngine: latents=output.trajectory_latents, multimodal_output={"audio": request_audio_payload}, final_output_type="audio", + stage_durations=output.stage_durations, ), ) else: @@ -239,6 +242,7 @@ class DiffusionEngine: latents=output.trajectory_latents, custom_output=output.custom_output or {}, multimodal_output=mm_output, + stage_durations=output.stage_durations, ), ) diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index 9ed3481d7..dc8a91b3b 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py @@ -26,6 +26,7 @@ from vllm.transformers_utils.configs.bagel import BagelConfig from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific @@ -147,7 +148,7 @@ class SiglipNaViTWrapper(nn.Module): return outputs.last_hidden_state.squeeze(0) -class BagelPipeline(nn.Module): +class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin): """Bagel generation pipeline (MoT) packaged for vllm-omni diffusion engine. This pipeline is self-contained and uses the ported Bagel core files. @@ -255,6 +256,9 @@ class BagelPipeline(nn.Module): ] self.to(self.device) + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) @staticmethod def _decode_image_from_latent( @@ -621,7 +625,9 @@ class BagelPipeline(nn.Module): ) img = self._decode_image_from_latent(self.bagel, self.vae, latents[0], image_shape) - return DiffusionOutput(output=img) + return DiffusionOutput( + output=img, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: state = self.state_dict() diff --git a/vllm_omni/diffusion/models/flux/pipeline_flux.py b/vllm_omni/diffusion/models/flux/pipeline_flux.py index 3955fee12..06af7c95d 100644 --- a/vllm_omni/diffusion/models/flux/pipeline_flux.py +++ b/vllm_omni/diffusion/models/flux/pipeline_flux.py @@ -27,6 +27,7 @@ from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_g from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.flux import FluxTransformer2DModel +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.quantization import get_vllm_quant_config_for_layers from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific @@ -131,7 +132,7 @@ def retrieve_timesteps( return timesteps, num_inference_steps -class FluxPipeline(nn.Module, CFGParallelMixin): +class FluxPipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin): def __init__( self, *, @@ -187,6 +188,10 @@ class FluxPipeline(nn.Module, CFGParallelMixin): ) self.default_sample_size = 128 + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) + def check_inputs( self, prompt, @@ -749,7 +754,9 @@ class FluxPipeline(nn.Module, CFGParallelMixin): latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor image = self.vae.decode(latents, return_dict=False)[0] - return DiffusionOutput(output=image) + return DiffusionOutput( + output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) diff --git a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py index 551be4f06..6261ff000 100644 --- a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py +++ b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py @@ -43,6 +43,7 @@ from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import ( Flux2Transformer2DModel, ) from vllm_omni.diffusion.models.interface import SupportImageInput +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.quantization import get_vllm_quant_config_for_layers from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs @@ -180,7 +181,7 @@ def compute_empirical_mu(image_seq_len: int, num_steps: int) -> float: return float(mu) -class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput): +class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, DiffusionPipelineProfilerMixin): """Flux2 klein pipeline for text-to-image generation.""" support_image_input = True @@ -244,6 +245,9 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput): self._num_timesteps = None self._current_timestep = None self._interrupt = False + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) @staticmethod def _get_qwen3_prompt_embeds( @@ -991,7 +995,9 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput): latents = latents.to(self.vae.dtype) image = self.vae.decode(latents, return_dict=False)[0] - return DiffusionOutput(output=image) + return DiffusionOutput( + output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) diff --git a/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py b/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py index 8fca50694..375f7e7b8 100644 --- a/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py +++ b/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py @@ -46,6 +46,7 @@ from vllm_omni.diffusion.models.glm_image.glm_image_transformer import ( GlmImageKVCache, GlmImageTransformer2DModel, ) +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.model_executor.model_loader.weight_utils import ( @@ -237,7 +238,7 @@ def retrieve_latents( raise AttributeError("Could not access latents of provided encoder_output") -class GlmImagePipeline(nn.Module): +class GlmImagePipeline(nn.Module, DiffusionPipelineProfilerMixin): """ GLM-Image Pipeline for text-to-image and image-to-image generation. @@ -320,6 +321,9 @@ class GlmImagePipeline(nn.Module): # Get transformer config for patch size self._patch_size = getattr(self.transformer, "patch_size", 2) + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) # ==================== Input Validation ==================== @@ -867,7 +871,9 @@ class GlmImagePipeline(nn.Module): latents = latents * latents_std + latents_mean image = self.vae.decode(latents, return_dict=False, generator=generator)[0] - return DiffusionOutput(output=image) + return DiffusionOutput( + output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load transformer weights.""" diff --git a/vllm_omni/diffusion/models/helios/pipeline_helios.py b/vllm_omni/diffusion/models/helios/pipeline_helios.py index 9fc6dcf7f..8112cfd96 100644 --- a/vllm_omni/diffusion/models/helios/pipeline_helios.py +++ b/vllm_omni/diffusion/models/helios/pipeline_helios.py @@ -27,6 +27,7 @@ from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineL from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel from vllm_omni.diffusion.models.helios.scheduling_helios import HeliosScheduler from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.platforms import current_omni_platform @@ -141,7 +142,7 @@ def get_helios_pre_process_func( return pre_process_func -class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): +class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin): """Helios text-to-video / image-to-video / video-to-video pipeline for vllm-omni. Supports T2V, I2V (with image input), and V2V (with video input). @@ -228,6 +229,9 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): self._guidance_scale = None self._num_timesteps = None self._current_timestep = None + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) @property def guidance_scale(self): @@ -647,7 +651,9 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): else: output = history_video - return DiffusionOutput(output=output) + return DiffusionOutput( + output=output, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + ) def _stage1_sample( self, diff --git a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py index 09f409f31..76d3efa2f 100644 --- a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py +++ b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py @@ -28,6 +28,7 @@ from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import LongCatImageTransformer2DModel +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, @@ -199,7 +200,7 @@ def get_prompt_language(prompt): return "en" -class LongCatImagePipeline(nn.Module, CFGParallelMixin): +class LongCatImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin): def __init__( self, *, @@ -250,6 +251,9 @@ class LongCatImagePipeline(nn.Module, CFGParallelMixin): self.default_sample_size = 128 self.tokenizer_max_length = 512 + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) def rewire_prompt(self, prompt, device): prompt = [prompt] if isinstance(prompt, str) else prompt @@ -670,7 +674,9 @@ class LongCatImagePipeline(nn.Module, CFGParallelMixin): image = self.vae.decode(latents, return_dict=False)[0] - return DiffusionOutput(output=image) + return DiffusionOutput( + output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + ) 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/longcat_image/pipeline_longcat_image_edit.py b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py index 3ba5e3448..7eccf6863 100644 --- a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py +++ b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py @@ -34,6 +34,7 @@ from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import ( LongCatImageTransformer2DModel, ) from vllm_omni.diffusion.models.longcat_image.pipeline_longcat_image import calculate_shift +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.model_executor.model_loader.weight_utils import ( @@ -220,7 +221,7 @@ def split_quotation(prompt, quote_pairs=None): return result -class LongCatImageEditPipeline(nn.Module, CFGParallelMixin, SupportImageInput): +class LongCatImageEditPipeline(nn.Module, CFGParallelMixin, SupportImageInput, DiffusionPipelineProfilerMixin): def __init__( self, *, @@ -279,6 +280,9 @@ class LongCatImageEditPipeline(nn.Module, CFGParallelMixin, SupportImageInput): self.default_sample_size = 128 self.tokenizer_max_length = 512 + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) def _encode_prompt(self, prompt, image): raw_vl_input = self.image_processor_vl(images=image, return_tensors="pt") @@ -703,7 +707,9 @@ class LongCatImageEditPipeline(nn.Module, CFGParallelMixin, SupportImageInput): latents = latents.to(dtype=self.vae.dtype) image = self.vae.decode(latents, return_dict=False)[0] - return DiffusionOutput(output=image) + return DiffusionOutput( + output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + ) 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/nextstep_1_1/pipeline_nextstep_1_1.py b/vllm_omni/diffusion/models/nextstep_1_1/pipeline_nextstep_1_1.py index 4a20c28c5..4fa56ea93 100644 --- a/vllm_omni/diffusion/models/nextstep_1_1/pipeline_nextstep_1_1.py +++ b/vllm_omni/diffusion/models/nextstep_1_1/pipeline_nextstep_1_1.py @@ -33,6 +33,7 @@ from vllm_omni.diffusion.models.nextstep_1_1.modeling_nextstep import ( NextStepConfig, NextStepModel, ) +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, @@ -130,7 +131,7 @@ def get_nextstep11_post_process_func(od_config: OmniDiffusionConfig): return post_process_func -class NextStep11Pipeline(nn.Module): +class NextStep11Pipeline(nn.Module, DiffusionPipelineProfilerMixin): """ NextStep-1.1 Pipeline for text-to-image generation. @@ -232,6 +233,9 @@ class NextStep11Pipeline(nn.Module): allow_patterns_overrides=["model-*.safetensors", "model.safetensors"], ) ] + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) @property def device(self): @@ -702,7 +706,9 @@ class NextStep11Pipeline(nn.Module): sampled_images = self.vae.decode(latents.to(self.vae.dtype)).sample sampled_images = sampled_images.detach().cpu().to(torch.float32) - return DiffusionOutput(output=sampled_images) + return DiffusionOutput( + output=sampled_images, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load model weights.""" diff --git a/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py b/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py index 963f1c483..eb62b83b3 100644 --- a/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py +++ b/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py @@ -39,6 +39,7 @@ from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.ovis_image.ovis_image_transformer import OvisImageTransformer2DModel +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific @@ -140,7 +141,7 @@ def retrieve_timesteps( return timesteps, num_inference_steps -class OvisImagePipeline(nn.Module, CFGParallelMixin): +class OvisImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin): def __init__( self, *, @@ -188,6 +189,9 @@ class OvisImagePipeline(nn.Module, CFGParallelMixin): self.user_prompt_begin_id = 28 self.tokenizer_max_length = 256 + self.user_prompt_begin_id self.default_sample_size = 128 + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) def _get_messages( self, @@ -734,7 +738,9 @@ class OvisImagePipeline(nn.Module, CFGParallelMixin): latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor image = self.vae.decode(latents, return_dict=False)[0] - return DiffusionOutput(output=image) + return DiffusionOutput( + output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py index 02758ee19..8ec3d6b9f 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py @@ -31,6 +31,7 @@ from vllm_omni.diffusion.models.qwen_image.cfg_parallel import ( from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( QwenImageTransformer2DModel, ) +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.quantization import get_vllm_quant_config_for_layers from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs @@ -236,7 +237,7 @@ def apply_rotary_emb_qwen( return x_out.type_as(x) -class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin): +class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin): def __init__( self, *, @@ -292,6 +293,10 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin): self.prompt_template_encode_start_idx = 34 self.default_sample_size = 128 + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) + def check_inputs( self, prompt, @@ -713,7 +718,9 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin): image = self.vae.decode(latents, return_dict=False)[0][:, :, 0] # processed_image = self.image_processor.postprocess(image, output_type=output_type) - return DiffusionOutput(output=image) + return DiffusionOutput( + output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py index bff248a44..f805a7e7c 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py @@ -35,6 +35,7 @@ from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import calculate_ from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( QwenImageTransformer2DModel, ) +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs from vllm_omni.inputs.data import OmniTextPrompt @@ -215,7 +216,7 @@ def retrieve_latents( raise AttributeError("Could not access latents of provided encoder_output") -class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMixin): +class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin): def __init__( self, *, @@ -266,6 +267,9 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi self.prompt_template_encode = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n<|im_start|>assistant\n" # noqa: E501 self.prompt_template_encode_start_idx = 64 self.default_sample_size = 128 + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) def check_inputs( self, @@ -814,7 +818,9 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi latents = latents / latents_std + latents_mean image = self.vae.decode(latents, return_dict=False)[0][:, :, 0] - return DiffusionOutput(output=image) + return DiffusionOutput( + output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py index 2fdc7003e..8e2ba90a4 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py @@ -38,6 +38,7 @@ from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import ( from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( QwenImageTransformer2DModel, ) +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs from vllm_omni.inputs.data import OmniTextPrompt @@ -166,7 +167,9 @@ def get_qwen_image_edit_plus_post_process_func( return post_process_func -class QwenImageEditPlusPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMixin): +class QwenImageEditPlusPipeline( + nn.Module, SupportImageInput, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin +): def __init__( self, *, @@ -224,6 +227,9 @@ class QwenImageEditPlusPipeline(nn.Module, SupportImageInput, QwenImageCFGParall ) self.prompt_template_encode_start_idx = 64 self.default_sample_size = 128 + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) def check_inputs( self, @@ -769,7 +775,9 @@ class QwenImageEditPlusPipeline(nn.Module, SupportImageInput, QwenImageCFGParall latents = latents / latents_std + latents_mean image = self.vae.decode(latents, return_dict=False)[0][:, :, 0] - return DiffusionOutput(output=image) + return DiffusionOutput( + output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py index dbe0bfe4f..6b041aece 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py @@ -34,6 +34,7 @@ from vllm_omni.diffusion.models.qwen_image.cfg_parallel import ( from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( QwenImageTransformer2DModel, ) +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs from vllm_omni.inputs.data import OmniTextPrompt @@ -193,7 +194,7 @@ def retrieve_latents( raise AttributeError("Could not access latents of provided encoder_output") -class QwenImageLayeredPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMixin): +class QwenImageLayeredPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin): color_format = "RGBA" def __init__( @@ -283,6 +284,9 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as self.default_sample_size = 128 self.stage = None + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) def check_inputs( self, @@ -843,7 +847,9 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as for bidx in range(b): images.append(image[bidx * f : (bidx + 1) * f]) - return DiffusionOutput(output=images) + return DiffusionOutput( + output=images, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) diff --git a/vllm_omni/diffusion/models/sd3/pipeline_sd3.py b/vllm_omni/diffusion/models/sd3/pipeline_sd3.py index 9c9722c36..86b2a187b 100644 --- a/vllm_omni/diffusion/models/sd3/pipeline_sd3.py +++ b/vllm_omni/diffusion/models/sd3/pipeline_sd3.py @@ -22,6 +22,7 @@ from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineL from vllm_omni.diffusion.models.sd3.sd3_transformer import ( SD3Transformer2DModel, ) +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, @@ -127,7 +128,7 @@ def retrieve_timesteps( return timesteps, num_inference_steps -class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin): +class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin): def __init__( self, *, @@ -187,6 +188,9 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin): self.default_sample_size = 128 self.patch_size = 2 self.output_type = self.od_config.output_type + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) def check_inputs( self, @@ -692,7 +696,9 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin): image = self.vae.decode(latents, return_dict=False)[0] - return DiffusionOutput(output=image) + return DiffusionOutput( + output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) diff --git a/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py b/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py index 22dfc06c5..7950fb491 100644 --- a/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py +++ b/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py @@ -29,6 +29,7 @@ from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.interface import SupportAudioOutput from vllm_omni.diffusion.models.stable_audio.stable_audio_transformer import StableAudioDiTModel +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs @@ -59,7 +60,7 @@ def get_stable_audio_post_process_func( return post_process_func -class StableAudioPipeline(nn.Module, SupportAudioOutput): +class StableAudioPipeline(nn.Module, SupportAudioOutput, DiffusionPipelineProfilerMixin): """ Pipeline for text-to-audio generation using Stable Audio Open. @@ -149,6 +150,9 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput): self._guidance_scale = None self._num_timesteps = None self._current_timestep = None + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) @property def guidance_scale(self): @@ -569,7 +573,9 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput): # Trim to requested length audio = audio[:, :, waveform_start:waveform_end] - return DiffusionOutput(output=audio) + return DiffusionOutput( + output=audio, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + ) 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.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index 9a0037a6a..d7d8bad52 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -25,6 +25,7 @@ from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineL from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.platforms import current_omni_platform @@ -189,7 +190,7 @@ def get_wan22_pre_process_func( return pre_process_func -class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): +class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin): def __init__( self, *, @@ -311,6 +312,10 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): self._num_timesteps = None self._current_timestep = None + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) + @property def guidance_scale(self): return self._guidance_scale @@ -699,7 +704,9 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): _t_pipeline_wall_ms - _t_stages_sum, ) - return DiffusionOutput(output=output) + return DiffusionOutput( + output=output, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + ) def predict_noise(self, current_model: nn.Module | None = None, **kwargs: Any) -> torch.Tensor: """ 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 fe15a24f5..1e8a94eb3 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 @@ -30,6 +30,7 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( load_transformer_config, retrieve_latents, ) +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.platforms import current_omni_platform @@ -141,7 +142,9 @@ def get_wan22_i2v_pre_process_func( return pre_process_func -class Wan22I2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, ProgressBarMixin): +class Wan22I2VPipeline( + nn.Module, SupportImageInput, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin +): """ Wan2.2 Image-to-Video Pipeline. @@ -249,6 +252,9 @@ class Wan22I2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, ProgressB self._guidance_scale_2 = None self._num_timesteps = None self._current_timestep = None + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) @property def guidance_scale(self): @@ -610,7 +616,9 @@ class Wan22I2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, ProgressB _t_pipeline_wall_ms - _t_stages_sum, ) - return DiffusionOutput(output=output) + return DiffusionOutput( + output=output, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + ) def predict_noise(self, current_model: nn.Module | None = None, **kwargs: Any) -> torch.Tensor: """ diff --git a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py index b16d24bf7..0ccb12cfe 100644 --- a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py +++ b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py @@ -37,6 +37,7 @@ from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineL from vllm_omni.diffusion.models.z_image.z_image_transformer import ( ZImageTransformer2DModel, ) +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.quantization import get_vllm_quant_config_for_layers from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.model_executor.model_loader.weight_utils import ( @@ -143,7 +144,7 @@ def retrieve_timesteps( return timesteps, num_inference_steps -class ZImagePipeline(nn.Module): +class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin): def __init__( self, *, @@ -185,6 +186,9 @@ class ZImagePipeline(nn.Module): self.vae_scale_factor = ( 2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 8 ) + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) def encode_prompt( self, @@ -640,7 +644,9 @@ class ZImagePipeline(nn.Module): image = self.vae.decode(latents, return_dict=False)[0] # image = self.image_processor.postprocess(image, output_type=output_type) - return DiffusionOutput(output=image) + return DiffusionOutput( + output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) diff --git a/vllm_omni/diffusion/profiler/diffusion_pipeline_profiler.py b/vllm_omni/diffusion/profiler/diffusion_pipeline_profiler.py new file mode 100644 index 000000000..399778cb4 --- /dev/null +++ b/vllm_omni/diffusion/profiler/diffusion_pipeline_profiler.py @@ -0,0 +1,101 @@ +import functools +import time +from collections.abc import Callable +from threading import Lock +from typing import Any + +from vllm.logger import init_logger + +from vllm_omni.platforms import current_omni_platform + +logger = init_logger(__name__) + + +def profiler(name: str, func: Callable, instance: Any) -> Callable: + """Timing a function execution.""" + + @functools.wraps(func) + def wrapper(*args, **kwargs) -> Any: + if name == f"{instance.__class__.__name__}.forward": + instance.clear_profiler_records() + if current_omni_platform.is_available(): + current_omni_platform.synchronize() + start_time = time.perf_counter() + try: + return func(*args, **kwargs) + finally: + if current_omni_platform.is_available(): + current_omni_platform.synchronize() + duration = time.perf_counter() - start_time + logger.info(f"[DiffusionPipelineProfiler] {name} took {duration:.6f}s") + # record the profiling data: duration of stages + with instance._profiler_lock: + instance._stage_durations[name] = duration + + return wrapper + + +def _get_attribute_by_path(obj: Any, path: str) -> tuple[Any, str]: + """Traverse an object by dotted path and return (parent_obj, attribute_name).""" + parts = path.split(".") + current = obj + + for part in parts[:-1]: + current = getattr(current, part, None) + if current is None: + return None, None + + return current, parts[-1] + + +def wrap_methods_by_paths(root_obj: Any, method_paths: list[str]) -> None: + """Wrap specified methods of an object with profiler.""" + for path in method_paths: + obj, method_name = _get_attribute_by_path(root_obj, path) + if not obj or not hasattr(obj, method_name): + logger.warning(f"[DiffusionPipelineProfiler] Method path {path} not found") + continue + + original_method = getattr(obj, method_name) + if not callable(original_method): + logger.warning(f"[DiffusionPipelineProfiler] Attribute {path} is not callable") + continue + + profiler_name = f"{root_obj.__class__.__name__}.{path}" + setattr(obj, method_name, profiler(profiler_name, original_method, root_obj)) + + +class DiffusionPipelineProfilerMixin: + _PROFILER_TARGETS = ["vae.encode", "vae.decode", "diffuse", "text_encoder.forward", "tokenizer.forward"] + + def setup_diffusion_pipeline_profiler( + self, profiler_targets: list[str] | None = None, enable_diffusion_pipeline_profiler: bool = False + ) -> None: + self.enable_diffusion_pipeline_profiler = enable_diffusion_pipeline_profiler + if not enable_diffusion_pipeline_profiler: + self.enable_diffusion_pipeline_profiler = enable_diffusion_pipeline_profiler + return + self._profiler_lock = Lock() + self._stage_durations: dict[str, float] = {} + targets = profiler_targets if profiler_targets is not None else self._PROFILER_TARGETS + if not targets: + return + + targets = ["forward"] + [ + t for t in targets if t != "forward" + ] # ensure "forward" implement 'clear_profiler_records' at first place + + targets = list(dict.fromkeys(targets)) + wrap_methods_by_paths( + self, + targets, + ) + + @property + def stage_durations(self) -> dict[str, float]: + with self._profiler_lock: + return self._stage_durations.copy() + + def clear_profiler_records(self) -> None: + with self._profiler_lock: + self._stage_durations.clear() diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 4a20e4bc4..50fb67597 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -844,6 +844,7 @@ class AsyncOmniEngine: "enable_sleep_mode": kwargs.get("enable_sleep_mode", False), "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True), "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4), + "enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False), }, "final_output": True, "final_output_type": "image", diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index a70ac64b6..6830fc44b 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -351,6 +351,13 @@ class OmniServeCommand(CLISubcommand): default=None, help="Maximum length for TTS voice style instructions (overrides stage config, default: 500).", ) + + # Enable diffusion pipeline profiling + omni_config_group.add_argument( + "--enable-diffusion-pipeline-profiler", + action="store_true", + help="Enable diffusion pipeline profiler to display stage durations.", + ) return serve_parser diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index fb5f3788d..0de412a96 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -238,6 +238,7 @@ class OmniBase: ) -> OmniRequestOutput | None: req_id = result.get("request_id") engine_outputs = result.get("engine_outputs") + stage_durations = getattr(result["engine_outputs"], "stage_durations", {}) finished = engine_outputs.finished submit_ts = result.get("stage_submit_ts") @@ -271,6 +272,7 @@ class OmniBase: final_output_type=stage_meta["final_output_type"], request_output=engine_outputs, images=images, + stage_durations=stage_durations, ) def shutdown(self) -> None: diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 1fc3cdb36..d0a2d5aaa 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -1906,6 +1906,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): choices: list[ChatCompletionResponseChoice] = [] final_res = omni_outputs.request_output + # Handle profiling data + stage_durations = omni_outputs.stage_durations + # Handle different image output formats images = [] @@ -1959,6 +1962,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): "image_url": { "url": f"data:image/png;base64,{img_base64}", }, + "stage_durations": stage_durations, } ) @@ -2160,6 +2164,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # Extract images from result # Handle nested OmniRequestOutput structure where images might be in request_output images = getattr(result.request_output, "images", []) + stage_durations = result.stage_durations # Convert images to base64 content image_contents: list[dict[str, Any]] = [] @@ -2174,6 +2179,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): "image_url": { "url": f"data:image/png;base64,{img_base64}", }, + "stage_durations": stage_durations, } ) diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py index be65ec98f..10b1e50ed 100644 --- a/vllm_omni/outputs.py +++ b/vllm_omni/outputs.py @@ -62,6 +62,9 @@ class OmniRequestOutput: _multimodal_output: dict[str, Any] = field(default_factory=dict) _custom_output: dict[str, Any] = field(default_factory=dict) + # profiling data + stage_durations: dict[str, float] = field(default_factory=dict) + @classmethod def from_pipeline( cls, @@ -98,6 +101,7 @@ class OmniRequestOutput: multimodal_output: dict[str, Any] | None = None, custom_output: dict[str, Any] | None = None, final_output_type: str = "image", + stage_durations: dict[str, float] | None = None, ) -> "OmniRequestOutput": """Create output from diffusion model. @@ -122,6 +126,7 @@ class OmniRequestOutput: metrics=metrics or {}, _multimodal_output=multimodal_output or {}, _custom_output=custom_output or {}, + stage_durations=stage_durations or {}, finished=True, ) @@ -271,6 +276,7 @@ class OmniRequestOutput: f"metrics={self.metrics}", f"multimodal_output={self._multimodal_output}", f"custom_output={self._custom_output}", + f"stage_durations={self.stage_durations}", ] return f"OmniRequestOutput({', '.join(parts)})"