577 lines
23 KiB
Python
577 lines
23 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
import time
|
|
from dataclasses import dataclass
|
|
from http import HTTPStatus
|
|
from typing import Any, cast
|
|
|
|
from fastapi import HTTPException
|
|
from PIL import Image
|
|
from vllm.engine.protocol import EngineClient
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
|
from vllm_omni.entrypoints.openai.protocol.videos import (
|
|
VideoData,
|
|
VideoGenerationRequest,
|
|
VideoGenerationResponse,
|
|
)
|
|
from vllm_omni.entrypoints.openai.stage_params import (
|
|
build_stage_sampling_params_list,
|
|
get_default_sampling_params_list,
|
|
)
|
|
from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request
|
|
from vllm_omni.entrypoints.openai.video_api_utils import _encode_video_bytes, encode_video_base64
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class ReferenceImage:
|
|
"""Reference class for tracking additional metadata if needed"""
|
|
|
|
data: Image.Image
|
|
|
|
|
|
@dataclass
|
|
class VideoGenerationArtifacts:
|
|
"""Normalized outputs and profiler metadata extracted from one request."""
|
|
|
|
videos: list[Any]
|
|
audios: list[Any | None]
|
|
audio_sample_rate: int
|
|
output_fps: int
|
|
stage_durations: dict[str, float]
|
|
peak_memory_mb: float
|
|
|
|
|
|
class OmniOpenAIServingVideo:
|
|
"""OpenAI-style video generation handler for omni diffusion models."""
|
|
|
|
def __init__(
|
|
self,
|
|
engine_client: EngineClient,
|
|
model_name: str | None = None,
|
|
stage_configs: list[Any] | None = None,
|
|
) -> None:
|
|
self._engine_client = engine_client
|
|
self._model_name = model_name
|
|
self._stage_configs = stage_configs
|
|
|
|
@property
|
|
def model_name(self) -> str | None:
|
|
return self._model_name
|
|
|
|
@property
|
|
def stage_configs(self) -> list[Any] | None:
|
|
return self._stage_configs
|
|
|
|
def set_stage_configs_if_missing(self, stage_configs: list[Any] | None) -> None:
|
|
if self._stage_configs is None and stage_configs is not None:
|
|
self._stage_configs = stage_configs
|
|
|
|
@classmethod
|
|
def for_diffusion(
|
|
cls,
|
|
diffusion_engine: EngineClient,
|
|
model_name: str,
|
|
stage_configs: list[Any] | None = None,
|
|
) -> OmniOpenAIServingVideo:
|
|
return cls(
|
|
diffusion_engine,
|
|
model_name=model_name,
|
|
stage_configs=stage_configs,
|
|
)
|
|
|
|
async def _run_and_extract(
|
|
self,
|
|
request: VideoGenerationRequest,
|
|
reference_id: str,
|
|
*,
|
|
reference_image: ReferenceImage | None = None,
|
|
) -> VideoGenerationArtifacts:
|
|
"""Run the generation pipeline and extract video/audio/profiler outputs."""
|
|
prompt: OmniTextPrompt = OmniTextPrompt(prompt=request.prompt)
|
|
if request.negative_prompt is not None:
|
|
prompt["negative_prompt"] = request.negative_prompt
|
|
|
|
gen_params = self._resolve_default_sampling_params()
|
|
|
|
input_image = None if reference_image is None else reference_image.data
|
|
vp = request.resolve_video_params()
|
|
if input_image is not None and vp.width is not None and vp.height is not None:
|
|
target_size = (vp.width, vp.height)
|
|
if input_image.size != target_size:
|
|
input_image = input_image.resize(target_size, Image.Resampling.LANCZOS)
|
|
if input_image is not None:
|
|
prompt["multi_modal_data"] = {"image": input_image}
|
|
if vp.width is not None and vp.height is not None:
|
|
gen_params.width = vp.width
|
|
gen_params.height = vp.height
|
|
if vp.num_frames is not None:
|
|
gen_params.num_frames = vp.num_frames
|
|
if vp.fps is not None:
|
|
gen_params.fps = vp.fps
|
|
gen_params.frame_rate = float(vp.fps)
|
|
provided_fields = request.model_fields_set
|
|
if "enable_frame_interpolation" in provided_fields:
|
|
gen_params.enable_frame_interpolation = request.enable_frame_interpolation
|
|
if "frame_interpolation_exp" in provided_fields:
|
|
gen_params.frame_interpolation_exp = request.frame_interpolation_exp
|
|
if "frame_interpolation_scale" in provided_fields:
|
|
gen_params.frame_interpolation_scale = request.frame_interpolation_scale
|
|
if "frame_interpolation_model_path" in provided_fields:
|
|
gen_params.frame_interpolation_model_path = request.frame_interpolation_model_path
|
|
|
|
if "num_inference_steps" in provided_fields and request.num_inference_steps is not None:
|
|
gen_params.num_inference_steps = request.num_inference_steps
|
|
if "guidance_scale" in provided_fields and request.guidance_scale is not None:
|
|
gen_params.guidance_scale = request.guidance_scale
|
|
if "guidance_scale_2" in provided_fields and request.guidance_scale_2 is not None:
|
|
gen_params.guidance_scale_2 = request.guidance_scale_2
|
|
if "true_cfg_scale" in provided_fields and request.true_cfg_scale is not None:
|
|
gen_params.true_cfg_scale = request.true_cfg_scale
|
|
if "seed" in provided_fields and request.seed is not None:
|
|
gen_params.seed = request.seed
|
|
if "boundary_ratio" in provided_fields and request.boundary_ratio is not None:
|
|
gen_params.boundary_ratio = request.boundary_ratio
|
|
|
|
logger.info(
|
|
"Boundary ratio parse: request=%s gen_params=%s",
|
|
request.boundary_ratio,
|
|
gen_params.boundary_ratio,
|
|
)
|
|
if "flow_shift" in provided_fields and request.flow_shift is not None:
|
|
gen_params.extra_args["flow_shift"] = request.flow_shift
|
|
|
|
# Apply model-specific extra parameters
|
|
if request.extra_params is not None:
|
|
if not isinstance(request.extra_params, dict):
|
|
raise HTTPException(
|
|
status_code=HTTPStatus.BAD_REQUEST.value,
|
|
detail="extra_params must be a JSON object/dict.",
|
|
)
|
|
# Merge extra_params into extra_args
|
|
gen_params.extra_args.update(request.extra_params)
|
|
logger.info("Applied extra_params: %s", request.extra_params)
|
|
|
|
self._apply_lora(request.lora, gen_params)
|
|
|
|
logger.info(
|
|
"Video sampling params: steps=%s guidance=%s guidance_2=%s seed=%s",
|
|
gen_params.num_inference_steps,
|
|
gen_params.guidance_scale,
|
|
gen_params.guidance_scale_2,
|
|
gen_params.seed,
|
|
)
|
|
|
|
result = await self._run_generation(prompt, gen_params, reference_id)
|
|
videos = self._extract_video_outputs(result)
|
|
audios = self._extract_audio_outputs(result, expected_count=len(videos))
|
|
audio_sample_rate = self._resolve_audio_sample_rate(result)
|
|
output_fps = (vp.fps or self._resolve_fps(result) or 24) * self._resolve_video_fps_multiplier(result)
|
|
return VideoGenerationArtifacts(
|
|
videos=videos,
|
|
audios=audios,
|
|
audio_sample_rate=audio_sample_rate,
|
|
output_fps=output_fps,
|
|
stage_durations=self._extract_stage_durations(result),
|
|
peak_memory_mb=self._extract_peak_memory_mb(result),
|
|
)
|
|
|
|
async def generate_videos(
|
|
self,
|
|
request: VideoGenerationRequest,
|
|
reference_id: str,
|
|
*,
|
|
reference_image: ReferenceImage | None = None,
|
|
) -> VideoGenerationResponse:
|
|
artifacts = await self._run_and_extract(request, reference_id, reference_image=reference_image)
|
|
|
|
video_codec_options = {"preset": "ultrafast", "threads": "0"}
|
|
if request.extra_params is not None and isinstance(request.extra_params, dict):
|
|
if "video_codec_options" in request.extra_params:
|
|
video_codec_options = request.extra_params["video_codec_options"]
|
|
|
|
_t_encode_start = time.perf_counter()
|
|
video_data = [
|
|
VideoData(
|
|
b64_json=(
|
|
encode_video_base64(video, fps=artifacts.output_fps, video_codec_options=video_codec_options)
|
|
if artifacts.audios[idx] is None
|
|
else encode_video_base64(
|
|
video,
|
|
fps=artifacts.output_fps,
|
|
audio=artifacts.audios[idx],
|
|
audio_sample_rate=artifacts.audio_sample_rate,
|
|
video_codec_options=video_codec_options,
|
|
)
|
|
)
|
|
)
|
|
for idx, video in enumerate(artifacts.videos)
|
|
]
|
|
_t_encode_ms = (time.perf_counter() - _t_encode_start) * 1000
|
|
logger.info("Video response encoding (MP4+base64): %.2f ms", _t_encode_ms)
|
|
return VideoGenerationResponse(
|
|
created=int(time.time()),
|
|
data=video_data,
|
|
stage_durations=artifacts.stage_durations,
|
|
peak_memory_mb=artifacts.peak_memory_mb,
|
|
)
|
|
|
|
async def generate_video_bytes(
|
|
self,
|
|
request: VideoGenerationRequest,
|
|
reference_id: str,
|
|
*,
|
|
reference_image: ReferenceImage | None = None,
|
|
) -> tuple[bytes, dict[str, float], float]:
|
|
"""Generate a video and return raw MP4 bytes, bypassing base64 encoding."""
|
|
artifacts = await self._run_and_extract(request, reference_id, reference_image=reference_image)
|
|
if len(artifacts.videos) > 1:
|
|
logger.warning(
|
|
"Video request %s generated %d outputs; returning only the first.",
|
|
reference_id,
|
|
len(artifacts.videos),
|
|
)
|
|
audio = artifacts.audios[0]
|
|
|
|
video_codec_options = {"preset": "ultrafast", "threads": "0"}
|
|
if request.extra_params is not None and isinstance(request.extra_params, dict):
|
|
if "video_codec_options" in request.extra_params:
|
|
video_codec_options = request.extra_params["video_codec_options"]
|
|
|
|
_t_encode_start = time.perf_counter()
|
|
video_bytes = _encode_video_bytes(
|
|
artifacts.videos[0],
|
|
fps=artifacts.output_fps,
|
|
**({"audio": audio, "audio_sample_rate": artifacts.audio_sample_rate} if audio is not None else {}),
|
|
video_codec_options=video_codec_options,
|
|
)
|
|
_t_encode_ms = (time.perf_counter() - _t_encode_start) * 1000
|
|
logger.info("Video response encoding (MP4 bytes): %.2f ms", _t_encode_ms)
|
|
return video_bytes, artifacts.stage_durations, artifacts.peak_memory_mb
|
|
|
|
@staticmethod
|
|
def _resolve_video_fps_multiplier(result: Any) -> int:
|
|
custom_output = getattr(result, "custom_output", None)
|
|
if isinstance(custom_output, dict):
|
|
multiplier = custom_output.get("video_fps_multiplier")
|
|
if multiplier is not None:
|
|
return int(multiplier)
|
|
request_output = getattr(result, "request_output", None)
|
|
if request_output is not None:
|
|
custom_output = getattr(request_output, "custom_output", None)
|
|
if isinstance(custom_output, dict):
|
|
multiplier = custom_output.get("video_fps_multiplier")
|
|
if multiplier is not None:
|
|
return int(multiplier)
|
|
return 1
|
|
|
|
def _resolve_default_sampling_params(self) -> OmniDiffusionSamplingParams:
|
|
default_sampling_params_list = getattr(self._engine_client, "default_sampling_params_list", None)
|
|
if default_sampling_params_list:
|
|
for params in default_sampling_params_list:
|
|
if isinstance(params, OmniDiffusionSamplingParams):
|
|
# Requests mutate sampling params in-place, including
|
|
# nested dict fields like extra_args. Deep-copy the stage
|
|
# defaults so one request cannot leak state into another.
|
|
return copy.deepcopy(params)
|
|
return OmniDiffusionSamplingParams()
|
|
|
|
@staticmethod
|
|
def _apply_lora(lora_body: Any, gen_params: OmniDiffusionSamplingParams) -> None:
|
|
try:
|
|
lora_request, lora_scale = parse_lora_request(lora_body)
|
|
except ValueError as e:
|
|
raise HTTPException(
|
|
status_code=HTTPStatus.BAD_REQUEST.value,
|
|
detail=str(e),
|
|
) from e
|
|
|
|
if lora_request is None:
|
|
return
|
|
|
|
gen_params.lora_request = lora_request
|
|
if lora_scale is not None:
|
|
gen_params.lora_scale = lora_scale
|
|
|
|
async def _run_generation(
|
|
self,
|
|
prompt: OmniTextPrompt,
|
|
gen_params: OmniDiffusionSamplingParams,
|
|
request_id: str,
|
|
) -> Any:
|
|
stage_configs = self._stage_configs or getattr(self._engine_client, "stage_configs", None)
|
|
|
|
if not stage_configs:
|
|
raise HTTPException(
|
|
status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
|
|
detail="Stage configs not found. Start server with an omni diffusion model.",
|
|
)
|
|
|
|
# Video generation endpoint only supports diffusion stages.
|
|
for stage in stage_configs:
|
|
stage_type = get_stage_type(stage)
|
|
if stage_type != "diffusion":
|
|
raise HTTPException(
|
|
status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
|
|
detail=f"Video generation only supports diffusion stages, found '{stage_type}' stage.",
|
|
)
|
|
|
|
# Common generation logic for both paths
|
|
engine_client = cast(AsyncOmni, self._engine_client)
|
|
sampling_params_list = build_stage_sampling_params_list(
|
|
list(stage_configs),
|
|
get_default_sampling_params_list(engine_client),
|
|
diffusion_params=gen_params,
|
|
replace_diffusion_params=True,
|
|
)
|
|
|
|
result = None
|
|
async for output in engine_client.generate(
|
|
prompt=prompt,
|
|
request_id=request_id,
|
|
sampling_params_list=sampling_params_list,
|
|
):
|
|
result = output
|
|
|
|
if result is None:
|
|
raise HTTPException(
|
|
status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
|
|
detail="No output generated from video generation pipeline.",
|
|
)
|
|
return result
|
|
|
|
@staticmethod
|
|
def _normalize_video_outputs(videos: Any) -> list[Any]:
|
|
if videos is None:
|
|
return []
|
|
if hasattr(videos, "ndim") and videos.ndim == 5:
|
|
return [videos[i] for i in range(videos.shape[0])]
|
|
if isinstance(videos, list):
|
|
if not videos:
|
|
return []
|
|
first = videos[0]
|
|
if hasattr(first, "ndim") and first.ndim == 5:
|
|
flattened: list[Any] = []
|
|
for item in videos:
|
|
if hasattr(item, "ndim") and item.ndim == 5:
|
|
flattened.extend([item[i] for i in range(item.shape[0])])
|
|
else:
|
|
flattened.append(item)
|
|
return flattened
|
|
if isinstance(first, list):
|
|
return videos
|
|
if hasattr(first, "ndim") and first.ndim == 3:
|
|
return [videos]
|
|
if isinstance(first, Image.Image):
|
|
return [videos]
|
|
return videos
|
|
return [videos]
|
|
|
|
def _extract_video_outputs(self, result: Any) -> list[Any]:
|
|
videos = None
|
|
if hasattr(result, "images") and result.images:
|
|
videos = result.images
|
|
elif hasattr(result, "request_output"):
|
|
request_output = result.request_output
|
|
if isinstance(request_output, dict) and request_output.get("images"):
|
|
videos = request_output["images"]
|
|
elif hasattr(request_output, "images") and request_output.images:
|
|
videos = request_output.images
|
|
elif hasattr(request_output, "multimodal_output") and request_output.multimodal_output:
|
|
videos = request_output.multimodal_output.get("video")
|
|
if videos is None and hasattr(result, "multimodal_output") and result.multimodal_output:
|
|
videos = result.multimodal_output.get("video")
|
|
|
|
normalized = self._normalize_video_outputs(videos)
|
|
if not normalized:
|
|
raise HTTPException(
|
|
status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
|
|
detail="No video outputs found in generation result.",
|
|
)
|
|
return normalized
|
|
|
|
@staticmethod
|
|
def _extract_audio_outputs(result: Any, expected_count: int) -> list[Any | None]:
|
|
audio = None
|
|
if hasattr(result, "multimodal_output") and result.multimodal_output:
|
|
audio = result.multimodal_output.get("audio")
|
|
elif hasattr(result, "request_output"):
|
|
request_output = result.request_output
|
|
if isinstance(request_output, dict) and request_output.get("multimodal_output"):
|
|
mm_output = request_output.get("multimodal_output") or {}
|
|
audio = mm_output.get("audio")
|
|
elif hasattr(request_output, "multimodal_output") and request_output.multimodal_output:
|
|
audio = request_output.multimodal_output.get("audio")
|
|
|
|
if audio is None:
|
|
return [None] * expected_count
|
|
|
|
if isinstance(audio, (list, tuple)):
|
|
if len(audio) == expected_count and any(hasattr(item, "shape") or hasattr(item, "ndim") for item in audio):
|
|
return list(audio)
|
|
if expected_count == 1:
|
|
return [audio]
|
|
|
|
if hasattr(audio, "ndim") and getattr(audio, "ndim", None) is not None and audio.ndim > 1:
|
|
first_dim = getattr(audio, "shape", [0])[0]
|
|
if first_dim == expected_count:
|
|
return [audio[i] for i in range(expected_count)]
|
|
|
|
if expected_count == 1:
|
|
return [audio]
|
|
|
|
return [audio] + [None] * max(expected_count - 1, 0)
|
|
|
|
def _resolve_audio_sample_rate(self, result: Any) -> int:
|
|
result_sample_rate = self._extract_audio_sample_rate_from_result(result)
|
|
if result_sample_rate is not None:
|
|
return result_sample_rate
|
|
|
|
model_config = getattr(self._engine_client, "model_config", None)
|
|
hf_config = getattr(model_config, "hf_config", None)
|
|
config_sample_rate = self._extract_audio_sample_rate_from_config(hf_config)
|
|
if config_sample_rate is not None:
|
|
return config_sample_rate
|
|
|
|
return 24000
|
|
|
|
@staticmethod
|
|
def _resolve_fps(result: Any) -> int | None:
|
|
"""Extract fps from multimodal_output if the model reported it."""
|
|
multimodal_output = getattr(result, "multimodal_output", None)
|
|
if isinstance(multimodal_output, dict):
|
|
fps = multimodal_output.get("fps")
|
|
if fps is not None:
|
|
try:
|
|
fps_val = fps.item() if hasattr(fps, "item") else int(fps)
|
|
if fps_val > 0:
|
|
return fps_val
|
|
except (TypeError, ValueError):
|
|
pass
|
|
|
|
request_output = getattr(result, "request_output", None)
|
|
if isinstance(request_output, dict):
|
|
mm = request_output.get("multimodal_output") or {}
|
|
if isinstance(mm, dict):
|
|
fps = mm.get("fps")
|
|
if fps is not None:
|
|
try:
|
|
fps_val = fps.item() if hasattr(fps, "item") else int(fps)
|
|
if fps_val > 0:
|
|
return fps_val
|
|
except (TypeError, ValueError):
|
|
pass
|
|
elif hasattr(request_output, "multimodal_output"):
|
|
mm = getattr(request_output, "multimodal_output", None)
|
|
if isinstance(mm, dict):
|
|
fps = mm.get("fps")
|
|
if fps is not None:
|
|
try:
|
|
fps_val = fps.item() if hasattr(fps, "item") else int(fps)
|
|
if fps_val > 0:
|
|
return fps_val
|
|
except (TypeError, ValueError):
|
|
pass
|
|
|
|
return None
|
|
|
|
@classmethod
|
|
def _extract_audio_sample_rate_from_result(cls, result: Any) -> int | None:
|
|
multimodal_output = getattr(result, "multimodal_output", None)
|
|
if isinstance(multimodal_output, dict):
|
|
sample_rate = cls._coerce_audio_sample_rate(
|
|
multimodal_output.get("audio_sample_rate")
|
|
or multimodal_output.get("sample_rate")
|
|
or multimodal_output.get("sampling_rate")
|
|
or multimodal_output.get("sr")
|
|
)
|
|
if sample_rate is not None:
|
|
return sample_rate
|
|
|
|
request_output = getattr(result, "request_output", None)
|
|
if isinstance(request_output, dict):
|
|
multimodal_output = request_output.get("multimodal_output") or {}
|
|
if isinstance(multimodal_output, dict):
|
|
return cls._coerce_audio_sample_rate(
|
|
multimodal_output.get("audio_sample_rate")
|
|
or multimodal_output.get("sample_rate")
|
|
or multimodal_output.get("sampling_rate")
|
|
or multimodal_output.get("sr")
|
|
)
|
|
elif hasattr(request_output, "multimodal_output"):
|
|
multimodal_output = getattr(request_output, "multimodal_output", None)
|
|
if isinstance(multimodal_output, dict):
|
|
return cls._coerce_audio_sample_rate(
|
|
multimodal_output.get("audio_sample_rate")
|
|
or multimodal_output.get("sample_rate")
|
|
or multimodal_output.get("sampling_rate")
|
|
or multimodal_output.get("sr")
|
|
)
|
|
|
|
return None
|
|
|
|
@classmethod
|
|
def _extract_audio_sample_rate_from_config(cls, config: Any) -> int | None:
|
|
if config is None:
|
|
return None
|
|
|
|
for attr_name in ("output_sampling_rate", "audio_sample_rate", "sample_rate", "sampling_rate"):
|
|
raw_value = config.get(attr_name) if isinstance(config, dict) else getattr(config, attr_name, None)
|
|
sample_rate = cls._coerce_audio_sample_rate(raw_value)
|
|
if sample_rate is not None:
|
|
return sample_rate
|
|
|
|
for component_name in ("vocoder", "audio_vae"):
|
|
component = (
|
|
config.get(component_name) if isinstance(config, dict) else getattr(config, component_name, None)
|
|
)
|
|
if component is None:
|
|
continue
|
|
|
|
sample_rate = cls._extract_audio_sample_rate_from_config(component)
|
|
if sample_rate is not None:
|
|
return sample_rate
|
|
|
|
component_config = (
|
|
component.get("config") if isinstance(component, dict) else getattr(component, "config", None)
|
|
)
|
|
sample_rate = cls._extract_audio_sample_rate_from_config(component_config)
|
|
if sample_rate is not None:
|
|
return sample_rate
|
|
|
|
return None
|
|
|
|
@staticmethod
|
|
def _coerce_audio_sample_rate(value: Any) -> int | None:
|
|
if value is None:
|
|
return None
|
|
|
|
try:
|
|
sample_rate = value.item() if hasattr(value, "item") else value
|
|
sample_rate = int(sample_rate)
|
|
except (TypeError, ValueError):
|
|
return None
|
|
|
|
return sample_rate if sample_rate > 0 else None
|
|
|
|
@staticmethod
|
|
def _extract_stage_durations(result: Any) -> dict[str, float]:
|
|
stage_durations = getattr(result, "stage_durations", None)
|
|
return stage_durations if isinstance(stage_durations, dict) else {}
|
|
|
|
@staticmethod
|
|
def _extract_peak_memory_mb(result: Any) -> float:
|
|
peak_memory_mb = getattr(result, "peak_memory_mb", 0.0)
|
|
try:
|
|
return float(peak_memory_mb or 0.0)
|
|
except (TypeError, ValueError):
|
|
return 0.0
|