Files
vllm-omni/vllm_omni/diffusion/diffusion_engine.py
T
c708aecb18 [Diffusion] [Model] Support AudioX (#2077)
Signed-off-by: Zhang Jian <jianmusings@gmail.com>
Signed-off-by: Zhang <jianmusings@gmail.com>
Signed-off-by: Zhang <zhang.jian@u.nus.edu>
Signed-off-by: Zhang Jian <e0322744@u.nus.edu>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: Zhang Jian <e0322744@u.nus.edu>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
2026-05-04 10:24:35 -04:00

567 lines
24 KiB
Python

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
import inspect
import queue
import threading
import time
from collections.abc import Iterable
from typing import Any
import numpy as np
import PIL.Image
import torch
from vllm.logger import init_logger
from vllm.v1.engine.exceptions import EngineDeadError
from vllm_omni.diffusion.data import (
DiffusionOutput,
DiffusionRequestAbortedError,
OmniDiffusionConfig,
)
from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
from vllm_omni.diffusion.registry import (
DiffusionModelRegistry,
get_diffusion_post_process_func,
get_diffusion_pre_process_func,
)
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched import RequestScheduler, SchedulerInterface, StepScheduler
from vllm_omni.diffusion.sched.interface import DiffusionRequestStatus
from vllm_omni.diffusion.worker.utils import RunnerOutput
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
from vllm_omni.outputs import OmniRequestOutput
logger = init_logger(__name__)
def supports_image_input(model_class_name: str) -> bool:
model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
if model_cls is None:
return False
return bool(getattr(model_cls, "support_image_input", False))
def supports_audio_input(model_class_name: str) -> bool:
model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
if model_cls is None:
return False
return bool(getattr(model_cls, "support_audio_input", False))
def image_color_format(model_class_name: str) -> str:
model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
return getattr(model_cls, "color_format", "RGB")
def supports_audio_output(model_class_name: str) -> bool:
model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
if model_cls is None:
return False
return bool(getattr(model_cls, "support_audio_output", False))
class DiffusionEngine:
"""The diffusion engine for vLLM-Omni diffusion models."""
def __init__(
self,
od_config: OmniDiffusionConfig,
scheduler: SchedulerInterface | None = None,
):
"""Initialize the diffusion engine.
Args:
config: The configuration for the diffusion engine.
"""
self.od_config = od_config
self.post_process_func = get_diffusion_post_process_func(od_config)
self.pre_process_func = get_diffusion_pre_process_func(od_config)
# Cache whether the model-specific postprocess accepts request-level
# sampling params so step() can support both legacy and extended hooks.
self._post_process_accepts_sampling_params = bool(
self.post_process_func is not None
and "sampling_params" in inspect.signature(self.post_process_func).parameters
)
executor_class = DiffusionExecutor.get_class(od_config)
self.executor = executor_class(od_config)
self.step_execution = bool(getattr(od_config, "step_execution", False))
self.scheduler: SchedulerInterface = scheduler or (
StepScheduler() if self.step_execution else RequestScheduler()
)
self.scheduler.initialize(od_config)
self._rpc_lock = threading.RLock()
self.abort_queue: queue.Queue[str] = queue.Queue()
self.execute_fn = self.executor.execute_step if self.step_execution else self.executor.execute_request
try:
self._dummy_run()
except Exception as e:
logger.error(f"Dummy run failed: {e}")
self.close()
raise e
def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
diffusion_engine_start_time = time.perf_counter()
# Apply pre-processing if available
preprocess_time = 0.0
if self.pre_process_func is not None:
preprocess_start_time = time.perf_counter()
request = self.pre_process_func(request)
preprocess_time = time.perf_counter() - preprocess_start_time
logger.info(f"Pre-processing completed in {preprocess_time:.4f} seconds")
exec_start_time = time.perf_counter()
output = self.add_req_and_wait_for_response(request)
exec_total_time = time.perf_counter() - exec_start_time
if output.aborted:
raise DiffusionRequestAbortedError(output.abort_message or "Diffusion request aborted.")
if output.error:
raise RuntimeError(output.error)
logger.info("Generation completed successfully.")
if output.output is None:
logger.warning("Output is None, returning empty OmniRequestOutput")
return [
OmniRequestOutput.from_diffusion(
request_id=request.request_ids[i] if i < len(request.request_ids) else "",
images=[],
prompt=prompt,
metrics={},
latents=None,
)
for i, prompt in enumerate(request.prompts)
]
# When CPU offload is enabled, move output to CPU before
# post-processing to avoid device OOM — model weights may still
# reside on the device and leave no headroom for intermediates.
output_data = output.output
if (
self.od_config.enable_cpu_offload
and isinstance(output_data, torch.Tensor)
and output_data.device.type != "cpu"
):
output_data = output_data.cpu()
postprocess_start_time = time.perf_counter()
if self.post_process_func is not None:
# Some video pipelines need request-level controls during
# postprocess (for example worker-side frame interpolation).
if self._post_process_accepts_sampling_params:
outputs = self.post_process_func(output_data, sampling_params=request.sampling_params)
else:
outputs = self.post_process_func(output_data)
else:
outputs = output_data
audio_payload = None
custom_output = output.custom_output or {}
model_audio_sample_rate = None
model_fps = None
if isinstance(outputs, dict):
audio_payload = outputs.get("audio")
custom_output.update(outputs.get("custom_output") or {})
model_audio_sample_rate = outputs.get("audio_sample_rate")
model_fps = outputs.get("fps")
outputs = outputs.get("video", outputs)
postprocess_time = time.perf_counter() - postprocess_start_time
logger.info(f"Post-processing completed in {postprocess_time:.4f} seconds")
step_total_ms = (time.perf_counter() - diffusion_engine_start_time) * 1000
logger.info(
"DiffusionEngine.step breakdown: preprocess=%.2f ms, "
"add_req_and_wait=%.2f ms, postprocess=%.2f ms, total=%.2f ms",
preprocess_time * 1000,
exec_total_time * 1000,
postprocess_time * 1000,
step_total_ms,
)
# Convert to OmniRequestOutput format
# Ensure outputs is a list
if not isinstance(outputs, list):
outputs = [outputs] if outputs is not None else []
metrics = {
"preprocess_time_ms": preprocess_time * 1000,
"diffusion_engine_exec_time_ms": exec_total_time * 1000,
"diffusion_engine_total_time_ms": step_total_ms,
"image_num": int(request.sampling_params.num_outputs_per_prompt),
"resolution": int(request.sampling_params.resolution),
"postprocess_time_ms": postprocess_time * 1000,
}
# Handle single request or multiple requests
is_audio_output = supports_audio_output(self.od_config.model_class_name)
if is_audio_output and model_audio_sample_rate is None:
model_cls = DiffusionModelRegistry._try_load_model_cls(self.od_config.model_class_name)
model_audio_sample_rate = getattr(model_cls, "audio_sample_rate", None)
def _audio_mm(payload: Any) -> dict[str, Any]:
mm: dict[str, Any] = {"audio": payload}
if model_audio_sample_rate is not None:
mm["audio_sample_rate"] = model_audio_sample_rate
return mm
if len(request.prompts) == 1:
# Single request: return single OmniRequestOutput
prompt = request.prompts[0]
request_id = request.request_ids[0] if request.request_ids else ""
if is_audio_output:
request_audio_payload = outputs[0] if len(outputs) == 1 else outputs
return [
OmniRequestOutput.from_diffusion(
request_id=request_id,
images=[],
prompt=prompt,
metrics=metrics,
latents=output.trajectory_latents,
trajectory_latents=output.trajectory_latents,
trajectory_timesteps=output.trajectory_timesteps,
trajectory_log_probs=output.trajectory_log_probs,
trajectory_decoded=output.trajectory_decoded,
multimodal_output=_audio_mm(request_audio_payload),
final_output_type="audio",
stage_durations=output.stage_durations,
peak_memory_mb=output.peak_memory_mb,
),
]
else:
mm_output = {}
if audio_payload is not None:
mm_output["audio"] = audio_payload
if model_audio_sample_rate is not None:
mm_output["audio_sample_rate"] = model_audio_sample_rate
if model_fps is not None:
mm_output["fps"] = model_fps
return [
OmniRequestOutput.from_diffusion(
request_id=request_id,
images=outputs,
prompt=prompt,
metrics=metrics,
latents=output.trajectory_latents,
trajectory_latents=output.trajectory_latents,
trajectory_timesteps=output.trajectory_timesteps,
trajectory_log_probs=output.trajectory_log_probs,
trajectory_decoded=output.trajectory_decoded,
custom_output=custom_output,
multimodal_output=mm_output,
stage_durations=output.stage_durations,
peak_memory_mb=output.peak_memory_mb,
),
]
else:
# Multiple requests: return list of OmniRequestOutput
# Split images based on num_outputs_per_prompt for each request
results = []
output_idx = 0
for i, prompt in enumerate(request.prompts):
request_id = request.request_ids[i] if i < len(request.request_ids) else ""
# Get images for this request
num_outputs = request.sampling_params.num_outputs_per_prompt
start_idx = output_idx
end_idx = start_idx + num_outputs
request_outputs = outputs[start_idx:end_idx] if output_idx < len(outputs) else []
output_idx = end_idx
if is_audio_output:
request_audio_payload = request_outputs[0] if len(request_outputs) == 1 else request_outputs
results.append(
OmniRequestOutput.from_diffusion(
request_id=request_id,
images=[],
prompt=prompt,
metrics=metrics,
latents=output.trajectory_latents,
trajectory_latents=output.trajectory_latents,
trajectory_timesteps=output.trajectory_timesteps,
trajectory_log_probs=output.trajectory_log_probs,
trajectory_decoded=output.trajectory_decoded,
multimodal_output=_audio_mm(request_audio_payload),
final_output_type="audio",
stage_durations=output.stage_durations,
peak_memory_mb=output.peak_memory_mb,
),
)
else:
mm_output = {}
if audio_payload is not None:
sliced_audio = audio_payload
if isinstance(audio_payload, (list, tuple)):
sliced_audio = audio_payload[start_idx:end_idx]
if len(sliced_audio) == 1:
sliced_audio = sliced_audio[0]
elif hasattr(audio_payload, "shape") and getattr(audio_payload, "shape", None) is not None:
if len(audio_payload.shape) > 0 and audio_payload.shape[0] >= end_idx:
sliced_audio = audio_payload[start_idx:end_idx]
if num_outputs == 1:
sliced_audio = sliced_audio[0]
mm_output["audio"] = sliced_audio
if model_audio_sample_rate is not None:
mm_output["audio_sample_rate"] = model_audio_sample_rate
if model_fps is not None:
mm_output["fps"] = model_fps
results.append(
OmniRequestOutput.from_diffusion(
request_id=request_id,
images=request_outputs,
prompt=prompt,
metrics=metrics,
latents=output.trajectory_latents,
trajectory_latents=output.trajectory_latents,
trajectory_timesteps=output.trajectory_timesteps,
trajectory_log_probs=output.trajectory_log_probs,
trajectory_decoded=output.trajectory_decoded,
custom_output=custom_output,
multimodal_output=mm_output,
stage_durations=output.stage_durations,
peak_memory_mb=output.peak_memory_mb,
),
)
return results
@staticmethod
def make_engine(
config: OmniDiffusionConfig,
scheduler: SchedulerInterface | None = None,
) -> DiffusionEngine:
"""Factory method to create a DiffusionEngine instance.
Args:
config: The configuration for the diffusion engine.
Returns:
An instance of DiffusionEngine.
"""
return DiffusionEngine(config, scheduler=scheduler)
def add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput:
with self._rpc_lock:
target_sched_req_id = self.scheduler.add_request(request)
# keep scheduling and executing until the target request is finished
while True:
self._process_aborts_queue()
sched_output = self.scheduler.schedule()
if sched_output.is_empty:
if target_sched_req_id in sched_output.finished_req_ids:
return self._finalize_finished_request(target_sched_req_id)
if not self.scheduler.has_requests():
raise RuntimeError("Diffusion scheduler has no runnable requests.")
continue
# NOTE: add_req_and_wait_for_response() is synchronous, and
# the scheduler currently enforces _max_batch_size = 1 (see
# vllm_omni/diffusion/sched/base_scheduler.py), so we directly
# take the single scheduled request here.
sched_req_id = sched_output.scheduled_req_ids[0]
try:
runner_output = self.execute_fn(sched_output)
except EngineDeadError:
raise
except Exception as exc:
logger.error("Execution failed for diffusion request %s", sched_req_id, exc_info=True)
runner_output = RunnerOutput(
req_id=sched_req_id,
step_index=None,
finished=True,
result=DiffusionOutput(error=str(exc)),
)
self._process_aborts_queue()
finished_req_ids = self.scheduler.update_from_output(sched_output, runner_output)
if target_sched_req_id in finished_req_ids:
return self._finalize_finished_request(
target_sched_req_id,
runner_output=runner_output,
missing_result_error="Diffusion execution finished without a final output.",
)
def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:
"""Start or stop profiling on all diffusion workers.
Args:
is_start: True to start profiling, False to stop.
profile_prefix: Optional prefix for trace filename.
"""
if is_start:
if profile_prefix is None:
profile_prefix = f"diffusion_{int(time.time())}"
logger.info(f"Starting diffusion profiling with prefix: {profile_prefix}")
else:
logger.info("Stopping diffusion profiling...")
try:
self.collective_rpc(method="profile", args=(is_start, profile_prefix))
except Exception as e:
action = "start" if is_start else "stop"
logger.error(f"Failed to {action} profiling on workers", exc_info=True)
if is_start:
raise RuntimeError(f"Could not {action} profiler: {e}") from e
def _dummy_run(self):
"""A dummy run to warm up the model."""
num_inference_steps = 1
height = 512
width = 512
if supports_image_input(self.od_config.model_class_name):
# Provide a dummy image input if the model supports it
color_format = image_color_format(self.od_config.model_class_name)
dummy_image = PIL.Image.new(color_format, (width, height))
else:
dummy_image = None
if supports_audio_input(self.od_config.model_class_name):
audio_sr = 16000
dummy_audio = np.random.randn(audio_sr * 2).astype(np.float32)
else:
dummy_audio = None
prompt: OmniTextPrompt = {
"prompt": "dummy run",
"multi_modal_data": {"image": dummy_image, "audio": dummy_audio},
}
req = OmniDiffusionRequest(
prompts=[prompt],
request_ids=["dummy_req_id"],
sampling_params=OmniDiffusionSamplingParams(
height=height,
width=width,
num_inference_steps=num_inference_steps,
# Keep warmup path minimal and robust across text encoders.
# Some models may fail when warmup implicitly triggers
# classifier-free guidance with an empty negative prompt.
guidance_scale=0.0,
num_outputs_per_prompt=1,
# Disable CFG for warmup to avoid triggering CFG parallel
# validation when cfg_parallel_size > 1.
extra_args={"cfg_text_scale": 1.0, "cfg_img_scale": 1.0},
),
)
logger.info("dummy run to warm up the model")
request = self.pre_process_func(req) if self.pre_process_func is not None else req
output = self.add_req_and_wait_for_response(request)
if output.error:
raise RuntimeError(f"Dummy run failed: {output.error}")
def collective_rpc(
self,
method: str,
timeout: float | None = None,
args: tuple = (),
kwargs: dict | None = None,
unique_reply_rank: int | None = None,
) -> Any:
"""Call a method on worker processes and get results immediately.
Args:
method: The method name (str) to execute on workers
timeout: Optional timeout in seconds
args: Positional arguments for the method
kwargs: Keyword arguments for the method
unique_reply_rank: If set, only get reply from this rank
Returns:
Single result if unique_reply_rank is provided, otherwise list of results
"""
assert isinstance(method, str), "Only string method names are supported for now"
deadline = None if timeout is None else time.monotonic() + timeout
acquired = False
try:
if deadline is None:
self._rpc_lock.acquire()
acquired = True
else:
lock_timeout = max(0, deadline - time.monotonic())
acquired = self._rpc_lock.acquire(timeout=lock_timeout)
if not acquired:
raise TimeoutError(f"RPC call to {method} timed out waiting for engine lock.")
rpc_timeout = None if deadline is None else max(0, deadline - time.monotonic())
if deadline is not None and rpc_timeout <= 0:
raise TimeoutError(f"RPC call to {method} timed out.")
return self.executor.collective_rpc(
method=method,
timeout=rpc_timeout,
args=args,
kwargs=kwargs,
unique_reply_rank=unique_reply_rank,
)
finally:
if acquired:
self._rpc_lock.release()
def close(self) -> None:
if hasattr(self, "scheduler"):
self.scheduler.close()
if hasattr(self, "executor"):
self.executor.shutdown()
def abort(self, request_id: str | Iterable[str]) -> None:
request_ids = [request_id] if isinstance(request_id, str) else list(request_id)
for req_id in request_ids:
self.abort_queue.put(req_id)
def _process_aborts_queue(self) -> None:
if self.abort_queue.empty():
return
request_ids: list[str] = []
while not self.abort_queue.empty():
ids = self.abort_queue.get_nowait()
request_ids.extend((ids,) if isinstance(ids, str) else ids)
self._abort_requests(request_ids)
def _abort_requests(self, request_ids: str | Iterable[str]) -> None:
request_ids = [request_ids] if isinstance(request_ids, str) else list(request_ids)
sched_req_ids: list[str] = []
for request_id in dict.fromkeys(request_ids):
sched_req_id = self.scheduler.get_sched_req_id(request_id)
if sched_req_id is not None:
sched_req_ids.append(sched_req_id)
for sched_req_id in dict.fromkeys(sched_req_ids):
if self.scheduler.get_request_state(sched_req_id) is not None:
self.scheduler.finish_requests(sched_req_id, DiffusionRequestStatus.FINISHED_ABORTED)
def _finalize_finished_request(
self,
sched_req_id: str,
runner_output: RunnerOutput | None = None,
missing_result_error: str = "Diffusion scheduler finished target request without execution output.",
) -> DiffusionOutput:
state = self.scheduler.get_request_state(sched_req_id)
popped_state = self.scheduler.pop_request_state(sched_req_id)
state = state or popped_state
if state is None:
raise RuntimeError(f"Diffusion scheduler lost state for request {sched_req_id}.")
if state.status == DiffusionRequestStatus.FINISHED_ABORTED:
request_id = state.req.request_ids[0] if state.req.request_ids else sched_req_id
return DiffusionOutput(
aborted=True,
abort_message=f"Request {request_id} aborted.",
)
if runner_output is not None and runner_output.result is not None:
return runner_output.result
return DiffusionOutput(error=missing_result_error)