vllm-omni/vllm_omni/experimental/ar_diffusion/consumer.py

# SPDX-License-Identifier: Apache-2.0
"""AsyncOmni consumer for typed realtime AR-Diffusion ticks."""

from __future__ import annotations

import copy
from collections.abc import AsyncIterator, Callable, Mapping, Sequence
from typing import Any, Protocol

from vllm_omni.experimental.ar_diffusion.tick_protocol import (
    ARDiffusionChunkMetadata,
    ARDiffusionTickRequest,
)
from vllm_omni.inputs.data import (
    OmniDiffusionSamplingParams,
    OmniPromptType,
    OmniSamplingParams,
)
from vllm_omni.outputs import OmniRequestOutput

AR_DIFFUSION_OUTPUT_METADATA_KEY = "ar_diffusion"


class ARDiffusionGenerateClient(Protocol):
    """Internal subset of AsyncOmni used to execute one tick."""

    def generate(
        self,
        prompt: OmniPromptType,
        *,
        request_id: str,
        sampling_params_list: Sequence[OmniSamplingParams],
        output_modalities: list[str] | None = None,
    ) -> AsyncIterator[OmniRequestOutput]: ...


def _metadata_mapping(output: OmniRequestOutput) -> Mapping[str, Any]:
    multimodal_output = output.multimodal_output
    if not isinstance(multimodal_output, Mapping):
        raise ValueError("AR-Diffusion output must contain a multimodal output mapping.")
    metadata = multimodal_output.get("metadata")
    if not isinstance(metadata, Mapping):
        raise ValueError("AR-Diffusion output must contain multimodal_output.metadata.")
    value = metadata.get(AR_DIFFUSION_OUTPUT_METADATA_KEY)
    if not isinstance(value, Mapping):
        raise ValueError(f"AR-Diffusion output metadata must contain {AR_DIFFUSION_OUTPUT_METADATA_KEY!r}.")
    return value


def _parse_chunk_metadata(output: OmniRequestOutput) -> ARDiffusionChunkMetadata:
    value = _metadata_mapping(output)
    session_id = value.get("session_id")
    request_id = value.get("request_id")
    chunk_index = value.get("chunk_index")
    applied_event_ids = value.get("applied_event_ids", ())
    if not isinstance(session_id, str) or not session_id.strip():
        raise ValueError("AR-Diffusion output metadata session_id must be a non-empty string.")
    if not isinstance(request_id, str) or not request_id.strip():
        raise ValueError("AR-Diffusion output metadata request_id must be a non-empty string.")
    if isinstance(chunk_index, bool) or not isinstance(chunk_index, int) or chunk_index < 0:
        raise ValueError("AR-Diffusion output metadata chunk_index must be a non-negative integer.")
    if not isinstance(applied_event_ids, Sequence) or isinstance(applied_event_ids, (str, bytes, bytearray)):
        raise ValueError("AR-Diffusion output metadata applied_event_ids must be a sequence.")
    event_ids = tuple(applied_event_ids)
    if any(isinstance(event_id, bool) or not isinstance(event_id, int) or event_id < 0 for event_id in event_ids):
        raise ValueError("AR-Diffusion output metadata applied_event_ids must contain non-negative integers.")
    if event_ids != tuple(sorted(set(event_ids))):
        raise ValueError("AR-Diffusion output metadata applied_event_ids must be unique and strictly increasing.")
    return ARDiffusionChunkMetadata(
        session_id=session_id,
        request_id=request_id,
        chunk_index=chunk_index,
        applied_event_ids=event_ids,
    )


class ARDiffusionOmniTickConsumer:
    """Converts a typed tick into one AsyncOmni diffusion request.

    ``prompt_provider`` packages session initialization data such as an initial
    image into the normal Omni prompt. Control ``track/schema/data`` remains in
    the typed tick and is not interpreted by this class.
    """

    def __init__(
        self,
        engine_client: ARDiffusionGenerateClient,
        *,
        prompt_provider: Callable[[ARDiffusionTickRequest], OmniPromptType],
        sampling_params_list: Sequence[OmniSamplingParams],
        diffusion_stage_id: int,
        output_modalities: Sequence[str] | None = None,
    ) -> None:
        templates = tuple(sampling_params_list)
        if not templates:
            raise ValueError("sampling_params_list must not be empty.")
        if (
            isinstance(diffusion_stage_id, bool)
            or not isinstance(diffusion_stage_id, int)
            or not 0 <= diffusion_stage_id < len(templates)
        ):
            raise ValueError("diffusion_stage_id must index sampling_params_list.")
        if not isinstance(templates[diffusion_stage_id], OmniDiffusionSamplingParams):
            raise TypeError("The selected diffusion stage must use OmniDiffusionSamplingParams.")
        stage_pools = getattr(getattr(engine_client, "engine", None), "stage_pools", None)
        if stage_pools is None or not 0 <= diffusion_stage_id < len(stage_pools):
            raise ValueError(
                "AR-Diffusion realtime execution requires an initialized AsyncOmni "
                "engine with a matching diffusion stage pool."
            )
        replica_count = getattr(stage_pools[diffusion_stage_id], "num_replicas", None)
        if isinstance(replica_count, bool) or not isinstance(replica_count, int) or replica_count < 1:
            raise ValueError("The selected AR-Diffusion stage must expose a positive num_replicas.")
        if replica_count != 1:
            raise ValueError(
                "AR-Diffusion realtime sessions currently require exactly one "
                "replica for the selected diffusion stage; session-affine "
                f"routing is not implemented (got num_replicas={replica_count})."
            )
        self._engine_client = engine_client
        self._prompt_provider = prompt_provider
        self._sampling_params_templates = templates
        self._diffusion_stage_id = diffusion_stage_id
        self._output_modalities = list(output_modalities) if output_modalities is not None else None

    def _sampling_params_for_tick(
        self,
        tick: ARDiffusionTickRequest,
    ) -> list[OmniSamplingParams]:
        params = [copy.copy(template) for template in self._sampling_params_templates]
        diffusion_params = params[self._diffusion_stage_id]
        assert isinstance(diffusion_params, OmniDiffusionSamplingParams)
        diffusion_params.extra_args = {
            **(diffusion_params.extra_args or {}),
            **tick.to_extra_args(),
        }
        return params

    async def execute_tick(self, tick: ARDiffusionTickRequest) -> OmniRequestOutput:
        prompt = self._prompt_provider(tick)
        final_output: OmniRequestOutput | None = None
        async for output in self._engine_client.generate(
            prompt,
            request_id=tick.request_id,
            sampling_params_list=self._sampling_params_for_tick(tick),
            output_modalities=self._output_modalities,
        ):
            if output.stage_id != self._diffusion_stage_id:
                continue
            if output.finished:
                final_output = output

        if final_output is None:
            raise RuntimeError(
                f"AR-Diffusion tick {tick.request_id!r} completed without a final "
                f"output from stage {self._diffusion_stage_id}."
            )
        if final_output.error:
            raise RuntimeError(f"AR-Diffusion tick {tick.request_id!r} failed: {final_output.error}")
        if final_output.request_id != tick.request_id:
            raise ValueError("AR-Diffusion output request_id does not match the submitted tick request_id.")
        # Parse here so malformed/missing standard metadata fails at the engine
        # boundary, before the session attempts to commit accepted events.
        self.chunk_metadata(final_output)
        return final_output

    def chunk_metadata(self, output: OmniRequestOutput) -> ARDiffusionChunkMetadata:
        return _parse_chunk_metadata(output)