vraitiandClaude Opus 4.6
c98517ce7f
Signed-off-by: vraiti <vraiti@redhat.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
376 lines
14 KiB
Python
376 lines
14 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Mapping
|
|
from dataclasses import dataclass, field
|
|
from typing import TypeGuard, cast
|
|
|
|
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
|
from vllm_omni.diffusion.io_support import supports_audio_output
|
|
from vllm_omni.diffusion.registry import DiffusionModelRegistry
|
|
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
|
from vllm_omni.inputs.data import OmniPromptType
|
|
from vllm_omni.outputs import OmniRequestOutput
|
|
from vllm_omni.outputs.output_metadata import (
|
|
DiffusionMetadata,
|
|
DiffusionMultimodalOutput,
|
|
DiffusionOutputEnvelope,
|
|
DiffusionPayload,
|
|
DiffusionPayloadValue,
|
|
DiffusionPostprocessRawOutput,
|
|
DiffusionTrajectoryPayload,
|
|
strip_internal_metadata,
|
|
validate_diffusion_metadata,
|
|
validate_public_diffusion_metadata,
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DiffusionPostprocessOutput:
|
|
outputs: DiffusionPayload
|
|
metadata: DiffusionMetadata = field(default_factory=dict)
|
|
primary_key: str | None = None
|
|
|
|
|
|
def _is_output_envelope(outputs: DiffusionPostprocessRawOutput) -> TypeGuard[DiffusionOutputEnvelope]:
|
|
return isinstance(outputs, dict) and isinstance(outputs.get("payload"), dict)
|
|
|
|
|
|
def normalize_diffusion_postprocess_output(
|
|
outputs: DiffusionPostprocessRawOutput,
|
|
) -> DiffusionPostprocessOutput:
|
|
"""Normalize diffusion postprocess output into payload and metadata."""
|
|
|
|
if _is_output_envelope(outputs):
|
|
payload: DiffusionPayload = outputs.get("payload") or {}
|
|
metadata = outputs.get("metadata") or {}
|
|
if not isinstance(metadata, dict):
|
|
metadata = {}
|
|
validate_diffusion_metadata(metadata)
|
|
|
|
public_metadata = strip_internal_metadata(metadata)
|
|
if "text" in payload and "text" not in public_metadata:
|
|
public_metadata = {**public_metadata, "text": {}}
|
|
validate_public_diffusion_metadata(public_metadata)
|
|
|
|
return DiffusionPostprocessOutput(
|
|
outputs=payload,
|
|
metadata=public_metadata,
|
|
primary_key=_infer_primary_payload_key(payload),
|
|
)
|
|
|
|
if isinstance(outputs, dict):
|
|
payload = {key: value for key, value in outputs.items() if key not in {"audio_sample_rate", "fps"}}
|
|
metadata = _metadata_from_legacy_payload(outputs)
|
|
if "text" in payload and "text" not in metadata:
|
|
metadata["text"] = {}
|
|
if metadata:
|
|
validate_public_diffusion_metadata(metadata)
|
|
return DiffusionPostprocessOutput(
|
|
outputs=payload,
|
|
metadata=metadata,
|
|
primary_key=_infer_primary_payload_key(payload),
|
|
)
|
|
|
|
return DiffusionPostprocessOutput(
|
|
outputs={"output": outputs},
|
|
primary_key="output",
|
|
)
|
|
|
|
|
|
def _metadata_from_legacy_payload(payload: dict[str, DiffusionPayloadValue]) -> DiffusionMetadata:
|
|
metadata: DiffusionMetadata = {}
|
|
audio_sample_rate = payload.get("audio_sample_rate")
|
|
if audio_sample_rate is not None:
|
|
metadata["audio"] = {"sample_rate": audio_sample_rate}
|
|
fps = payload.get("fps")
|
|
if fps is not None:
|
|
metadata["video"] = {"fps": fps}
|
|
return metadata
|
|
|
|
|
|
def _infer_primary_payload_key(payload: DiffusionPayload) -> str | None:
|
|
for key in ("video", "image", "text", "audio", "output"):
|
|
if key in payload:
|
|
return key
|
|
if set(payload).issubset({"actions", "trajectory"}):
|
|
return None
|
|
if "actions" in payload:
|
|
return None
|
|
return next(iter(payload), None)
|
|
|
|
|
|
def format_empty_diffusion_outputs(
|
|
request: OmniDiffusionRequest,
|
|
*,
|
|
finished: bool = True,
|
|
) -> list[OmniRequestOutput]:
|
|
return [
|
|
OmniRequestOutput.from_diffusion(
|
|
request_id=request.request_id,
|
|
images=[],
|
|
prompt=request.prompt,
|
|
metrics={},
|
|
latents=None,
|
|
finished=finished,
|
|
)
|
|
]
|
|
|
|
|
|
def format_diffusion_outputs(
|
|
*,
|
|
request: OmniDiffusionRequest,
|
|
od_config: OmniDiffusionConfig,
|
|
diffusion_output: DiffusionOutput,
|
|
output_data: DiffusionPostprocessRawOutput,
|
|
postprocess_output: DiffusionPostprocessOutput,
|
|
) -> list[OmniRequestOutput]:
|
|
"""Convert a finished diffusion model output into API-facing outputs."""
|
|
|
|
primary_payload = _primary_payload(postprocess_output)
|
|
outputs = _ensure_list(primary_payload)
|
|
metrics = {
|
|
"image_num": int(request.sampling_params.num_outputs_per_prompt),
|
|
"resolution": (
|
|
int(request.sampling_params.resolution) if request.sampling_params.resolution is not None else None
|
|
),
|
|
"width": request.sampling_params.width,
|
|
"height": request.sampling_params.height,
|
|
"num_inference_steps": request.sampling_params.num_inference_steps,
|
|
}
|
|
|
|
# Only the primary payload determines the response type. Some image models
|
|
# include reasoning text in metadata, but their final output is still an image.
|
|
is_text_output = postprocess_output.primary_key == "text"
|
|
|
|
is_audio_output = supports_audio_output(od_config.model_class_name)
|
|
audio_sample_rate = _metadata_audio_sample_rate(postprocess_output.metadata)
|
|
if is_audio_output and audio_sample_rate is None:
|
|
model_cls = DiffusionModelRegistry._try_load_model_cls(od_config.model_class_name)
|
|
audio_sample_rate = getattr(model_cls, "audio_sample_rate", None)
|
|
|
|
stream_metadata = _stream_metadata_from_diffusion_output(diffusion_output)
|
|
if stream_metadata:
|
|
existing_stream = postprocess_output.metadata.get("stream")
|
|
merged_stream = dict(existing_stream) if isinstance(existing_stream, Mapping) else {}
|
|
merged_stream.update(stream_metadata)
|
|
postprocess_output.metadata["stream"] = merged_stream
|
|
|
|
return _format_single_prompt_output(
|
|
request=request,
|
|
prompt=request.prompt,
|
|
diffusion_output=diffusion_output,
|
|
outputs=outputs,
|
|
metrics=metrics,
|
|
postprocess_output=postprocess_output,
|
|
is_text_output=is_text_output,
|
|
is_audio_output=is_audio_output,
|
|
audio_sample_rate=audio_sample_rate,
|
|
finished=diffusion_output.finished,
|
|
)
|
|
|
|
|
|
def _stream_metadata_from_diffusion_output(diffusion_output: DiffusionOutput) -> DiffusionMetadata | None:
|
|
has_stream_data = (
|
|
diffusion_output.total_chunks > 1
|
|
or not diffusion_output.finished
|
|
or bool(diffusion_output.started_event_ids)
|
|
or bool(diffusion_output.active_event_ids)
|
|
or bool(diffusion_output.completed_event_ids)
|
|
)
|
|
if not has_stream_data:
|
|
return None
|
|
return cast(
|
|
DiffusionMetadata,
|
|
{
|
|
"generation_chunk_index": diffusion_output.chunk_index,
|
|
"started_event_ids": diffusion_output.started_event_ids,
|
|
"active_event_ids": diffusion_output.active_event_ids,
|
|
"completed_event_ids": diffusion_output.completed_event_ids,
|
|
},
|
|
)
|
|
|
|
|
|
def _ensure_list(outputs: DiffusionPayloadValue) -> list[DiffusionPayloadValue]:
|
|
if isinstance(outputs, list):
|
|
return outputs
|
|
return [outputs] if outputs is not None else []
|
|
|
|
|
|
def _primary_payload(postprocess_output: DiffusionPostprocessOutput) -> DiffusionPayloadValue | None:
|
|
if postprocess_output.primary_key is None:
|
|
return []
|
|
return postprocess_output.outputs.get(postprocess_output.primary_key)
|
|
|
|
|
|
def _metadata_audio_sample_rate(metadata: DiffusionMetadata) -> int | None:
|
|
audio_metadata = metadata.get("audio")
|
|
if not isinstance(audio_metadata, dict):
|
|
return None
|
|
sample_rate = audio_metadata.get("sample_rate")
|
|
return sample_rate if isinstance(sample_rate, int) else None
|
|
|
|
|
|
def _metadata_video_fps(metadata: DiffusionMetadata) -> float | None:
|
|
video_metadata = metadata.get("video")
|
|
if not isinstance(video_metadata, dict):
|
|
return None
|
|
fps = video_metadata.get("fps")
|
|
return fps if isinstance(fps, (int, float)) and not isinstance(fps, bool) else None
|
|
|
|
|
|
def _format_audio_multimodal_output(
|
|
payload: DiffusionPayloadValue,
|
|
audio_sample_rate: int | None,
|
|
metadata: DiffusionMetadata | None = None,
|
|
) -> DiffusionMultimodalOutput:
|
|
mm_output: DiffusionMultimodalOutput = {"audio": payload}
|
|
if metadata:
|
|
mm_output["metadata"] = metadata
|
|
if audio_sample_rate is not None:
|
|
mm_output["audio_sample_rate"] = audio_sample_rate
|
|
return mm_output
|
|
|
|
|
|
def _has_non_audio_postprocess_payload(postprocess_output: DiffusionPostprocessOutput) -> bool:
|
|
return (
|
|
"video" in postprocess_output.outputs
|
|
or "image" in postprocess_output.outputs
|
|
or "text" in postprocess_output.metadata
|
|
or "actions" in postprocess_output.outputs
|
|
or "trajectory" in postprocess_output.outputs
|
|
or _metadata_video_fps(postprocess_output.metadata) is not None
|
|
)
|
|
|
|
|
|
def _build_multimodal_output(
|
|
postprocess_output: DiffusionPostprocessOutput,
|
|
audio_sample_rate: int | None,
|
|
) -> DiffusionMultimodalOutput:
|
|
mm_output: DiffusionMultimodalOutput = {}
|
|
if postprocess_output.metadata:
|
|
mm_output["metadata"] = postprocess_output.metadata
|
|
for key, value in postprocess_output.outputs.items():
|
|
if key in {"audio", "actions", "trajectory"}:
|
|
mm_output[key] = value
|
|
if audio_sample_rate is not None:
|
|
mm_output["audio_sample_rate"] = audio_sample_rate
|
|
if (fps := _metadata_video_fps(postprocess_output.metadata)) is not None:
|
|
mm_output["fps"] = fps
|
|
return mm_output
|
|
|
|
|
|
def _format_single_prompt_output(
|
|
*,
|
|
request: OmniDiffusionRequest,
|
|
prompt: OmniPromptType,
|
|
diffusion_output: DiffusionOutput,
|
|
outputs: list[DiffusionPayloadValue],
|
|
metrics: dict[str, object],
|
|
postprocess_output: DiffusionPostprocessOutput,
|
|
is_text_output: bool,
|
|
is_audio_output: bool,
|
|
audio_sample_rate: int | None,
|
|
finished: bool = True,
|
|
) -> list[OmniRequestOutput]:
|
|
request_id = request.request_id
|
|
mm_output = _build_multimodal_output(postprocess_output, audio_sample_rate)
|
|
if is_text_output:
|
|
mm_output["text"] = outputs[0] if len(outputs) == 1 else outputs
|
|
trajectory_payload = _trajectory_payload(postprocess_output, diffusion_output)
|
|
trajectory_latents = trajectory_payload.get("latents")
|
|
trajectory_timesteps = trajectory_payload.get("timesteps")
|
|
trajectory_log_probs = trajectory_payload.get("log_probs")
|
|
trajectory_decoded = trajectory_payload.get("decoded")
|
|
|
|
if is_text_output:
|
|
return [
|
|
OmniRequestOutput.from_diffusion(
|
|
request_id=request_id,
|
|
images=[],
|
|
prompt=prompt,
|
|
metrics=metrics,
|
|
multimodal_output=mm_output,
|
|
final_output_type="text",
|
|
stage_durations=diffusion_output.stage_durations,
|
|
peak_memory_mb=diffusion_output.peak_memory_mb,
|
|
finished=finished,
|
|
),
|
|
]
|
|
|
|
if is_audio_output and not _has_non_audio_postprocess_payload(postprocess_output):
|
|
request_audio_payload = postprocess_output.outputs.get("audio")
|
|
if request_audio_payload is None:
|
|
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=trajectory_latents,
|
|
trajectory_latents=trajectory_latents,
|
|
trajectory_timesteps=trajectory_timesteps,
|
|
trajectory_log_probs=trajectory_log_probs,
|
|
trajectory_decoded=trajectory_decoded,
|
|
multimodal_output=_format_audio_multimodal_output(
|
|
request_audio_payload,
|
|
audio_sample_rate,
|
|
postprocess_output.metadata,
|
|
),
|
|
final_output_type="audio",
|
|
stage_durations=diffusion_output.stage_durations,
|
|
peak_memory_mb=diffusion_output.peak_memory_mb,
|
|
finished=finished,
|
|
),
|
|
]
|
|
|
|
return [
|
|
OmniRequestOutput.from_diffusion(
|
|
request_id=request_id,
|
|
images=outputs,
|
|
prompt=prompt,
|
|
metrics=metrics,
|
|
latents=trajectory_latents,
|
|
trajectory_latents=trajectory_latents,
|
|
trajectory_timesteps=trajectory_timesteps,
|
|
trajectory_log_probs=trajectory_log_probs,
|
|
trajectory_decoded=trajectory_decoded,
|
|
multimodal_output=mm_output,
|
|
stage_durations=diffusion_output.stage_durations,
|
|
peak_memory_mb=diffusion_output.peak_memory_mb,
|
|
finished=finished,
|
|
),
|
|
]
|
|
|
|
|
|
def _trajectory_payload(
|
|
postprocess_output: DiffusionPostprocessOutput,
|
|
diffusion_output: DiffusionOutput,
|
|
) -> DiffusionTrajectoryPayload:
|
|
trajectory: DiffusionTrajectoryPayload = {}
|
|
payload = postprocess_output.outputs.get("trajectory")
|
|
if isinstance(payload, Mapping):
|
|
for source_key, target_key in (
|
|
("latents", "latents"),
|
|
("timesteps", "timesteps"),
|
|
("log_probs", "log_probs"),
|
|
("decoded", "decoded"),
|
|
):
|
|
value = payload.get(source_key)
|
|
if value is not None:
|
|
trajectory[target_key] = value
|
|
|
|
fallback_fields = (
|
|
("latents", diffusion_output.trajectory_latents),
|
|
("timesteps", diffusion_output.trajectory_timesteps),
|
|
("log_probs", diffusion_output.trajectory_log_probs),
|
|
("decoded", diffusion_output.trajectory_decoded),
|
|
)
|
|
for key, value in fallback_fields:
|
|
if key not in trajectory and value is not None:
|
|
trajectory[key] = value
|
|
return trajectory
|