import asyncio
import base64
import json
import time
import uuid
from collections.abc import AsyncGenerator, AsyncIterator, Callable
from dataclasses import fields, is_dataclass
from datetime import datetime, timedelta, timezone
from io import BytesIO
from typing import Any, Final, cast

import jinja2
import torch
from fastapi import Request
from openai.types.chat.chat_completion_audio import ChatCompletionAudio as OpenAIChatCompletionAudio
from PIL import Image
from pydantic import TypeAdapter
from vllm.entrypoints.chat_utils import (
    ChatCompletionMessageParam,
    ChatTemplateContentFormatOption,
    ConversationMessage,
    get_history_tool_calls_cnt,
    make_tool_call_id,
)

from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
from vllm_omni.entrypoints.async_omni import AsyncOmni
from vllm_omni.entrypoints.openai.diffusion_request_utils import (
    apply_normalized_diffusion_request_extra_args,
    normalize_diffusion_request_args,
)
from vllm_omni.entrypoints.openai.protocol.chat_completion import OmniChatCompletionResponse
from vllm_omni.entrypoints.utils import coerce_param_message_types
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
from vllm_omni.metrics import definitions as _metric_defs
from vllm_omni.metrics.modality import (
    observe_audio_first_packet,
    observe_audio_streaming_finalize,
)
from vllm_omni.model_extras import (
    get_extra_body_params,
    get_extra_output_params,
)

try:
    import soundfile
except ImportError:
    soundfile = None


from vllm.entrypoints.generate.base.serving import clamp_prompt_logprobs
from vllm.entrypoints.launcher import terminate_if_errored
from vllm.entrypoints.openai.chat_completion.protocol import (
    ChatCompletionNamedToolChoiceParam,
    ChatCompletionRequest,
    ChatCompletionResponse,
    ChatCompletionResponseChoice,
    ChatCompletionResponseStreamChoice,
    ChatMessage,
)
from vllm.entrypoints.openai.chat_completion.serving import (
    OpenAIServingChat,
    _get_mm_token_counts,
    _make_prompt_tokens_details,
)
from vllm.entrypoints.openai.engine.protocol import (
    DeltaFunctionCall,
    DeltaMessage,
    DeltaToolCall,
    ErrorInfo,
    ErrorResponse,
    FunctionCall,
    FunctionDefinition,
    RequestResponseMetadata,
    ToolCall,
    UsageInfo,
)
from vllm.entrypoints.openai.parser.harmony_utils import (
    get_streamable_parser_for_assistant,
)
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
from vllm.entrypoints.serve.engine.typing import ChatLikeRequest
from vllm.entrypoints.serve.utils.api_utils import should_include_usage
from vllm.entrypoints.serve.utils.tool_calls_utils import maybe_filter_parallel_tool_calls
from vllm.inputs import PromptType
from vllm.logger import init_logger
from vllm.multimodal.media.connector import MediaConnector
from vllm.outputs import RequestOutput

try:
    from vllm.parser.mistral import MistralToolCall
except ImportError:
    # Accelerator images can lag the core vLLM release by one version.
    from vllm.tool_parsers.mistral_tool_parser import MistralToolCall


from vllm.reasoning import ReasoningParser
from vllm.renderers import BaseRenderer, merge_kwargs
from vllm.renderers.inputs import TokPrompt
from vllm.sampling_params import RequestOutputKind, SamplingParams
from vllm.tokenizers import TokenizerLike
from vllm.tokenizers import TokenizerLike as AnyTokenizer
from vllm.tokenizers.mistral import (
    MistralTokenizer,
    maybe_serialize_tool_calls,
    truncate_tool_call_ids,
    validate_request_params,
)
from vllm.tool_parsers import ToolParser
from vllm.tool_parsers.streaming import extract_required_tool_call_streaming
from vllm.utils.collection_utils import as_list
from vllm.v1.engine.exceptions import EngineDeadError

from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
from vllm_omni.entrypoints.openai.image_api_utils import (
    encode_image_base64_with_compression,
    validate_layered_layers,
)
from vllm_omni.entrypoints.openai.protocol import OmniChatCompletionStreamResponse
from vllm_omni.entrypoints.openai.protocol.audio import (
    DEFAULT_AUDIO_FORMAT,
    SUPPORTED_CHAT_AUDIO_FORMATS,
    AudioResponse,
    CreateAudio,
)
from vllm_omni.entrypoints.openai.protocol.images import (
    ImageData,
    ImageEditARDeltaChunk,
    ImageEditImageChunk,
    ImageEditStreamError,
)
from vllm_omni.entrypoints.openai.stage_params import (
    build_stage_sampling_params_list,
    clone_sampling_params,
    get_default_sampling_params_list,
)
from vllm_omni.entrypoints.openai.utils import (
    get_stage_type,
    get_supported_speakers_from_hf_config,
    is_single_stage_diffusion,
    parse_lora_request,
    resolve_diffusion_od_config,
    validate_requested_speaker,
)
from vllm_omni.errors import OmniClientError
from vllm_omni.lora.request import LoRARequest
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.outputs.output_metadata import DiffusionMetadataMapping, DiffusionMetadataValue
from vllm_omni.utils.audio import audio_chunk_pcm_bytes, audio_chunk_sample_rate

logger = init_logger(__name__)


async def _identity_async(value: Any) -> Any:
    return value


class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
    """OpenAI-compatible chat serving for both LLM and Diffusion models.

    This class extends OpenAIServingChat to support:
    - Standard LLM chat completions
    - Diffusion model image generation via chat interface

    For diffusion mode, use the `for_diffusion` class method to create an instance.
    """

    # Diffusion mode attributes
    _diffusion_mode: bool = False
    _diffusion_engine: AsyncOmni | None = None
    _diffusion_model_name: str = ""
    _supported_speakers: set[str] | None = None
    _diffusion_extra_body_params: frozenset[str] | None = None
    _diffusion_extra_output_params: frozenset[str] | None = None
    # Common fields stay at the request root, but still participate in
    # conflicts with model-specific compatibility forms.
    _diffusion_common_root_fields = frozenset(
        {
            "height",
            "width",
            "num_outputs_per_prompt",
            "seed",
            "quality",
            "num_inference_steps",
            "guidance_scale",
            "true_cfg_scale",
            "num_frames",
            "guidance_scale_2",
            "layers",
            "resolution",
        }
    )
    _diffusion_existing_control_fields = frozenset(
        {
            "size",
            "negative_prompt",
            "lora",
            "modalities",
        }
    )
    _diffusion_root_field_aliases = {"cfg_scale": "true_cfg_scale"}
    _diffusion_serving_root_fields = _diffusion_common_root_fields | _diffusion_existing_control_fields

    # Harmony flag (always False for vllm-omni models)
    use_harmony: bool = False

    @property
    def tool_call_id_type(self) -> str:
        """Return the tool call ID type, delegating to model config.

        Upstream vLLM removed the stored ``tool_call_id_type`` attribute
        from ``OpenAIServingChat`` after the ParserManager refactor; the
        field is now resolved on demand via ``get_tool_call_id_type``.
        """
        try:
            from vllm.entrypoints.chat_utils import get_tool_call_id_type

            return get_tool_call_id_type(self.model_config)
        except Exception:
            return "random"

    def _should_stream_with_auto_tool_parsing(self, request: ChatCompletionRequest) -> bool:
        """Check if streamed tokens should go through the tool-call parser.

        We only want to do this IF user-provided tools are set, a tool parser
        is configured, "auto" tool choice is enabled, and the request's tool
        choice field indicates that "auto" tool choice should be used.

        This method existed in upstream vLLM commit 91df0fad4 (OpenAIServingChat)
        but was removed in the Harmony refactoring (PR #45171, #45104).  Omni's
        independently-maintained ``chat_completion_stream_generator`` still calls
        it, so we keep a local copy.
        """
        # parser_cls may be None (no tool parser configured); guard accordingly
        tp_cls = self.parser_cls.tool_parser_cls if self.parser_cls is not None else None
        return request.tools and tp_cls is not None and self.enable_auto_tools and request.tool_choice in ["auto", None]

    @staticmethod
    def _stage_get(obj: Any, key: str, default: Any = None) -> Any:
        if obj is None:
            return default
        if hasattr(obj, "get"):
            return obj.get(key, default)
        return getattr(obj, key, default)

    def _has_minicpmo45_stage(self) -> bool:
        for stage in getattr(self.engine_client, "stage_configs", []) or []:
            engine_args = self._stage_get(stage, "engine_args")
            model_arch = self._stage_get(engine_args, "model_arch")
            if model_arch == "MiniCPMO45OmniForConditionalGeneration":
                return True
        return False

    def _fix_minicpmo45_audio_stream_output_kinds(
        self,
        sampling_params_list: list[Any],
        output_modalities: list[str] | tuple[str, ...] | None,
    ) -> list[Any]:
        """Keep MiniCPM-o 4.5's thinker final-only during audio streaming.

        Chat ``stream=true`` normally coerces every AR stage to DELTA output.
        That is correct for token-level text streaming, but MiniCPM-o 4.5's
        Talker consumes the completed TTS span plus aligned thinker hidden
        states. If Stage0 is DELTA, llm2tts receives one generated token at a
        time and repeatedly starts tiny, independent TTS streams. The API then
        never represents one coherent audio response.

        For this model, stream the audio stage, not the thinker-to-talker
        boundary.
        """
        if not output_modalities or "audio" not in output_modalities:
            return sampling_params_list
        if not self._has_minicpmo45_stage():
            return sampling_params_list

        stage_configs = getattr(self.engine_client, "stage_configs", []) or []
        for idx, stage in enumerate(stage_configs):
            if idx >= len(sampling_params_list):
                break
            sp = sampling_params_list[idx]
            if not hasattr(sp, "output_kind"):
                continue

            engine_args = self._stage_get(stage, "engine_args")
            if self._stage_get(engine_args, "model_arch") != "MiniCPMO45OmniForConditionalGeneration":
                continue

            model_stage = self._stage_get(engine_args, "model_stage")
            if model_stage == "llm":
                sp.output_kind = RequestOutputKind.FINAL_ONLY
            elif model_stage == "tts":
                sp.output_kind = RequestOutputKind.DELTA
        return sampling_params_list

    @classmethod
    def for_diffusion(
        cls,
        diffusion_engine: AsyncOmni,
        model_name: str,
    ) -> "OmniOpenAIServingChat":
        """Create a chat serving instance for diffusion models.

        Args:
            diffusion_engine: The async diffusion engine
            model_name: Name of the model being served

        Returns:
            OmniOpenAIServingChat instance configured for diffusion mode

        Note:
            Request-level parameters (num_inference_steps, guidance_scale, seed,
            height, width, num_frames, fps, etc.) are passed per-request via the API.
        """
        instance = cls.__new__(cls)
        instance._diffusion_mode = True
        instance._diffusion_engine = diffusion_engine
        instance._diffusion_model_name = model_name
        instance._diffusion_extra_body_params = None
        instance._diffusion_extra_output_params = None
        instance.engine_client = None
        instance.has_kv_connector = False
        # Extra body/output params are resolved lazily on first use; see
        # _get_diffusion_extra_body_params / _get_diffusion_extra_output_params.
        return instance

    def _get_diffusion_extra_body_params(self) -> frozenset[str]:
        """Return model-specific extra_body params from the extra registry."""
        if self._diffusion_extra_body_params is not None:
            return self._diffusion_extra_body_params

        params: frozenset[str] = frozenset()
        try:
            od_config = resolve_diffusion_od_config(self.engine_client, self._diffusion_engine)
            if od_config is not None and getattr(od_config, "model_class_name", None):
                params = get_extra_body_params(od_config.model_class_name)
        except Exception as e:
            logger.warning("Failed to read model extra_body params: %s", e)

        self._diffusion_extra_body_params = params
        return params

    def _normalize_diffusion_request_args(
        self,
        request: ChatCompletionRequest,
    ) -> tuple[dict[str, object], dict[str, Any]]:
        """Normalize extras and preserve validated common root fields."""
        root = dict(request.model_extra or {})
        nested = root.pop("extra_body", None)
        if nested is None:
            nested = getattr(request, "extra_body", None)
        explicit = getattr(request, "model_fields_set", None)
        if explicit is None:
            explicit = getattr(request, "__fields_set__", ())
        explicit_root_args = {
            key: getattr(request, key) for key in explicit if key != "extra_body" and hasattr(request, key)
        }
        return normalize_diffusion_request_args(
            root=root,
            nested=nested,
            explicit_root_args=explicit_root_args,
            serving_root_fields=self._diffusion_serving_root_fields,
            registered_extra_fields=self._get_diffusion_extra_body_params(),
            root_field_aliases=self._diffusion_root_field_aliases,
        )

    def _get_diffusion_extra_output_params(
        self,
        output: object,
    ) -> dict[str, DiffusionMetadataValue] | None:
        """Pick model-specific extra output keys from diffusion metadata."""
        metadata: DiffusionMetadataMapping = {}
        mm_output = getattr(output, "multimodal_output", None)
        if isinstance(mm_output, dict):
            raw_metadata = mm_output.get("metadata")
            if isinstance(raw_metadata, dict):
                metadata = raw_metadata

        if self._diffusion_extra_output_params is None:
            params: frozenset[str] = frozenset()
            try:
                od_config = resolve_diffusion_od_config(self.engine_client, self._diffusion_engine)
                if od_config is not None and getattr(od_config, "model_class_name", None):
                    params = get_extra_output_params(od_config.model_class_name)
            except Exception as e:
                logger.warning("Failed to read model extra output params: %s", e)
            self._diffusion_extra_output_params = params

        if not self._diffusion_extra_output_params:
            return None
        flat_metadata: dict[str, DiffusionMetadataValue] = {}
        for section in metadata.values():
            if isinstance(section, dict):
                flat_metadata.update(section)
        out = {k: flat_metadata[k] for k in self._diffusion_extra_output_params if k in flat_metadata}
        return out or None

    @staticmethod
    def _get_diffusion_text_output(output: object) -> str:
        mm_output = getattr(output, "multimodal_output", None)
        if isinstance(mm_output, dict) and mm_output.get("text") is not None:
            return str(mm_output["text"])
        return ""

    def _get_supported_speakers(self) -> set[str]:
        """Load supported speakers from model config (cached)."""
        if self._supported_speakers is not None:
            return self._supported_speakers
        try:
            self._supported_speakers = get_supported_speakers_from_hf_config(self.model_config.hf_config)
            return self._supported_speakers
        except Exception as e:
            logger.warning("Could not load speakers from model config: %s", e)
        self._supported_speakers = set()
        return self._supported_speakers

    @staticmethod
    def _truthy_extra_body_flag(request: Any, key: str) -> bool:
        if isinstance(request, dict):
            extra_body = request
            model_extra = {}
        else:
            extra_body = getattr(request, "extra_body", None) or {}
            model_extra = getattr(request, "model_extra", None) or {}
        value = extra_body.get(key, model_extra.get(key))
        if isinstance(value, str):
            return value.lower() in {"1", "true", "yes", "on"}
        return bool(value)

    @classmethod
    def _filter_stage_metrics_detail(cls, metrics: dict[str, Any] | None, request: Any) -> dict[str, Any] | None:
        if not metrics:
            return metrics
        if cls._truthy_extra_body_flag(request, "return_stage_metrics"):
            return metrics
        return None

    async def create_chat_completion(
        self,
        request: ChatCompletionRequest,
        raw_request: Request | None = None,
    ) -> AsyncGenerator[str, None] | ChatCompletionResponse | ErrorResponse:
        """
        Chat Completion API similar to OpenAI's API.

        See https://platform.openai.com/docs/api-reference/chat/create
        for the API specification. This API mimics the OpenAI
        Chat Completion API.

        For diffusion models, this generates images and returns them
        in a chat completion response format.
        """
        return await self._with_kv_transfer_rejection_cleanup(
            self._create_chat_completion(request, raw_request), request, raw_request
        )

    async def _create_chat_completion(
        self,
        request: ChatCompletionRequest,
        raw_request: Request | None = None,
    ) -> AsyncGenerator[str, None] | ChatCompletionResponse | ErrorResponse:
        stage_configs = getattr(self.engine_client, "stage_configs", ()) or ()
        serves_diffusion = self._diffusion_mode or any(
            get_stage_type(stage_config) == "diffusion" for stage_config in stage_configs
        )
        normalized_extra_args: dict[str, object] = {}
        diffusion_request_args: dict[str, Any] = {}
        if serves_diffusion:
            try:
                normalized_extra_args, diffusion_request_args = self._normalize_diffusion_request_args(request)
            except ValueError as exc:
                return self._create_error_response(str(exc), status_code=400)

        # Handle diffusion mode
        if self._diffusion_mode:
            return await self._create_diffusion_chat_completion(
                request,
                normalized_extra_args,
                diffusion_request_args,
                raw_request,
            )

        request_timestamp = time.time()
        if raw_request is not None:
            request_timestamp = float(getattr(raw_request.state, "request_timestamp", request_timestamp))

        error_check_ret = await self._check_model(request)
        if error_check_ret is not None:
            logger.error("Error with model %s", error_check_ret)
            return error_check_ret

        # If the engine is dead, raise the engine's DEAD_ERROR.
        # This is required for the streaming case, where we return a
        # success status before we actually start generating text :).
        if self.engine_client.errored:
            raise self.engine_client.dead_error

        try:
            lora_request = self._maybe_get_adapters(request, supports_default_mm_loras=True)

            model_name = self.models.model_name(lora_request)

            renderer = self.renderer
            tokenizer = renderer.get_tokenizer()
            if tokenizer is None:
                tokenizer = await self.engine_client.get_tokenizer()

            reasoning_parser: ReasoningParser | None = None
            if self.parser_cls is not None and self.parser_cls.reasoning_parser_cls is not None:
                chat_template_kwargs = self._effective_chat_template_kwargs(request)
                reasoning_parser = self.parser_cls.reasoning_parser_cls(
                    tokenizer,
                    chat_template_kwargs=chat_template_kwargs,  # type: ignore[call-arg]
                )

            tool_parser = self.parser_cls.tool_parser_cls if self.parser_cls is not None else None

            if isinstance(tokenizer, MistralTokenizer):
                # because of issues with pydantic we need to potentially
                # re-serialize the tool_calls field of the request
                # for more info: see comment in `maybe_serialize_tool_calls`
                maybe_serialize_tool_calls(request)
                truncate_tool_call_ids(request)
                validate_request_params(request)

            # Check if tool parsing is unavailable (common condition)
            tool_parsing_unavailable = (
                tool_parser is None and not isinstance(tokenizer, MistralTokenizer) and not self.use_harmony
            )

            # Validate tool_choice when tool parsing is required but unavailable
            if tool_parsing_unavailable and request.tool_choice not in (
                None,
                "none",
            ):
                if request.tool_choice == "auto" and not self.enable_auto_tools:
                    # for hf tokenizers, "auto" tools requires
                    # --enable-auto-tool-choice and --tool-call-parser
                    return self.create_error_response(
                        '"auto" tool choice requires --enable-auto-tool-choice and --tool-call-parser to be set'
                    )
                elif request.tool_choice != "auto":
                    # "required" or named tool requires tool parser
                    return self.create_error_response(
                        f'tool_choice="{request.tool_choice}" requires --tool-call-parser to be set'
                    )

            if request.tools is None or (request.tool_choice == "none" and self.exclude_tools_when_tool_choice_none):
                tool_dicts = None
            else:
                tool_dicts = [tool.model_dump() for tool in request.tools]

            if not self.use_harmony:
                error_check_ret = self.online_renderer.validate_chat_template(
                    request_chat_template=request.chat_template,
                    chat_template_kwargs=request.chat_template_kwargs,
                    trust_request_chat_template=self.trust_request_chat_template,
                )
                if error_check_ret is not None:
                    return error_check_ret

                # Effective kwargs fold request.chat_template_kwargs, reasoning_effort,
                # and server defaults — mirrors OpenAIServingChat._effective_chat_template_kwargs.
                merged_template_kwargs = self._effective_chat_template_kwargs(request)
                conversation, engine_prompts = await self._preprocess_chat(
                    request,
                    request.messages,
                    default_template=request.chat_template or self.chat_template,
                    default_template_content_format=self.chat_template_content_format,
                    default_template_kwargs=merged_template_kwargs,
                    tool_dicts=tool_dicts,
                    tool_parser=tool_parser,
                    # OMNI: Additional parameters
                    renderer=renderer,
                    add_generation_prompt=request.add_generation_prompt,
                    continue_final_message=request.continue_final_message,
                    documents=getattr(request, "documents", None),
                    add_special_tokens=request.add_special_tokens,
                )
            else:
                should_include_tools = tool_dicts is not None
                conversation, engine_prompts = self.online_renderer._make_request_with_harmony(
                    request, should_include_tools
                )

        except (ValueError, TypeError, RuntimeError, jinja2.TemplateError) as e:
            logger.exception("Error in preprocessing prompt inputs")
            message = str(e)
            if e.__cause__ is not None:
                message = f"{message} {e.__cause__}"
            return self.create_error_response(message)

        request_id = f"chatcmpl-{self._base_request_id(raw_request, request.request_id)}"

        request_metadata = RequestResponseMetadata(request_id=request_id)
        if raw_request:
            raw_request.state.request_metadata = request_metadata

        # Some models will return a list like ["text", None, "audio"], better
        # to strip None in the list
        engine_output_modalities = [x for x in self.engine_client.output_modalities if x is not None]
        output_modalities = getattr(request, "modalities", engine_output_modalities)
        request.modalities = output_modalities if output_modalities is not None else engine_output_modalities

        if not isinstance(request.modalities, list) or not all(isinstance(m, str) for m in request.modalities):
            return self.create_error_response("'modalities' must be a list of strings.")
        allowed_modalities = set(engine_output_modalities)
        if is_single_stage_diffusion(self.engine_client):
            allowed_modalities.add("text")
        unsupported = set(request.modalities) - allowed_modalities
        if unsupported:
            return self.create_error_response(
                f"Unsupported output modalities {', '.join(sorted(unsupported))} for this model. "
                f"Supported modalities: {', '.join(sorted(allowed_modalities))}",
            )

        if request.modalities and "audio" in request.modalities:
            audio_format_check = self._resolve_audio_format(request)
            if isinstance(audio_format_check, ErrorResponse):
                return audio_format_check

        num_inference_steps = None
        extra_body = diffusion_request_args
        # Omni multistage image generation: Stage-0 (AR) should receive a clean
        # text prompt (and optional conditioning image/size) so the model's own
        # processor can construct the correct inputs.
        # If we pass pre-tokenized chat-template ids, GLM-Image can become
        # effectively unconditioned and produce nonsense images.
        if request.modalities and ("image" in request.modalities):
            try:
                extracted_prompt, reference_images = self._extract_diffusion_prompt_and_images_from_messages(
                    request.messages
                )
                if not extracted_prompt:
                    return self.create_error_response("No text prompt found in messages")

                height, width = self._resolve_height_width_from_extra_body(extra_body)

                num_inference_steps = extra_body.get("num_inference_steps")
                if num_inference_steps is not None:
                    try:
                        num_inference_steps = int(num_inference_steps)
                    except Exception:
                        num_inference_steps = None

                negative_prompt = extra_body.get("negative_prompt")
                engine_prompt_image: dict[str, Any] | None = None
                if reference_images:
                    # Best-effort decode first reference image for i2i.
                    try:
                        img_bytes = base64.b64decode(reference_images[0])
                        img = Image.open(BytesIO(img_bytes))
                        engine_prompt_image = {"img2img": img}
                    except Exception:
                        engine_prompt_image = None

                # Override the prompts produced by chat-template preprocessing.
                is_img2img = engine_prompt_image is not None
                tprompt: OmniTextPrompt = {"prompt": extracted_prompt}
                if is_img2img:
                    tprompt["modalities"] = ["img2img"]
                else:
                    tprompt["modalities"] = ["image"]
                if negative_prompt is not None:
                    tprompt["negative_prompt"] = negative_prompt
                # Always attach mm_processor_kwargs (possibly empty) so
                # OmniInputPreprocessor._process_text routes through the
                # multimodal processor path. Without it, the preprocessor
                # falls back to plain _tokenize_prompt and AR-based image-gen
                # models like GLM-Image never see their image-generation
                # scaffold.
                mm_processor_kwargs: dict[str, Any] = {}
                if height is not None:
                    mm_processor_kwargs["target_h"] = height
                if width is not None:
                    mm_processor_kwargs["target_w"] = width
                # Pass output modalities so model-specific MM processors can
                # detect image-generation requests and apply their own prompt
                # rewrites (e.g. query-token expansion, placeholder injection).
                mm_processor_kwargs["modalities"] = ["img2img"] if is_img2img else ["image"]
                tprompt["mm_processor_kwargs"] = mm_processor_kwargs
                if engine_prompt_image is not None:
                    tprompt["multi_modal_data"] = engine_prompt_image
                    # Provide multi_modal_uuids so that newer vLLM versions
                    # can validate multi_modal_data / multi_modal_uuids
                    # consistency.  After the multimodal processor consumes
                    # the image data, the uuids remain as a stable reference.
                    tprompt["multi_modal_uuids"] = {
                        k: [f"{request_id}-{k}-{i}" for i in range(len(v))]
                        if isinstance(v, list)
                        else [f"{request_id}-{k}-0"]
                        for k, v in engine_prompt_image.items()
                    }

                engine_prompts = [tprompt]
                # Store height/width for applying to diffusion stage sampling params later
                _image_gen_height = height
                _image_gen_width = width
            except Exception as e:
                logger.warning("Failed to build image-generation prompt for omni multistage: %s", e)
                _image_gen_height = None
                _image_gen_width = None
        elif request.modalities and ("text" in request.modalities) and is_single_stage_diffusion(self.engine_client):
            # Single-stage diffusion text output (img2text / text2text).
            # Build a diffusion-style prompt with modalities=["text"] so the
            # pipeline routes to its text generation path.
            try:
                extracted_prompt, reference_images = self._extract_diffusion_prompt_and_images_from_messages(
                    request.messages
                )
                if not extracted_prompt:
                    return self.create_error_response("No text prompt found in messages")

                tprompt: OmniTextPrompt = {"prompt": extracted_prompt}
                tprompt["modalities"] = ["text"]

                if reference_images:
                    try:
                        img_bytes = base64.b64decode(reference_images[0])
                        img = Image.open(BytesIO(img_bytes))
                        tprompt["multi_modal_data"] = {"image": img}
                        tprompt["multi_modal_uuids"] = {"image": [f"{request_id}-image-0"]}
                    except Exception:
                        pass

                engine_prompts = [tprompt]
            except Exception as e:
                logger.warning("Failed to build text-output prompt for single-stage diffusion: %s", e)
            _image_gen_height = None
            _image_gen_width = None
        else:
            _image_gen_height = None
            _image_gen_width = None

        # Schedule the request and get the result generator.
        generators: list[AsyncGenerator[RequestOutput, None]] = []
        mm_token_counts: dict[str, int] | None = None
        try:
            for i, engine_prompt in enumerate(engine_prompts):
                if self.enable_prompt_tokens_details:
                    mm_token_counts = _get_mm_token_counts(engine_prompt)
                if hasattr(request, "sampling_params_list"):
                    sampling_params_list = self._to_sampling_params_list(request.sampling_params_list)
                else:
                    # Use standard OpenAI API parameters for comprehension stage
                    sampling_params_list = self._build_sampling_params_list_from_request(request)

                # If this is a streaming (output) request, coerce cumulative outputs
                # to delta to ensure emitted outputs are correctly drained. Otherwise
                # convert cumulative to Final Only to ensure the output is correct.
                sampling_params_list = coerce_param_message_types(sampling_params_list, request.stream)
                if request.stream:
                    sampling_params_list = self._fix_minicpmo45_audio_stream_output_kinds(
                        sampling_params_list,
                        output_modalities,
                    )

                comprehension_idx = (
                    self._get_comprehension_stage_index()
                    if _image_gen_height is not None and _image_gen_width is not None
                    else None
                )
                # Apply user-specified overrides to diffusion stage(s) for image generation
                for idx, sp in enumerate(sampling_params_list):
                    if idx == comprehension_idx:
                        extra_args = dict(getattr(sp, "extra_args", {}) or {})
                        extra_args["target_h"] = int(_image_gen_height)
                        extra_args["target_w"] = int(_image_gen_width)
                        sp.extra_args = extra_args
                    if hasattr(sp, "height") and _image_gen_height is not None:
                        sp.height = _image_gen_height
                    if hasattr(sp, "width") and _image_gen_width is not None:
                        sp.width = _image_gen_width
                    if hasattr(sp, "num_inference_steps") and num_inference_steps is not None:
                        sp.num_inference_steps = num_inference_steps
                    if isinstance(sp, OmniDiffusionSamplingParams):
                        quality = extra_body.get("quality")
                        self._set_if_supported(
                            sp,
                            quality=quality,
                        )
                        apply_normalized_diffusion_request_extra_args(sp, normalized_extra_args)
                    else:
                        apply_declared_extra_args(
                            sp,
                            self._get_diffusion_extra_body_params(),
                            extra_body,
                        )

                self._log_inputs(
                    request_id,
                    engine_prompt,
                    params_list=sampling_params_list,
                    lora_request=lora_request,
                )

                generator = self.engine_client.generate(
                    prompt=engine_prompt,
                    request_id=request_id,
                    sampling_params_list=sampling_params_list,
                    output_modalities=output_modalities,
                    arrival_time=request_timestamp,
                    lora_request=lora_request,
                )

                generators.append(generator)
        except ValueError as e:
            return self.create_error_response(e)

        assert len(generators) == 1
        (result_generator,) = generators

        # Streaming response
        if request.stream:
            return self.chat_completion_stream_generator(
                request,
                result_generator,
                request_id,
                model_name,
                conversation,
                tokenizer,
                request_metadata,
                reasoning_parser,
                mm_token_counts=mm_token_counts,
                raw_request=raw_request,
            )

        try:
            return await self.chat_completion_full_generator(
                request,
                result_generator,
                request_id,
                model_name,
                conversation,
                tokenizer,
                request_metadata,
                reasoning_parser,
                mm_token_counts=mm_token_counts,
            )
        except ValueError as e:
            return self.create_error_response(e)

    async def _preprocess_chat(
        self,
        request: ChatLikeRequest | ResponsesRequest,
        messages: list[ChatCompletionMessageParam],
        default_template: str | None,
        default_template_content_format: ChatTemplateContentFormatOption,
        default_template_kwargs: dict[str, Any] | None = None,
        tool_dicts: list[dict[str, Any]] | None = None,
        tool_parser: Callable[[TokenizerLike], ToolParser] | None = None,
        # OMNI: Additional parameters for backward compatibility
        renderer: BaseRenderer | None = None,
        add_generation_prompt: bool = True,
        continue_final_message: bool = False,
        documents: list[dict[str, str]] | None = None,
        add_special_tokens: bool = False,
    ) -> tuple[list[ConversationMessage], list[TokPrompt]]:
        if renderer is None:
            renderer = self.renderer

        # Keep OMNI compatibility args wired while delegating rendering
        # to the upstream async renderer pipeline.
        default_template_kwargs = merge_kwargs(
            default_template_kwargs,
            dict(
                tools=tool_dicts,
                documents=documents,
                add_generation_prompt=add_generation_prompt,
                continue_final_message=continue_final_message,
                add_special_tokens=add_special_tokens,
                tokenize=isinstance(renderer.tokenizer, MistralTokenizer),
            ),
        )

        tok_params = request.build_tok_params(self.model_config)
        mm_config = self.model_config.multimodal_config
        chat_params = request.build_chat_params(
            default_template,
            default_template_content_format,
        ).with_defaults(
            default_template_kwargs,
            default_media_io_kwargs=(mm_config.media_io_kwargs if mm_config else None),
            default_mm_processor_kwargs=getattr(request, "mm_processor_kwargs", None),
        )

        deferred_multi_modal_data: dict[str, Any] | None = None
        if self._needs_multistage_multimodal_split():
            messages, deferred_multi_modal_data = await self._prepare_multistage_multimodal_inputs(
                messages,
                request,
            )

        (conversation,), (engine_prompt,) = await renderer.render_chat_async(
            [messages],
            chat_params,
            tok_params,
            prompt_extras={
                k: v for k in ("mm_processor_kwargs", "cache_salt") if (v := getattr(request, k, None)) is not None
            },
        )

        tokenizer = renderer.get_tokenizer()

        # tool parsing is done only if a tool_parser has been set and if
        # tool_choice is not "none" (if tool_choice is "none" but a tool_parser
        # is set, we want to prevent parsing a tool_call hallucinated by the LLM
        should_parse_tools = tool_parser is not None and (
            hasattr(request, "tool_choice") and request.tool_choice != "none"
        )

        if should_parse_tools:
            if not isinstance(request, ChatCompletionRequest):
                msg = "Tool usage is only supported for Chat Completions API"
                raise NotImplementedError(msg)

            request = tool_parser(tokenizer).adjust_request(  # type: ignore
                request=request
            )

        # Preserve a clean text prompt for downstream stages (e.g., GLM-Image diffusion).
        # For image generation, we want the raw user caption instead of a rendered template.
        # But for multimodal comprehension (img2text), we MUST keep the rendered prompt
        # containing image tokens.
        req_modalities = getattr(request, "modalities", [])
        if req_modalities and ("image" in req_modalities):
            extracted_prompt, _ = self._extract_diffusion_prompt_and_images_from_messages(messages)
            if extracted_prompt:
                engine_prompt["prompt"] = extracted_prompt

        mm_processor_kwargs = getattr(request, "mm_processor_kwargs", None)
        if mm_processor_kwargs is not None:
            engine_prompt["mm_processor_kwargs"] = mm_processor_kwargs

        if hasattr(request, "cache_salt") and request.cache_salt is not None:
            engine_prompt["cache_salt"] = request.cache_salt

        additional_information = getattr(request, "additional_information", None)
        if isinstance(additional_information, dict):
            prompt_additional_information = self._ensure_prompt_additional_information(engine_prompt)
            prompt_additional_information.update(additional_information)

        if deferred_multi_modal_data:
            prompt_additional_information = self._ensure_prompt_additional_information(engine_prompt)
            prompt_additional_information["deferred_multi_modal_data"] = deferred_multi_modal_data

        speaker = getattr(request, "voice", None) or getattr(request, "speaker", None)
        normalized = validate_requested_speaker(speaker, self._get_supported_speakers())
        if normalized is not None:
            prompt_additional_information = self._ensure_prompt_additional_information(engine_prompt)
            prompt_additional_information["speaker"] = [normalized]

        language = getattr(request, "language", None)
        if language is not None and isinstance(language, str) and language.strip():
            prompt_additional_information = self._ensure_prompt_additional_information(engine_prompt)
            prompt_additional_information["language"] = [language.strip()]

        # Style instruction — used by Ming-flash-omni instruct TTS path
        # (ming_task="instruct").  For the omni speech path the thinker2talker
        # bridge drops this field to match upstream omni_audio_generation
        # which hardcodes instruction=None.
        instructions = getattr(request, "instructions", None)
        if instructions is not None and isinstance(instructions, str) and instructions.strip():
            prompt_additional_information = self._ensure_prompt_additional_information(engine_prompt)
            prompt_additional_information["instruction"] = instructions.strip()

        return conversation, [engine_prompt]

    @staticmethod
    def _ensure_prompt_additional_information(engine_prompt: dict[str, Any]) -> dict[str, Any]:
        additional_information = engine_prompt.get("additional_information")
        if not isinstance(additional_information, dict):
            additional_information = {}
            engine_prompt["additional_information"] = additional_information
        return additional_information

    def _needs_multistage_multimodal_split(self) -> bool:
        return bool(self._deferred_multimodal_modalities())

    def _deferred_multimodal_modalities(self) -> set[str]:
        stage_configs = list(getattr(self.engine_client, "stage_configs", []) or [])
        if len(stage_configs) < 2:
            return set()

        first_stage_modalities = self._stage_input_modalities(stage_configs[0])
        if not first_stage_modalities:
            return set()

        downstream_modalities: set[str] = set()
        for stage in stage_configs[1:]:
            downstream_modalities.update(self._stage_input_modalities(stage))
        return downstream_modalities - first_stage_modalities

    @staticmethod
    def _stage_input_modalities(stage: Any) -> set[str]:
        engine_args = getattr(stage, "engine_args", None)
        explicit = (
            getattr(stage, "input_modalities", None)
            or getattr(stage, "modalities", None)
            or getattr(engine_args, "input_modalities", None)
            or getattr(engine_args, "modalities", None)
        )
        if explicit:
            return {str(modality) for modality in as_list(explicit)}

        model_stage = str(getattr(engine_args, "model_stage", None) or getattr(stage, "model_stage", "")).lower()
        if model_stage in {"asr", "stt"} or model_stage.endswith("_asr"):
            return {"audio"}
        if any(name in model_stage for name in ("vision", "vl", "aura")):
            return {"image", "video"}
        if any(name in model_stage for name in ("tts", "talker", "code2wav")):
            return set()

        if getattr(stage, "requires_multimodal_data", False):
            return {"audio", "image", "video"}
        return set()

    async def _prepare_multistage_multimodal_inputs(
        self,
        messages: list[ChatCompletionMessageParam],
        request: ChatLikeRequest | ResponsesRequest,
    ) -> tuple[list[ChatCompletionMessageParam], dict[str, Any] | None]:
        """Hide modalities unsupported by stage 0 and stash them for downstream stages."""
        deferred_modalities = self._deferred_multimodal_modalities()
        deferred_parts: dict[str, list[Any]] = {modality: [] for modality in deferred_modalities}
        stripped_messages: list[ChatCompletionMessageParam] = []

        for message in messages:
            content = message.get("content") if isinstance(message, dict) else None
            if not isinstance(content, list):
                stripped_messages.append(message)
                continue

            stripped_content: list[Any] = []
            changed = False
            for part in content:
                modality, payload = self._deferred_multimodal_part(part, deferred_modalities)
                if modality is not None:
                    if payload is not None:
                        deferred_parts.setdefault(modality, []).append(payload)
                    changed = True
                    continue
                stripped_content.append(part)

            if changed:
                stripped_message = dict(message)
                stripped_message["content"] = stripped_content
                stripped_messages.append(cast(ChatCompletionMessageParam, stripped_message))
            else:
                stripped_messages.append(message)

        deferred_parts = {modality: parts for modality, parts in deferred_parts.items() if parts}
        if not deferred_parts:
            return messages, None

        media_connector = MediaConnector(
            media_io_kwargs=getattr(request, "media_io_kwargs", None),
            allowed_local_media_path=getattr(self.model_config, "allowed_local_media_path", "") or "",
            allowed_media_domains=getattr(self.model_config, "allowed_media_domains", None),
        )
        multi_modal_data: dict[str, Any] = {}
        for modality, parts in deferred_parts.items():
            multi_modal_data[modality] = await self._materialize_deferred_multimodal_parts(
                media_connector,
                modality,
                parts,
            )
        return stripped_messages, multi_modal_data

    @staticmethod
    def _deferred_multimodal_part(part: Any, deferred_modalities: set[str]) -> tuple[str | None, Any | None]:
        if not isinstance(part, dict):
            return None, None
        part_type = part.get("type")
        if part_type in {"video_url", "video"} and "video" in deferred_modalities:
            video = part.get("video_url", part.get("video"))
            if isinstance(video, dict):
                video = video.get("url")
            return "video", video
        if part_type in {"image_url", "image_pil", "image"} and "image" in deferred_modalities:
            image = part.get("image_pil", part.get("image_url", part.get("image")))
            if isinstance(image, dict):
                image = image.get("url")
            return "image", image
        if part_type in {"audio_url", "input_audio", "audio"} and "audio" in deferred_modalities:
            audio = part.get("audio_url", part.get("input_audio", part.get("audio")))
            if isinstance(audio, dict):
                audio = audio.get("url") or audio.get("data")
            return "audio", audio
        return None, None

    @staticmethod
    async def _materialize_deferred_multimodal_parts(
        media_connector: MediaConnector,
        modality: str,
        parts: list[Any],
    ) -> list[Any]:
        if modality == "video":
            return list(await asyncio.gather(*(media_connector.fetch_video_async(part) for part in parts)))
        if modality == "image":
            fetch_image = getattr(media_connector, "fetch_image_async", None)
            if fetch_image is None:
                return parts
            return list(
                await asyncio.gather(
                    *(fetch_image(part) if isinstance(part, str) else _identity_async(part) for part in parts)
                )
            )
        if modality == "audio":
            fetch_audio = getattr(media_connector, "fetch_audio_async", None)
            if fetch_audio is None:
                return parts
            return list(
                await asyncio.gather(
                    *(fetch_audio(part) if isinstance(part, str) else _identity_async(part) for part in parts)
                )
            )
        return parts

    async def _inject_audio_from_video_urls(
        self,
        messages: list[ChatCompletionMessageParam],
    ) -> list[ChatCompletionMessageParam]:
        """Pre-extract audio from video URLs and inject as audio_url content items.

        When use_audio_in_video=True, the qwen2_5_omni_thinker multimodal
        processor requires that the number of audio items equals the number of
        video items (it subtracts mm_counts["video"] from mm_counts["audio"]).
        The client only sends video_url items; this method adds the matching
        audio_url items on the server side before the renderer processes them.
        """
        import io

        from vllm_omni.entrypoints.chat_utils import extract_audio_from_video_async

        new_messages: list[ChatCompletionMessageParam] = []
        for msg in messages:
            content = msg.get("content") if isinstance(msg, dict) else getattr(msg, "content", None)
            if not isinstance(content, list):
                new_messages.append(msg)
                continue

            video_urls = [
                part.get("video_url", {}).get("url")
                for part in content
                if isinstance(part, dict) and part.get("type") == "video_url" and part.get("video_url", {}).get("url")
            ]

            if not video_urls:
                new_messages.append(msg)
                continue

            audios = await asyncio.gather(*(extract_audio_from_video_async(u) for u in video_urls))

            audio_items: list[dict] = []
            for audio_array, sample_rate in audios:
                buf = io.BytesIO()
                if soundfile is not None:
                    soundfile.write(buf, audio_array, samplerate=int(sample_rate), format="WAV")
                else:
                    import struct

                    import numpy as np

                    audio_np = np.asarray(audio_array, dtype=np.float32)
                    sr = int(sample_rate)
                    num_channels = 1
                    bits_per_sample = 32
                    num_frames = len(audio_np)
                    data_size = num_frames * num_channels * (bits_per_sample // 8)
                    # Write minimal RIFF/WAV header
                    buf.write(b"RIFF")
                    buf.write(struct.pack("<I", 36 + data_size))
                    buf.write(b"WAVE")
                    buf.write(b"fmt ")
                    buf.write(
                        struct.pack(
                            "<IHHIIHH",
                            16,
                            3,
                            num_channels,
                            sr,
                            sr * num_channels * (bits_per_sample // 8),
                            num_channels * (bits_per_sample // 8),
                            bits_per_sample,
                        )
                    )
                    buf.write(b"data")
                    buf.write(struct.pack("<I", data_size))
                    buf.write(audio_np.tobytes())

                audio_b64 = base64.b64encode(buf.getvalue()).decode()
                audio_items.append(
                    {
                        "type": "audio_url",
                        "audio_url": {"url": f"data:audio/wav;base64,{audio_b64}"},
                    }
                )

            new_content = list(content) + audio_items
            if isinstance(msg, dict):
                new_msg = {**msg, "content": new_content}
            else:
                new_msg = msg.model_copy(update={"content": new_content})
            new_messages.append(new_msg)

        return new_messages

    def _to_sampling_params_list(self, sampling_params_list: list[dict]) -> list[Any]:
        """Convert request dicts to stage-typed sampling params objects.

        For diffusion stages, build ``OmniDiffusionSamplingParams`` so
        downstream ``StageDiffusionClient._sampling_params_to_dict`` (which
        requires a dataclass) works. For LLM stages build ``SamplingParams``.
        If callers provide params for fewer stages than the native pipeline has
        (for example AURA has three semantic models but four engine stages),
        append cloned deploy defaults for the omitted tail stages.
        """
        stage_configs = list(getattr(self.engine_client, "stage_configs", []) or [])
        default_params_list = list(getattr(self.engine_client, "default_sampling_params_list", []) or [])
        final_sampling_params_list: list[Any] = []
        for idx, sampling_params in enumerate(sampling_params_list):
            stage_type = get_stage_type(stage_configs[idx]) if idx < len(stage_configs) else "llm"
            target_cls = OmniDiffusionSamplingParams if stage_type == "diffusion" else SamplingParams
            if isinstance(sampling_params, dict):
                final_sampling_params_list.append(target_cls(**sampling_params))
            elif isinstance(sampling_params, target_cls):
                final_sampling_params_list.append(sampling_params)
            elif isinstance(sampling_params, SamplingParams | OmniDiffusionSamplingParams):
                # Cross-typed (e.g. user passed SamplingParams but this is a
                # diffusion stage) — rebuild via a dict round-trip so we end
                # up with the correct target class.
                as_dict = {
                    f.name: getattr(sampling_params, f.name)
                    for f in (fields(sampling_params) if is_dataclass(sampling_params) else [])
                } or sampling_params.__dict__
                final_sampling_params_list.append(target_cls(**as_dict))
            else:
                raise ValueError(f"Invalid sampling params: {sampling_params}")
        for idx in range(len(final_sampling_params_list), len(stage_configs)):
            if idx < len(default_params_list):
                final_sampling_params_list.append(clone_sampling_params(default_params_list[idx]))
            else:
                final_sampling_params_list.append(SamplingParams())
        return final_sampling_params_list

    def _get_comprehension_stage_index(self) -> int:
        for idx, stage in enumerate(self.engine_client.stage_configs):
            if stage.is_comprehension:
                return idx
        raise ValueError("No comprehension stage (is_comprehension=True) found in stage configs")

    # OpenAI API standard sampling parameters that can be safely overridden.
    # These are the most commonly used parameters with compatible types
    # between ChatCompletionRequest and SamplingParams.
    # Users who need more control can use sampling_params_list in extra_body.
    _OPENAI_SAMPLING_FIELDS: set[str] = {
        "temperature",
        "top_p",
        "top_k",
        "max_tokens",
        "min_tokens",
        "seed",
        "ignore_eos",
        "stop",
        "stop_token_ids",
        "frequency_penalty",
        "presence_penalty",
    }

    def _apply_request_overrides(
        self,
        default_params: SamplingParams,
        request: ChatCompletionRequest,
    ) -> SamplingParams:
        """Clone default params and override with user-provided request values.

        Starts with YAML defaults and only overrides fields that the user
        explicitly provided (non-None values) in the request.

        max_tokens is NOT computed dynamically — it uses the deploy YAML default.

        Args:
            default_params: Default SamplingParams from the resolved pipeline and deploy configuration.
            request: The chat completion request containing user-provided values.

        Returns:
            New SamplingParams with YAML defaults overridden by request values.
        """
        params = default_params.clone()

        # Only apply fields explicitly provided by user, not protocol defaults.
        # Pydantic v2 uses `model_fields_set`; keep v1 fallback for compatibility.
        explicit_fields = getattr(request, "model_fields_set", None)
        if explicit_fields is None:
            explicit_fields = getattr(request, "__fields_set__", set())

        for field_name in self._OPENAI_SAMPLING_FIELDS:
            if field_name not in explicit_fields:
                continue

            value = getattr(request, field_name, None)
            if (value is not None and not isinstance(value, list)) or (isinstance(value, list) and len(value) > 0):
                setattr(params, field_name, value)

        return params

    @staticmethod
    def _set_if_supported(obj: Any, **kwargs: Any) -> None:
        for key, value in kwargs.items():
            if value is not None and hasattr(obj, key):
                setattr(obj, key, value)

    def _should_check_for_unstreamed_tool_arg_tokens(
        self,
        delta_message: Any,
        output: Any,
    ) -> bool:
        """Check whether the streaming generator should flush unstreamed
        tool-arg tokens at finish-time.

        This method was moved from OpenAIServingChat to the tool-parser layer
        by upstream commit 9affc17a05.  Omni's independently-maintained
        ``chat_completion_stream_generator`` still calls it, so we keep a
        local copy with the pre-9affc17a05 semantics.
        """
        return (
            output.finish_reason is not None
            and self.enable_auto_tools
            and self.parser_cls is not None
            and self.parser_cls.tool_parser_cls is not None
            and delta_message is not None
            and delta_message.tool_calls
        )

    def _build_sampling_params_list_from_request(
        self,
        request: ChatCompletionRequest,
    ) -> list[SamplingParams]:
        """Build sampling_params_list using standard OpenAI API parameters.

        For the comprehension stage, starts with YAML defaults and overrides with
        user-provided request values. For other stages, uses cloned YAML defaults.

        This approach ensures all YAML defaults (including seed, detokenize, etc.)
        are preserved while allowing users to override specific parameters.

        Args:
            request: The chat completion request containing OpenAI API parameters.

        Returns:
            List of SamplingParams, one for each stage.
        """
        default_params_list = self.engine_client.default_sampling_params_list
        comprehension_idx = self._get_comprehension_stage_index()

        sampling_params_list = []
        for idx, default_params in enumerate(default_params_list):
            if isinstance(default_params, dict):
                default_params = SamplingParams(**default_params)
            if idx == comprehension_idx:
                params = self._apply_request_overrides(default_params, request)
                sampling_params_list.append(params)
            else:
                # For other stages, clone default params
                sampling_params_list.append(default_params.clone())

        return sampling_params_list

    def _log_inputs(
        self,
        request_id: str,
        inputs: PromptType | TokPrompt,
        params_list: list[SamplingParams] | None,
        lora_request: LoRARequest | None,
    ) -> None:
        if self.request_logger is None:
            return
        components = self._extract_prompt_components(inputs)
        self.request_logger.log_inputs(
            request_id,
            components.text,
            components.token_ids,
            components.embeds,
            params=params_list,
            lora_request=lora_request,
        )

    async def chat_completion_stream_generator(
        self,
        request: ChatCompletionRequest,
        result_generator: AsyncIterator[RequestOutput],
        request_id: str,
        model_name: str,
        conversation: list[ConversationMessage],
        tokenizer: AnyTokenizer,
        request_metadata: RequestResponseMetadata,
        reasoning_parser: ReasoningParser | None = None,
        mm_token_counts: dict[str, int] | None = None,
        raw_request: Request | None = None,
    ):
        created_time = int(time.time())
        chunk_object_type: Final = "chat.completion.chunk"
        first_iteration_dict = {}
        assert hasattr(request, "modalities") and request.modalities is not None, (
            "Streaming request must specify output modalities"
        )
        for modality in request.modalities:
            first_iteration_dict[modality] = True

        # Send response for each token for each request.n (index)
        num_choices = 1 if request.n is None else request.n
        previous_num_tokens = [0] * num_choices
        finish_reason_sent = [False] * num_choices
        modality_finished: list[set[str]] = [set() for _ in range(num_choices)]
        modality_seen: list[set[str]] = [set() for _ in range(num_choices)]
        stop_reason_emitted: list[bool] = [False] * num_choices
        num_prompt_tokens = 0
        num_cached_tokens = None
        num_cache_creation_tokens = None
        if self.use_harmony:
            harmony_parsers = [get_streamable_parser_for_assistant() for _ in range(num_choices)]
            harmony_tools_streamed = [False] * num_choices
        tools_streamed = [False] * num_choices

        if isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam):
            tool_choice_function_name = request.tool_choice.function.name
        else:
            tool_choice_function_name = None

        # Determine whether tools are in use with "auto" tool choice
        tool_choice_auto = not tool_choice_function_name and self._should_stream_with_auto_tool_parsing(request)

        all_previous_token_ids: list[list[int]] | None
        function_name_returned = [False] * num_choices
        if self.tool_call_id_type == "kimi_k2":
            history_tool_call_cnt = get_history_tool_calls_cnt(conversation)
        else:
            history_tool_call_cnt = 0

        # Always track previous_texts for comprehensive output logging
        previous_texts = [""] * num_choices

        # Only one of these will be used, thus previous_texts and
        # all_previous_token_ids will not be used twice in the same iteration.
        if tool_choice_auto or reasoning_parser:
            # These are only required in "auto" tool choice case
            all_previous_token_ids = [[]] * num_choices
            # For reasoning parser and tool call all enabled
            added_content_delta_arr = [False] * num_choices
            reasoning_end_arr = [False] * num_choices
        else:
            all_previous_token_ids = None
        # Prepare the tool parser if it's needed
        try:
            tool_parser_cls = self.parser_cls.tool_parser_cls if self.parser_cls is not None else None
            if tool_choice_auto and tool_parser_cls is not None:
                tool_parsers: list[ToolParser | None] = [
                    tool_parser_cls(tokenizer, request.tools) for _ in range(num_choices)
                ]
            else:
                tool_parsers = [None] * num_choices
        except Exception as e:
            logger.exception("Error in tool parser creation.")
            data = self.create_streaming_error_response(e)
            yield f"data: {data}\n\n"
            yield "data: [DONE]\n\n"
            return

        stream_options = request.stream_options
        include_usage, include_continuous_usage = should_include_usage(stream_options, self.enable_force_include_usage)

        last_metrics: dict[str, Any] | None = None
        # Hold a strong reference to the audio request state so the
        # streaming-finalize hook below survives the inner generator's
        # _log_summary_and_cleanup, which pops request_states[request_id]
        # before this outer block runs.
        req_state_audio_ref: Any = None
        try:
            async for omni_res in result_generator:
                final_output_type = omni_res.final_output_type
                res = omni_res
                if final_output_type not in first_iteration_dict:
                    logger.warning(f"final output type: {final_output_type} is not needed by the request")
                    continue

                # Track which modalities have actually appeared in the stream.
                # This is used to determine when all *produced* modalities have
                # finished, which may be a subset of request.modalities when the
                # engine does not produce every requested modality.
                for output in res.outputs:
                    modality_seen[output.index].add(final_output_type)

                if omni_res.metrics:
                    last_metrics = omni_res.metrics

                # Initialize role before conditional blocks to avoid UnboundLocalError
                # when handling audio/image responses
                role = self.get_chat_request_role(request)

                # Compute prompt_text once at first iteration (upstream #42052)
                prompt_text = getattr(res, "prompt", None) if getattr(request, "return_prompt_text", None) else None

                # We need to do it here, because if there are exceptions in
                # the result_generator, it needs to be sent as the FIRST
                # response (by the try...catch).
                if first_iteration_dict[final_output_type] and final_output_type == "text":
                    # NOTE: prompt token IDs / cached tokens are based on the first iteration
                    # of the stage that producing text output for consistency for current
                    # non-stream behaviors.
                    if res.prompt_token_ids is not None:
                        num_prompt_tokens = len(res.prompt_token_ids)
                        if res.encoder_prompt_token_ids is not None:
                            num_prompt_tokens += len(res.encoder_prompt_token_ids)

                    num_cached_tokens = res.num_cached_tokens
                    num_cache_creation_tokens = res.num_cache_creation_tokens
                    # Send first response for each choice with role
                    # NOTE: num_choices defaults to 1 so this usually executes once per request
                    for i in range(num_choices):
                        choice_data = ChatCompletionResponseStreamChoice(
                            index=i,
                            delta=DeltaMessage(
                                role=role,
                                content="",
                            ),
                            logprobs=None,
                            finish_reason=None,
                        )

                        # return prompt_token_ids at the first chunk ever
                        chunk = OmniChatCompletionStreamResponse(
                            id=request_id,
                            object=chunk_object_type,
                            created=created_time,
                            choices=[choice_data],
                            model=model_name,
                            prompt_token_ids=(res.prompt_token_ids if request.return_token_ids else None),
                            prompt_text=prompt_text,
                            modality=final_output_type,
                        )

                        # if continuous usage stats are requested, add it
                        if include_continuous_usage:
                            chunk.usage = UsageInfo(
                                prompt_tokens=num_prompt_tokens,
                                completion_tokens=0,
                                total_tokens=num_prompt_tokens,
                            )

                        data = chunk.model_dump_json(exclude_unset=True)
                        yield f"data: {data}\n\n"

                    # Send response to echo the input portion of the
                    # last message
                    if request.echo:
                        last_msg_content: str | list[dict[str, str]] = ""
                        if conversation and "content" in conversation[-1] and conversation[-1].get("role") == role:
                            last_msg_content = conversation[-1]["content"] or ""

                        if last_msg_content:
                            for i in range(num_choices):
                                choice_data = ChatCompletionResponseStreamChoice(
                                    index=i,
                                    delta=DeltaMessage(content=last_msg_content),
                                    logprobs=None,
                                    finish_reason=None,
                                )
                                chunk = OmniChatCompletionStreamResponse(
                                    id=request_id,
                                    object=chunk_object_type,
                                    created=created_time,
                                    choices=[choice_data],
                                    model=model_name,
                                    modality=final_output_type,
                                )
                                if include_continuous_usage:
                                    chunk.usage = UsageInfo(
                                        prompt_tokens=num_prompt_tokens,
                                        completion_tokens=0,
                                        total_tokens=num_prompt_tokens,
                                    )

                                data = chunk.model_dump_json(exclude_unset=True)
                                yield f"data: {data}\n\n"
                    first_iteration_dict[final_output_type] = False

                if final_output_type == "text":
                    for output in res.outputs:
                        i = output.index
                        tool_parser = tool_parsers[i]

                        if finish_reason_sent[i]:
                            continue

                        if request.logprobs and request.top_logprobs is not None:
                            assert output.logprobs is not None, "Did not output logprobs"
                            logprobs = self._create_chat_logprobs(
                                token_ids=output.token_ids,
                                top_logprobs=output.logprobs,
                                tokenizer=tokenizer,
                                num_output_top_logprobs=request.top_logprobs,
                                return_as_token_id=request.return_tokens_as_token_ids,
                            )
                        else:
                            logprobs = None

                        if self.use_harmony:
                            harmony_parser = harmony_parsers[i]
                            prev_recipient = harmony_parser.current_recipient
                            delta_text = ""
                            for token_id in output.token_ids:
                                harmony_parser.process(token_id)
                                delta_text += harmony_parser.last_content_delta or ""
                            cur_channel = harmony_parser.current_channel
                            cur_recipient = harmony_parser.current_recipient
                        else:
                            delta_text = output.text or ""

                        if not delta_text and not output.token_ids and not previous_num_tokens[i]:
                            # Chunked prefill case, don't return empty chunks
                            continue

                        delta_message: DeltaMessage | None

                        # just update previous_texts and previous_token_ids
                        if tool_choice_auto or reasoning_parser:
                            assert previous_texts is not None
                            assert all_previous_token_ids is not None
                            previous_text = previous_texts[i]
                            previous_token_ids = all_previous_token_ids[i]
                            current_text = previous_text + delta_text
                            # avoid the None + list error.
                            if previous_token_ids:
                                current_token_ids = previous_token_ids + as_list(output.token_ids)
                            else:
                                current_token_ids = as_list(output.token_ids)

                        if self.use_harmony:
                            if cur_channel == "final":
                                delta_message = DeltaMessage(content=delta_text)
                            elif cur_channel == "analysis":
                                if request.include_reasoning:
                                    delta_message = DeltaMessage(reasoning=delta_text)
                                else:
                                    delta_message = None
                            elif (
                                cur_channel == "commentary" and cur_recipient and cur_recipient.startswith("functions.")
                            ):
                                # Count completed tool calls to determine index
                                base_index = 0
                                for msg in harmony_parser.messages:
                                    if (
                                        msg.channel == "commentary"
                                        and msg.recipient
                                        and msg.recipient.startswith("functions.")
                                    ):
                                        base_index += 1

                                if prev_recipient != cur_recipient:
                                    tool_name = cur_recipient.split("functions.", 1)[1]
                                    delta_message = DeltaMessage(
                                        tool_calls=[
                                            DeltaToolCall(
                                                id=make_tool_call_id(),
                                                type="function",
                                                function=DeltaFunctionCall(
                                                    name=tool_name,
                                                    arguments="",
                                                ),
                                                index=base_index,
                                            )
                                        ]
                                    )
                                elif delta_text:
                                    delta_message = DeltaMessage(
                                        tool_calls=[
                                            DeltaToolCall(
                                                index=base_index,
                                                function=DeltaFunctionCall(arguments=delta_text),
                                            )
                                        ]
                                    )
                                else:
                                    delta_message = None

                                if delta_message is not None:
                                    harmony_tools_streamed[i] = True
                            else:
                                delta_message = None
                        # handle streaming deltas for tools with named tool_choice
                        elif tool_choice_function_name:
                            if (
                                reasoning_parser
                                and not reasoning_end_arr[i]
                                and not reasoning_parser.is_reasoning_end(previous_token_ids)
                            ):
                                assert reasoning_parser is not None
                                delta_message = reasoning_parser.extract_reasoning_streaming(
                                    previous_text,
                                    current_text,
                                    delta_text,
                                    previous_token_ids,
                                    current_token_ids,
                                    output.token_ids,
                                )
                                # When encountering think end id in delta_token_ids
                                # or think end id in prompt_token_ids
                                # i.e {"enable_thinking": False},
                                # set reasoning status to end.
                                # Only keep 'content', remove 'reasoning'.
                                if reasoning_parser.is_reasoning_end(as_list(output.token_ids)) or (
                                    res.prompt_token_ids and reasoning_parser.is_reasoning_end(res.prompt_token_ids)
                                ):
                                    reasoning_end_arr[i] = True
                                    if delta_message and delta_message.content:
                                        # This need to be added to next `delta_text`
                                        current_text = delta_message.content
                                        delta_message.content = None
                                    else:
                                        current_text = ""
                            else:
                                # Just to add remaining `content`
                                if reasoning_parser:
                                    delta_text = previous_text + delta_text
                                    current_text = ""

                                if function_name_returned[i]:
                                    delta_tool_call = DeltaToolCall(
                                        function=DeltaFunctionCall(arguments=delta_text),
                                        index=i,
                                    )
                                else:
                                    delta_tool_call = DeltaToolCall(
                                        id=make_tool_call_id(),
                                        type="function",
                                        function=DeltaFunctionCall(
                                            name=tool_choice_function_name,
                                            arguments=delta_text,
                                        ),
                                        index=i,
                                    )
                                    function_name_returned[i] = True

                                delta_message = DeltaMessage(
                                    tool_calls=[
                                        delta_tool_call,
                                    ]
                                )
                                tools_streamed[i] = True

                        elif request.tool_choice == "required":
                            assert previous_texts is not None
                            previous_text = previous_texts[i]
                            current_text = previous_text + delta_text
                            fn_name_returned = function_name_returned[i]
                            output_token_ids = as_list(output.token_ids)

                            if (
                                reasoning_parser is not None
                                and not reasoning_end_arr[i]
                                and res.prompt_token_ids
                                and reasoning_parser.is_reasoning_end(res.prompt_token_ids)
                            ):
                                reasoning_end_arr[i] = True

                            if reasoning_parser and not reasoning_end_arr[i]:
                                delta_message = reasoning_parser.extract_reasoning_streaming(
                                    previous_text,
                                    current_text,
                                    delta_text,
                                    previous_token_ids,
                                    current_token_ids,
                                    output_token_ids,
                                )
                                if reasoning_parser.is_reasoning_end(output_token_ids):
                                    reasoning_end_arr[i] = True
                                    if delta_message and delta_message.content:
                                        current_text = delta_message.content
                                        delta_message.content = None
                                    else:
                                        # reasoning ended
                                        current_text = ""

                            else:
                                # either finished reasoning or no reasoning at all
                                content = current_text

                                delta_message, function_name_returned[i] = extract_required_tool_call_streaming(
                                    previous_text=previous_text,
                                    current_text=content,
                                    delta_text=delta_text,
                                    function_name_returned=fn_name_returned,
                                    tool_call_idx=history_tool_call_cnt,
                                    tool_call_id_type=self.tool_call_id_type,
                                )
                                if (
                                    delta_message
                                    and delta_message.tool_calls
                                    and delta_message.tool_calls[0].id is not None
                                ):
                                    history_tool_call_cnt += 1
                                    tools_streamed[i] = True

                        # handle streaming deltas for tools with "auto" tool choice
                        # and reasoning parser
                        elif tool_choice_auto and reasoning_parser:
                            assert tool_parser is not None
                            assert reasoning_parser is not None
                            assert added_content_delta_arr is not None
                            assert reasoning_end_arr is not None
                            output_token_ids = as_list(output.token_ids)
                            if not reasoning_end_arr[i]:
                                delta_message = reasoning_parser.extract_reasoning_streaming(
                                    previous_text,
                                    current_text,
                                    delta_text,
                                    previous_token_ids,
                                    current_token_ids,
                                    output_token_ids,
                                )
                                # When encountering think end id in prompt_token_ids
                                # i.e {"enable_thinking": False},
                                # set reasoning status to end.
                                # Remove the text and token ids related
                                # to 'reasoning'.
                                if res.prompt_token_ids and reasoning_parser.is_reasoning_end(res.prompt_token_ids):
                                    reasoning_end_arr[i] = True
                                    current_token_ids = output_token_ids
                                    if delta_message and delta_message.content:
                                        current_text = delta_message.content
                                        delta_message.content = None
                                    else:
                                        current_text = ""
                                # When encountering think end id in delta_token_ids,
                                # set reasoning status to end.
                                # Remove the text and token ids related
                                # to 'reasoning'.
                                if reasoning_parser.is_reasoning_end(output_token_ids):
                                    reasoning_end_arr[i] = True
                                    current_token_ids = reasoning_parser.extract_content_ids(output_token_ids)
                                    if delta_message and delta_message.content:
                                        current_text = delta_message.content
                                        delta_message.content = None
                                    else:
                                        current_text = ""

                            # handle tool calls only after reasoning is done,
                            else:
                                delta_token_ids = output_token_ids
                                # First time to tool call,
                                # add the remaining text and token ids
                                # to delta from previous
                                if not added_content_delta_arr[i]:
                                    added_content_delta_arr[i] = True
                                    previous_text = ""
                                    previous_token_ids = []
                                    delta_text = current_text
                                    delta_token_ids = current_token_ids

                                delta_message = tool_parser.extract_tool_calls_streaming(
                                    previous_text=previous_text,
                                    current_text=current_text,
                                    delta_text=delta_text,
                                    previous_token_ids=previous_token_ids,
                                    current_token_ids=current_token_ids,
                                    delta_token_ids=delta_token_ids,
                                    request=request,
                                )
                                if delta_message and delta_message.tool_calls:
                                    tools_streamed[i] = True
                        # when only tool calls
                        elif tool_choice_auto:
                            assert tool_parser is not None
                            delta_message = tool_parser.extract_tool_calls_streaming(
                                previous_text=previous_text,
                                current_text=current_text,
                                delta_text=delta_text,
                                previous_token_ids=previous_token_ids,
                                current_token_ids=current_token_ids,
                                delta_token_ids=output.token_ids,
                                request=request,
                            )
                            if delta_message and delta_message.tool_calls:
                                tools_streamed[i] = True

                        # when only reasoning
                        elif reasoning_parser:
                            delta_message = reasoning_parser.extract_reasoning_streaming(
                                previous_text,
                                current_text,
                                delta_text,
                                previous_token_ids,
                                current_token_ids,
                                output.token_ids,
                            )
                        # handle streaming just a content delta
                        else:
                            delta_message = DeltaMessage(content=delta_text)

                        # update the previous values for the next iteration
                        if (tool_choice_auto or reasoning_parser) and not self.use_harmony:
                            assert previous_texts is not None
                            assert all_previous_token_ids is not None
                            previous_texts[i] = current_text
                            all_previous_token_ids[i] = current_token_ids
                        else:
                            # Update for comprehensive logging even in simple case
                            assert previous_texts is not None
                            previous_texts[i] += delta_text

                        # set the previous values for the next iteration
                        previous_num_tokens[i] += len(output.token_ids)

                        # if the message delta is None (e.g. because it was a
                        # "control token" for tool calls or the parser otherwise
                        # wasn't ready to send a token, then
                        #   get the next token without streaming a chunk
                        if delta_message is None:
                            if output.finish_reason is None and not request.return_token_ids:
                                continue
                            delta_message = DeltaMessage()

                        # Log streaming delta if output logging is enabled
                        if self.enable_log_outputs and self.request_logger:
                            delta_content = ""
                            if delta_message.content:
                                delta_content = delta_message.content
                            elif delta_message.tool_calls:
                                delta_content = "".join(
                                    tc.function.arguments
                                    for tc in delta_message.tool_calls
                                    if tc.function and tc.function.arguments
                                )

                            if delta_content:
                                self.request_logger.log_outputs(
                                    request_id=request_id,
                                    outputs=delta_content,
                                    output_token_ids=as_list(output.token_ids),
                                    finish_reason=output.finish_reason,
                                    is_streaming=True,
                                    delta=True,
                                )

                        if output.finish_reason is None:
                            # Send token-by-token response for each request.n
                            choice_data = ChatCompletionResponseStreamChoice(
                                index=i,
                                delta=delta_message,
                                logprobs=logprobs,
                                finish_reason=None,
                                token_ids=(as_list(output.token_ids) if request.return_token_ids else None),
                            )

                        # if the model is finished generating
                        else:
                            # check to make sure we haven't "forgotten" to stream
                            #   any tokens that were generated but previously
                            #   matched by partial json parsing
                            # only happens if we are NOT using structured outputs
                            auto_tools_called = False
                            if tool_parser:
                                auto_tools_called = len(tool_parser.prev_tool_call_arr) > 0
                                index = len(tool_parser.prev_tool_call_arr) - 1 if auto_tools_called else 0
                            else:
                                index = 0

                            if self._should_check_for_unstreamed_tool_arg_tokens(delta_message, output) and tool_parser:
                                latest_delta_len = 0
                                if (
                                    isinstance(
                                        delta_message.tool_calls[0].function,
                                        DeltaFunctionCall,
                                    )
                                ) and isinstance(delta_message.tool_calls[0].function.arguments, str):
                                    latest_delta_len = len(delta_message.tool_calls[0].function.arguments)

                                # get the expected call based on partial JSON
                                # parsing which "autocompletes" the JSON.
                                # Tool parsers (e.g. Qwen3Coder) store
                                # arguments as a JSON string in
                                # prev_tool_call_arr. Calling json.dumps()
                                # on an already-serialized string would
                                # double-serialize it (e.g. '{"k":1}' becomes
                                # '"{\\"k\\":1}"'), which then causes the
                                # replace() below to fail and append the
                                # entire double-serialized string as a
                                # spurious final delta.
                                args = tool_parser.prev_tool_call_arr[index].get("arguments", {})
                                if isinstance(args, str):
                                    expected_call = args
                                else:
                                    expected_call = json.dumps(args, ensure_ascii=False)

                                # get what we've streamed so far for arguments
                                # for the current tool
                                actual_call = tool_parser.streamed_args_for_tool[index]
                                if latest_delta_len > 0:
                                    actual_call = actual_call[:-latest_delta_len]

                                # check to see if there's anything left to stream
                                remaining_call = expected_call.replace(actual_call, "", 1)
                                # set that as a delta message
                                delta_message = DeltaMessage(
                                    tool_calls=[
                                        DeltaToolCall(
                                            index=index,
                                            function=DeltaFunctionCall(arguments=remaining_call).model_dump(
                                                exclude_none=True
                                            ),
                                        )
                                    ]
                                )

                            # Send the finish response for each request.n only once
                            # In OpenAI's API, when a tool is called, the
                            # finish_reason is:
                            # "tool_calls" for "auto" or "required" tool calls,
                            # and "stop" for named tool calls.
                            if (
                                auto_tools_called
                                or (tools_streamed[i] and not tool_choice_function_name)
                                or (self.use_harmony and harmony_tools_streamed[i])
                            ):
                                finish_reason_ = "tool_calls"
                            else:
                                finish_reason_ = output.finish_reason if output.finish_reason else "stop"
                            # Only emit finish_reason on the last modality to
                            # comply with OpenAI streaming spec: exactly one
                            # chunk per choice carries finish_reason="stop".
                            modality_finished[i].add("text")
                            if modality_seen[i] < set(request.modalities) or not all(
                                m in modality_finished[i] for m in modality_seen[i]
                            ):
                                finish_reason_ = None
                            else:
                                stop_reason_emitted[i] = True
                            choice_data = ChatCompletionResponseStreamChoice(
                                index=i,
                                delta=delta_message,
                                logprobs=logprobs,
                                finish_reason=finish_reason_,
                                stop_reason=output.stop_reason,
                                token_ids=(as_list(output.token_ids) if request.return_token_ids else None),
                            )

                            finish_reason_sent[i] = True

                        choice_data = maybe_filter_parallel_tool_calls(choice_data, request)
                        chunk = OmniChatCompletionStreamResponse(
                            id=request_id,
                            object=chunk_object_type,
                            created=created_time,
                            choices=[choice_data],
                            model=model_name,
                            modality=final_output_type,
                            metrics=self._filter_stage_metrics_detail(omni_res.metrics, request),
                        )

                        # handle usage stats if requested & if continuous
                        if include_continuous_usage:
                            completion_tokens = previous_num_tokens[i]
                            chunk.usage = UsageInfo(
                                prompt_tokens=num_prompt_tokens,
                                completion_tokens=completion_tokens,
                                total_tokens=num_prompt_tokens + completion_tokens,
                            )

                        data = chunk.model_dump_json(exclude_unset=True)
                        yield f"data: {data}\n\n"

                elif final_output_type == "audio":
                    # Observe audio_ttfp_s on first audio packet for this request_id
                    # (once-per-request guard via first_audio_ts). The same hook
                    # also captures (stage, replica) for the streaming-continuity
                    # emit at request finalize. self.engine_client.request_states
                    # is keyed by the internal UUID-suffixed id (set by
                    # AsyncOmni.generate via _get_unique_request_id), but the
                    # `request_id` we have here is the external (user-visible) id,
                    # so resolve via the external_request_id field.
                    req_state = next(
                        (s for s in self.engine_client.request_states.values() if s.external_request_id == request_id),
                        None,
                    )
                    if req_state is not None and req_state_audio_ref is None:
                        req_state_audio_ref = req_state
                    now_ts = time.time()
                    if req_state is not None and req_state.first_audio_ts is None:
                        req_state.first_audio_ts = now_ts
                        replica_id = getattr(omni_res, "replica_id", None)
                        if replica_id is None:
                            stage_pools = getattr(self.engine_client.engine, "stage_pools", None)
                            # Fallback for older outputs. The orchestrator binds
                            # requests by internal id, but can release that
                            # binding before the API sees the final audio chunk,
                            # so prefer the OutputMessage replica copied onto
                            # OmniRequestOutput above.
                            replica_id = (
                                stage_pools[omni_res.stage_id].get_bound_replica_id(req_state.request_id)
                                if stage_pools is not None and 0 <= omni_res.stage_id < len(stage_pools)
                                else None
                            )
                        req_state.audio_emit_stage_id = omni_res.stage_id
                        req_state.audio_emit_replica_id = replica_id
                        observe_audio_first_packet(
                            self.engine_client.mod_metrics,
                            stage_id=omni_res.stage_id,
                            replica_id=replica_id,
                            arrival_ts=req_state.request_arrival_ts,
                            now_ts=now_ts,
                        )

                    role = self.get_chat_request_role(request)
                    choices_data = self._create_audio_choice(omni_res, role, request, stream=True)
                    if isinstance(choices_data, ErrorResponse):
                        logger.error(
                            "Skipping audio chunk for request %s: %s",
                            request_id,
                            choices_data.error.message,
                        )
                        continue
                    # Only emit finish_reason on the last modality to
                    # comply with OpenAI streaming spec.
                    for choice in choices_data:
                        if choice.finish_reason is not None:
                            modality_finished[choice.index].add("audio")
                        if modality_seen[choice.index] < set(request.modalities) or not all(
                            m in modality_finished[choice.index] for m in modality_seen[choice.index]
                        ):
                            choice.finish_reason = None
                        else:
                            stop_reason_emitted[choice.index] = True
                    # Record per-chunk PCM byte count + arrival timestamp for
                    # audio_underrun_s / audio_continuity_ok_total at finalize.
                    if req_state is not None and req_state.request_arrival_ts > 0:
                        chunk_bytes = audio_chunk_pcm_bytes(omni_res)
                        if chunk_bytes > 0:
                            req_state.audio_chunk_arrivals_s.append(max(now_ts - req_state.request_arrival_ts, 0.0))
                            req_state.audio_chunk_bytes.append(chunk_bytes)
                            if req_state.audio_sample_rate is None:
                                req_state.audio_sample_rate = audio_chunk_sample_rate(omni_res)
                    chunk = OmniChatCompletionStreamResponse(
                        id=request_id,
                        object=chunk_object_type,
                        created=created_time,
                        choices=choices_data,
                        model=model_name,
                        modality=final_output_type,
                        metrics=self._filter_stage_metrics_detail(omni_res.metrics, request),
                    )
                    chunk.usage = UsageInfo(
                        prompt_tokens=num_prompt_tokens,
                        completion_tokens=0,
                        total_tokens=num_prompt_tokens,
                    )
                    data = chunk.model_dump_json(exclude_unset=True)
                    yield f"data: {data}\n\n"

                elif final_output_type == "image":
                    role = self.get_chat_request_role(request)
                    choices_data = []
                    for choice in self._create_image_choice(omni_res, role, request, stream=True):
                        delta = DeltaMessage.model_construct(role=role)
                        object.__setattr__(delta, "content", choice.message.content)
                        if hasattr(delta, "__pydantic_fields_set__"):
                            delta.__pydantic_fields_set__.add("content")
                        stream_choice = ChatCompletionResponseStreamChoice(
                            index=choice.index,
                            delta=delta,
                            logprobs=None,
                            finish_reason=choice.finish_reason,
                            stop_reason=choice.stop_reason,
                        )
                        if stream_choice.finish_reason is not None:
                            modality_finished[stream_choice.index].add("image")
                        if modality_seen[stream_choice.index] < set(request.modalities) or not all(
                            m in modality_finished[stream_choice.index] for m in modality_seen[stream_choice.index]
                        ):
                            stream_choice.finish_reason = None
                        else:
                            stop_reason_emitted[stream_choice.index] = True
                        choices_data.append(stream_choice)
                    chunk = OmniChatCompletionStreamResponse(
                        id=request_id,
                        object=chunk_object_type,
                        created=created_time,
                        choices=choices_data,
                        model=model_name,
                        modality=final_output_type,
                        metrics=self._filter_stage_metrics_detail(omni_res.metrics, request),
                    )
                    # NOTE: Currently usage is only set the text stages to align with the behavior
                    # of the full generator. TODO (Alex): Add support for usage on all stages for
                    # both streaming and non-streaming.
                    data = chunk.model_dump_json(exclude_unset=True)
                    yield f"data: {data}\n\n"

                else:
                    logger.warning(f"Unsupported streaming final output type: {final_output_type}")
                    continue

            # Fallback: if a choice had modalities finish but no finish_reason="stop"
            # was emitted (e.g. request.modalities included a modality the engine
            # never produced), emit a final stop chunk for that choice.
            for i in range(num_choices):
                if modality_finished[i] and not stop_reason_emitted[i]:
                    stop_choice = ChatCompletionResponseStreamChoice(
                        index=i,
                        delta=DeltaMessage(),
                        finish_reason="stop",
                    )
                    stop_chunk = OmniChatCompletionStreamResponse(
                        id=request_id,
                        object=chunk_object_type,
                        created=created_time,
                        choices=[stop_choice],
                        model=model_name,
                    )
                    data = stop_chunk.model_dump_json(exclude_unset=True)
                    yield f"data: {data}\n\n"
            # Emit audio_underrun_s + audio_continuity_ok_total once per
            # request after the audio chunk stream is exhausted. The
            # captured reference (req_state_audio_ref) outlives the inner
            # _log_summary_and_cleanup that pops request_states[request_id].
            if (
                req_state_audio_ref is not None
                and req_state_audio_ref.audio_chunk_arrivals_s
                and req_state_audio_ref.audio_emit_replica_id is not None
            ):
                observe_audio_streaming_finalize(
                    self.engine_client.mod_metrics,
                    stage_id=req_state_audio_ref.audio_emit_stage_id or 0,
                    replica_id=req_state_audio_ref.audio_emit_replica_id,
                    chunk_arrival_times_s=req_state_audio_ref.audio_chunk_arrivals_s,
                    chunk_bytes=req_state_audio_ref.audio_chunk_bytes,
                    sample_rate=(req_state_audio_ref.audio_sample_rate or _metric_defs.DEFAULT_AUDIO_SAMPLE_RATE),
                )

            # once the final token is handled, if stream_options.include_usage
            # is sent, send the usage
            if include_usage:
                completion_tokens = sum(previous_num_tokens)
                final_usage = UsageInfo(
                    prompt_tokens=num_prompt_tokens,
                    completion_tokens=completion_tokens,
                    total_tokens=num_prompt_tokens + completion_tokens,
                )
                final_usage.prompt_tokens_details = _make_prompt_tokens_details(
                    self.enable_prompt_tokens_details,
                    num_cached_tokens,
                    num_cache_creation_tokens,
                    mm_token_counts,
                )

                final_usage_chunk = OmniChatCompletionStreamResponse(
                    id=request_id,
                    object=chunk_object_type,
                    created=created_time,
                    choices=[],
                    model=model_name,
                    usage=final_usage,
                    metrics=self._filter_stage_metrics_detail(last_metrics, request),
                )
                final_usage_data = final_usage_chunk.model_dump_json(exclude_unset=True, exclude_none=True)
                yield f"data: {final_usage_data}\n\n"

            # report to FastAPI middleware aggregate usage across all choices
            num_completion_tokens = sum(previous_num_tokens)
            request_metadata.final_usage_info = UsageInfo(
                prompt_tokens=num_prompt_tokens,
                completion_tokens=num_completion_tokens,
                total_tokens=num_prompt_tokens + num_completion_tokens,
            )

            # Log complete streaming response if output logging is enabled
            if self.enable_log_outputs and self.request_logger:
                # Log the complete response for each choice
                for i in range(num_choices):
                    full_text = (
                        previous_texts[i]
                        if previous_texts and i < len(previous_texts)
                        else f"<streaming_complete: {previous_num_tokens[i]} tokens>"
                    )
                    self.request_logger.log_outputs(
                        request_id=request_id,
                        outputs=full_text,
                        output_token_ids=None,  # Consider also logging all token IDs
                        finish_reason="streaming_complete",
                        is_streaming=True,
                        delta=False,
                    )

        except EngineDeadError as e:
            logger.error(
                "EngineDeadError during streaming for request %s: %s",
                request_id,
                e,
            )
            data = self.create_streaming_error_response(e)
            yield f"data: {data}\n\n"
            # Actively signal shutdown instead of waiting for the watchdog
            # (5s polling interval).
            if raw_request is not None:
                terminate_if_errored(
                    server=raw_request.app.state.server,
                    engine=self.engine_client,
                )
        except Exception as e:
            logger.exception("Error in chat completion stream generator.")
            data = self.create_streaming_error_response(e)
            yield f"data: {data}\n\n"
        # Send the final done message after all response.n are finished
        yield "data: [DONE]\n\n"

    async def chat_completion_full_generator(
        self,
        request: ChatCompletionRequest,
        result_generator: AsyncIterator[RequestOutput],
        request_id: str,
        model_name: str,
        conversation: list[ConversationMessage],
        tokenizer: TokenizerLike,
        request_metadata: RequestResponseMetadata,
        reasoning_parser: ReasoningParser | None = None,
        mm_token_counts: dict[str, int] | None = None,
    ) -> ErrorResponse | OmniChatCompletionResponse:
        created_time = int(time.time())
        final_res: RequestOutput | None = None

        final_outputs: list[OmniRequestOutput] = []
        try:
            async for res in result_generator:
                final_outputs.append(res)
        except asyncio.CancelledError:
            return self.create_error_response("Client disconnected")
        except ValueError as e:
            return self.create_error_response(e)

        assert final_outputs is not None

        choices: list[ChatCompletionResponseChoice] = []

        usage = UsageInfo(prompt_tokens=0, completion_tokens=0, total_tokens=0)
        role = self.get_chat_request_role(request)
        prompt_logprobs = None
        prompt_token_ids = None
        kv_transfer_params = None
        response_metrics: dict[str, Any] | None = None

        # Build requested modalities set for filtering
        requested_modalities = (
            set(request.modalities) if hasattr(request, "modalities") and request.modalities else None
        )

        for omni_outputs in final_outputs:
            choices_data = []
            if not omni_outputs.finished:
                continue

            # Filter outputs based on requested modalites
            if requested_modalities is not None and omni_outputs.final_output_type not in requested_modalities:
                logger.warning(f"final output type: {omni_outputs.final_output_type} is not needed by the request")
                continue

            if omni_outputs.final_output_type == "text":
                if omni_outputs.stage_id is not None or omni_outputs.outputs:
                    (
                        choices_data,
                        usage,
                        prompt_logprobs,
                        prompt_token_ids,
                        kv_transfer_params,
                    ) = self._create_text_choice(
                        request,
                        omni_outputs,
                        tokenizer,
                        conversation,
                        role,
                        reasoning_parser,
                        mm_token_counts=mm_token_counts,
                    )
                    final_res = omni_outputs
                else:
                    # Diffusion pipeline text output (e.g. single-stage
                    # img2text / text2text) — no AR request_output, so build
                    # a simple text choice from diffusion multimodal output.
                    text_body = self._get_diffusion_text_output(omni_outputs)
                    message = ChatMessage(role=role, content=text_body)
                    choices_data = [
                        ChatCompletionResponseChoice(
                            index=0,
                            message=message,
                            logprobs=None,
                            finish_reason="stop",
                            stop_reason=None,
                        )
                    ]
            elif omni_outputs.final_output_type == "audio":
                choices_data = self._create_audio_choice(omni_outputs, role, request, stream=False)
                if isinstance(choices_data, ErrorResponse):
                    return choices_data
            elif omni_outputs.final_output_type == "image":
                choices_data = self._create_image_choice(omni_outputs, role, request, stream=False)
            else:
                logger.warning(f"Unsupported final output type: {omni_outputs.final_output_type}")
                continue
            if omni_outputs.metrics:
                response_metrics = dict(omni_outputs.metrics)
            if omni_outputs.final_output_type == "image":
                # Expose diffusion profiler metrics on the top-level response for benchmarks / clients.
                if response_metrics is None:
                    response_metrics = {}
                response_metrics.setdefault("stage_durations", omni_outputs.stage_durations or {})
                response_metrics.setdefault("peak_memory_mb", float(omni_outputs.peak_memory_mb or 0.0))
                extra = self._get_diffusion_extra_output_params(omni_outputs)
                if extra:
                    response_metrics.update(extra)
            choices.extend(choices_data)

        response_metrics = self._filter_stage_metrics_detail(response_metrics, request)

        # Compute prompt_text for non-streaming response (upstream #42052)
        prompt_text = (
            getattr(final_res, "prompt", None)
            if final_res is not None and getattr(request, "return_prompt_text", None)
            else None
        )

        response = OmniChatCompletionResponse(
            id=request_id,
            created=created_time,
            model=model_name,
            choices=choices,
            usage=usage,
            prompt_logprobs=prompt_logprobs,
            prompt_token_ids=prompt_token_ids,
            prompt_text=prompt_text,
            kv_transfer_params=kv_transfer_params,
            metrics=response_metrics,
        )

        # Log complete response if output logging is enabled
        if self.enable_log_outputs and self.request_logger:
            for choice in choices:
                output_text = ""
                if choice.message.content:
                    output_text = choice.message.content
                elif choice.message.tool_calls:
                    # For tool calls, log the function name and arguments
                    tool_call_descriptions = []
                    for tc in choice.message.tool_calls:
                        if hasattr(tc.function, "name") and hasattr(tc.function, "arguments"):
                            tool_call_descriptions.append(f"{tc.function.name}({tc.function.arguments})")
                    tool_calls_str = ", ".join(tool_call_descriptions)
                    output_text = f"[tool_calls: {tool_calls_str}]"

                if output_text:
                    # Get the corresponding output token IDs
                    output_token_ids = None
                    if choice.index < len(final_res.outputs):
                        output_token_ids = final_res.outputs[choice.index].token_ids

                    self.request_logger.log_outputs(
                        request_id=request_id,
                        outputs=output_text,
                        output_token_ids=output_token_ids,
                        finish_reason=choice.finish_reason,
                        is_streaming=False,
                        delta=False,
                    )

        return response

    def _create_text_choice(
        self,
        request: ChatCompletionRequest,
        omni_outputs: OmniRequestOutput,
        tokenizer: TokenizerLike,
        conversation: list[ConversationMessage],
        role: str,
        reasoning_parser: ReasoningParser | None = None,
        mm_token_counts: dict[str, int] | None = None,
    ):
        final_res = omni_outputs
        if self.tool_call_id_type == "kimi_k2":
            history_tool_call_cnt = get_history_tool_calls_cnt(conversation)
        else:
            history_tool_call_cnt = 0

        choices: list[ChatCompletionResponseChoice] = []

        for output in final_res.outputs:
            token_ids = output.token_ids
            out_logprobs = output.logprobs
            tool_call_info = None

            if request.logprobs and request.top_logprobs is not None:
                assert out_logprobs is not None, "Did not output logprobs"
                logprobs = self._create_chat_logprobs(
                    token_ids=token_ids,
                    top_logprobs=out_logprobs,
                    num_output_top_logprobs=request.top_logprobs,
                    tokenizer=tokenizer,
                    return_as_token_id=request.return_tokens_as_token_ids,
                )
            else:
                logprobs = None

            if self.use_harmony:
                # Use upstream's HarmonyParser (inherited as self.parser_cls for
                # gpt_oss/harmony models) instead of vendoring a copy of the
                # pre-refactor parse_chat_output. We take only reasoning+content;
                # tool calls are extracted by omni's own path below.
                if self.parser_cls is not None:
                    parser = self.parser_cls(tokenizer, request.tools)
                    reasoning, content, _ = parser.parse(
                        "",
                        request,
                        model_output_token_ids=token_ids,
                    )
                else:
                    reasoning, content = None, None
                if not request.include_reasoning:
                    reasoning = None

                if self.parser_cls is not None and self.parser_cls.tool_parser_cls is not None:
                    tool_parser = self.parser_cls.tool_parser_cls(tokenizer, request.tools)
                    # NOTE: We use token_ids for openai tool parser
                    tool_call_info = tool_parser.extract_tool_calls(
                        "",
                        request=request,
                        token_ids=token_ids,  # type: ignore
                    )
                    content = tool_call_info.content
                    message = ChatMessage(
                        role=role,
                        reasoning=reasoning,
                        content=content,
                        tool_calls=tool_call_info.tool_calls,
                    )
                else:
                    message = ChatMessage(
                        role=role,
                        reasoning=reasoning,
                        content=content,
                    )

                choice_data = ChatCompletionResponseChoice(
                    index=output.index,
                    message=message,
                    logprobs=logprobs,
                    finish_reason=(
                        "tool_calls"
                        if (tool_call_info is not None and tool_call_info.tools_called)
                        else (output.finish_reason if output.finish_reason else "stop")
                    ),
                    stop_reason=output.stop_reason,
                )
                choices.append(choice_data)
                continue

            if reasoning_parser:
                # If the reasoning parser is enabled,
                # tool calls are extracted exclusively from the content.
                reasoning, content = reasoning_parser.extract_reasoning(output.text, request=request)
                if not request.include_reasoning:
                    reasoning = None
            else:
                reasoning = None
                content = output.text

            auto_tools_called = False
            # if auto tools are not enabled, and a named tool choice using
            #   outlines is not being used
            if (not self.enable_auto_tools or self.parser_cls is None or self.parser_cls.tool_parser_cls is None) and (
                not isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam)
                and request.tool_choice != "required"
            ):
                message = ChatMessage(role=role, reasoning=reasoning, content=content)

            # if the request uses tools and specified a tool choice
            elif request.tool_choice and type(request.tool_choice) is ChatCompletionNamedToolChoiceParam:
                tool_call_class = MistralToolCall if isinstance(tokenizer, MistralTokenizer) else ToolCall
                message = ChatMessage(
                    role=role,
                    reasoning=reasoning,
                    content="",
                    tool_calls=[
                        tool_call_class(
                            function=FunctionCall(
                                name=request.tool_choice.function.name,
                                arguments=content,
                            )
                        )
                    ],
                )

            elif request.tool_choice and request.tool_choice == "required":
                tool_call_class = MistralToolCall if isinstance(tokenizer, MistralTokenizer) else ToolCall

                # the fields of FunctionDefinition are a superset of the
                # tool call outputs and can be used for parsing
                assert content is not None
                tool_calls = TypeAdapter(list[FunctionDefinition]).validate_json(content)
                tool_call_ids = []
                for tool_call in tool_calls:
                    tool_call_ids.append(
                        make_tool_call_id(
                            id_type=self.tool_call_id_type,
                            func_name=tool_call.name,
                            idx=history_tool_call_cnt,
                        )
                    )
                    history_tool_call_cnt += 1
                message = ChatMessage(
                    role=role,
                    content="",
                    tool_calls=[
                        tool_call_class(
                            id=tool_call_ids[i],
                            function=FunctionCall(
                                name=tool_call.name,
                                arguments=json.dumps(tool_call.parameters, ensure_ascii=False),
                            ),
                        )
                        for i, tool_call in enumerate(tool_calls)
                    ],
                    reasoning=reasoning,
                )

            # if the request doesn't use tool choice
            # OR specifies to not use a tool
            elif not request.tool_choice or request.tool_choice == "none":
                message = ChatMessage(role=role, reasoning=reasoning, content=content)

            # handle when there are tools and tool choice is auto
            elif (
                request.tools
                and (request.tool_choice == "auto" or request.tool_choice is None)
                and self.enable_auto_tools
                and self.parser_cls is not None
                and self.parser_cls.tool_parser_cls is not None
            ):
                try:
                    tool_parser = self.parser_cls.tool_parser_cls(tokenizer, request.tools)
                except RuntimeError as e:
                    logger.exception("Error in tool parser creation.")
                    return self.create_error_response(e)

                tool_call_info = tool_parser.extract_tool_calls(content if content is not None else "", request=request)
                # In the OpenAI API the finish_reason is "tools_called"
                # if the tool choice is auto and the model produced a tool
                # call. The same is not true for named function calls
                auto_tools_called = tool_call_info.tools_called
                if tool_call_info.tools_called:
                    message = ChatMessage(
                        role=role,
                        reasoning=reasoning,
                        content=tool_call_info.content,
                        tool_calls=tool_call_info.tool_calls,
                    )

                else:
                    # FOR NOW make it a chat message; we will have to detect
                    # the type to make it later.
                    ret_content = content

                    # try to use content return from tool parser first,
                    # tool parser may do some modify for the content.
                    if tool_call_info.content and len(tool_call_info.content) > 0:
                        ret_content = tool_call_info.content
                    message = ChatMessage(
                        role=role,
                        reasoning=reasoning,
                        content=ret_content,
                    )

            # undetermined case that is still important to handle
            else:
                logger.error(
                    "Error in chat_completion_full_generator - cannot determine if tools should be extracted. "
                    "Returning a standard chat completion."
                )
                message = ChatMessage(role=role, reasoning=reasoning, content=content)

            choice_data = ChatCompletionResponseChoice(
                index=output.index,
                message=message,
                logprobs=logprobs,
                finish_reason=(
                    "tool_calls" if auto_tools_called else output.finish_reason if output.finish_reason else "stop"
                ),
                stop_reason=output.stop_reason,
                token_ids=(as_list(output.token_ids) if request.return_token_ids else None),
            )
            choices.append(choice_data)

        if request.echo:
            last_msg_content: str | list[dict[str, str]] = ""
            if conversation and "content" in conversation[-1] and conversation[-1].get("role") == role:
                last_msg_content = conversation[-1]["content"] or ""
            if isinstance(last_msg_content, list):
                last_msg_content = "\n".join(msg["text"] for msg in last_msg_content)

            for choice in choices:
                full_message = last_msg_content + (choice.message.content or "")
                choice.message.content = full_message

        assert final_res.prompt_token_ids is not None
        num_prompt_tokens = len(final_res.prompt_token_ids)
        if final_res.encoder_prompt_token_ids is not None:
            num_prompt_tokens += len(final_res.encoder_prompt_token_ids)
        num_generated_tokens = sum(len(output.token_ids) for output in final_res.outputs)
        usage = UsageInfo(
            prompt_tokens=num_prompt_tokens,
            completion_tokens=num_generated_tokens,
            total_tokens=num_prompt_tokens + num_generated_tokens,
        )
        usage.prompt_tokens_details = _make_prompt_tokens_details(
            self.enable_prompt_tokens_details,
            final_res.num_cached_tokens,
            final_res.num_cache_creation_tokens,
            mm_token_counts,
        )

        prompt_logprobs = clamp_prompt_logprobs(final_res.prompt_logprobs)
        prompt_token_ids = final_res.prompt_token_ids if request.return_token_ids else None
        kv_transfer_params = final_res.kv_transfer_params

        return choices, usage, prompt_logprobs, prompt_token_ids, kv_transfer_params

    def _create_audio_choice(
        self, omni_outputs: OmniRequestOutput, role: str, request: ChatCompletionRequest, stream: bool = False
    ) -> list[ChatCompletionResponseChoice] | list[ChatCompletionResponseStreamChoice] | ErrorResponse:
        choices: list[ChatCompletionResponseChoice] = []
        final_res = omni_outputs
        # OMNI: Access multimodal_output from CompletionOutput (outputs[0]), not from RequestOutput
        # Reference: examples/offline_inference/qwen3_omni/end2end.py line 421
        # The attribute is attached dynamically when stage audio arrives; fall
        # back to the no-audio error response instead of an AttributeError 500
        # when the pipeline produced no audio for this request.
        mm_output = getattr(final_res.outputs[0], "multimodal_output", None) or {}
        audio_data = mm_output.get("audio")
        if isinstance(audio_data, list):
            if not audio_data:
                audio_tensor = None
            elif stream:
                audio_tensor = audio_data[-1]
            else:
                audio_tensor = torch.cat(audio_data, dim=-1)
        else:
            audio_tensor = audio_data
        if audio_tensor is None:
            if not stream:
                return self._create_error_response("Audio generation completed but no audio was produced.")
            # A streamed message can legitimately carry no waveform: the talker
            # may emit an empty speech segment, and Stage 2 then reports the
            # chunk with metadata only. Emit an empty audio delta so the
            # caller's finish_reason bookkeeping still runs and the stream
            # terminates normally instead of raising mid-generation.
            return [
                ChatCompletionResponseStreamChoice(
                    index=output.index,
                    delta=DeltaMessage(role=role, content=""),
                    logprobs=None,
                    finish_reason=output.finish_reason,
                    stop_reason=output.stop_reason,
                    token_ids=(as_list(output.token_ids) if request.return_token_ids else None),
                )
                for output in final_res.outputs
            ]
        audio_tensor = audio_tensor.detach().cpu().float().numpy()

        # Ensure audio is 1D (flatten if needed)
        if audio_tensor.ndim > 1:
            audio_tensor = audio_tensor.flatten()

        # Prefer the talker-reported sample rate when present. Qwen3-Omni
        # omits "sr" and runs at 24kHz; Ming-flash-omni surfaces a 44.1kHz
        # AudioVAE rate via multimodal_output["sr"].
        sr_raw = mm_output.get("sr")
        if isinstance(sr_raw, (list, tuple)):
            sr_raw = next((item for item in sr_raw if item is not None), None)
        if sr_raw is None:
            sample_rate = 24000
        elif hasattr(sr_raw, "item"):
            sample_rate = int(sr_raw.item())
        else:
            sample_rate = int(sr_raw)

        audio_format = self._resolve_audio_format(request)
        if isinstance(audio_format, ErrorResponse):
            return audio_format

        audio_obj = CreateAudio(
            audio_tensor=audio_tensor,
            sample_rate=sample_rate,
            response_format=audio_format,
            speed=1.0,
            base64_encode=True,
        )

        audio_response: AudioResponse = self.create_audio(audio_obj)
        audio_base64 = audio_response.audio_data

        # Generate unique ID for the audio
        audio_id = f"audio-{uuid.uuid4().hex[:16]}"

        # Set expiration time (e.g., 24 hours from now) as Unix timestamp
        expires_at = int((datetime.now(timezone.utc) + timedelta(hours=24)).timestamp())

        # Create OpenAIChatCompletionAudio object with all required fields
        audio_obj = OpenAIChatCompletionAudio(
            id=audio_id,
            data=audio_base64,
            expires_at=expires_at,
            transcript="",  # Empty transcript if not available
        )

        for output in final_res.outputs:
            if stream:
                choice_data = ChatCompletionResponseStreamChoice(
                    index=output.index,
                    delta=DeltaMessage(role=role, content=audio_base64),
                    logprobs=None,
                    finish_reason=output.finish_reason,
                    stop_reason=output.stop_reason,
                    token_ids=(as_list(output.token_ids) if request.return_token_ids else None),
                )
            else:
                choice_data = ChatCompletionResponseChoice(
                    index=output.index,
                    message=ChatMessage(role=role, audio=audio_obj),
                    logprobs=None,
                    finish_reason="stop",
                    stop_reason=output.stop_reason,
                )
            choices.append(choice_data)
        return choices

    def _create_image_choice(
        self, omni_outputs: OmniRequestOutput, role: str, request: ChatCompletionRequest, stream: bool = False
    ):
        """Create chat completion response choices for image output.

        Converts image tensor or PIL Image output from diffusion models
        into base64-encoded image data for API response.

        Args:
            omni_outputs: Output containing image data from diffusion stage
            role: The role for the response message (e.g., "assistant")

        Returns:
            List of ChatCompletionResponseChoice with image content
        """
        from PIL import Image

        choices: list[ChatCompletionResponseChoice] = []

        # Handle profiling data
        stage_durations = omni_outputs.stage_durations
        peak_memory_mb = omni_outputs.peak_memory_mb

        # Handle different image output formats
        images = []

        # First check omni_outputs.images directly (for diffusion mode via from_diffusion)
        if omni_outputs.images:
            images = omni_outputs.images
        # Fall back to completion outputs for pipeline mode (multimodal_output
        # is attached to CompletionOutput by AR stages).
        elif omni_outputs.outputs:
            completion_output = omni_outputs.outputs[0]
            final_res = omni_outputs
            if hasattr(completion_output, "multimodal_output") and completion_output.multimodal_output:
                image_data = completion_output.multimodal_output.get("image")
                if image_data is not None:
                    if isinstance(image_data, Image.Image):
                        images.append(image_data)
                    elif hasattr(image_data, "cpu"):  # Tensor
                        import numpy as np

                        # Convert tensor to PIL Image
                        img_array = image_data.float().detach().cpu().numpy()
                        # Handle different tensor formats (CHW -> HWC)
                        if img_array.ndim == 3 and img_array.shape[0] in [1, 3, 4]:
                            img_array = np.transpose(img_array, (1, 2, 0))
                        # Normalize to 0-255
                        if img_array.max() <= 1.0:
                            img_array = (img_array * 255).astype(np.uint8)
                        else:
                            img_array = img_array.astype(np.uint8)
                        # Handle grayscale
                        if img_array.ndim == 2:
                            images.append(Image.fromarray(img_array, mode="L"))
                        elif img_array.shape[-1] == 1:
                            images.append(Image.fromarray(img_array.squeeze(-1), mode="L"))
                        elif img_array.shape[-1] == 3:
                            images.append(Image.fromarray(img_array, mode="RGB"))
                        elif img_array.shape[-1] == 4:
                            images.append(Image.fromarray(img_array, mode="RGBA"))
            elif hasattr(final_res, "images") and final_res.images:
                images = final_res.images

        # Convert images to base64
        image_contents = []
        for img in images:
            with BytesIO() as buffer:
                img.save(buffer, format="PNG")
                img_bytes = buffer.getvalue()
            img_base64 = base64.b64encode(img_bytes).decode("utf-8")
            image_contents.append(
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/png;base64,{img_base64}",
                    },
                    "stage_durations": stage_durations,
                    "peak_memory_mb": peak_memory_mb,
                }
            )

        # Create message content
        if len(image_contents) == 1:
            content = image_contents
        elif len(image_contents) > 1:
            content = image_contents
        else:
            content = [{"type": "text", "text": "Image generation completed but no images were produced."}]

        # Create response choice
        # Use model_construct to bypass validation for multimodal content
        # (ChatMessage.content only accepts str, but we need list for images)
        # Then use object.__setattr__ to directly set the field, bypassing Pydantic's type checking
        import warnings as warnings_module

        with warnings_module.catch_warnings():
            warnings_module.filterwarnings("ignore", category=UserWarning, module="pydantic")
            message = ChatMessage.model_construct(role=role)
            object.__setattr__(message, "content", content)
            # Mark content as set in fields_set to ensure proper serialization
            if hasattr(message, "__pydantic_fields_set__"):
                message.__pydantic_fields_set__.add("content")
        choice_data = ChatCompletionResponseChoice(
            index=0,
            message=message,
            logprobs=None,
            finish_reason="stop",
            stop_reason=None,
        )
        choices.append(choice_data)

        return choices

    # ==================== Diffusion Mode Methods ====================
    def _build_multistage_generation_inputs(
        self,
        *,
        engine: AsyncOmni,
        prompt: str,
        extra_body: dict[str, Any],
        reference_images: list[Image.Image],
        gen_params: OmniDiffusionSamplingParams,
        tokenizer: Any = None,
    ) -> tuple[OmniTextPrompt, list[Any]]:
        """Build the shared multistage generation prompt and stage params."""
        stage_configs = getattr(engine, "stage_configs", None) or []
        default_params_list = get_default_sampling_params_list(engine)

        height = gen_params.height
        width = gen_params.width
        seed = gen_params.seed
        generator_device = gen_params.generator_device
        num_outputs_per_prompt = gen_params.num_outputs_per_prompt
        num_inference_steps = extra_body.get("num_inference_steps")
        quality = extra_body.get("quality")
        guidance_scale = extra_body.get("guidance_scale")
        true_cfg_scale = extra_body.get("true_cfg_scale") or extra_body.get("cfg_scale")
        negative_prompt = extra_body.get("negative_prompt")
        num_frames = extra_body.get("num_frames")
        guidance_scale_2 = extra_body.get("guidance_scale_2")
        lora_body = extra_body.get("lora")
        layers = extra_body.get("layers")
        resolution = extra_body.get("resolution")
        bot_task = extra_body.get("bot_task")
        use_system_prompt = extra_body.get("use_system_prompt") or extra_body.get("sys_type")
        custom_system_prompt = extra_body.get("system_prompt")

        engine_prompt_data: dict[str, Any] | None = None
        modalities = ["image"]
        if reference_images:
            if len(reference_images) == 1:
                engine_prompt_data = {"img2img": reference_images[0]}
                modalities = ["img2img"]
            else:
                engine_prompt_data = {"image": reference_images}

        prompt_token_ids: list[int] | None = None
        system_prompt_type: str | None = None
        build_kwargs: dict[str, Any] = {}
        ar_stop_token_ids: list[int] | None = None

        if bot_task is not None or use_system_prompt is not None or custom_system_prompt is not None:
            from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import (
                build_prompt,
                build_prompt_tokens,
                resolve_stop_token_ids,
            )

            build_kwargs: dict[str, Any] = {
                "task": "it2i" if reference_images else "t2i",
                "sys_type": use_system_prompt,
                "custom_system_prompt": custom_system_prompt,
                "num_images": len(reference_images) if reference_images else 1,
            }

            if bot_task is not None:
                build_kwargs["bot_task"] = bot_task
            elif "bot_task" in extra_body:
                # Explicit None from the caller is plain-mode; omitted lets
                # each task fall back to its default trigger.
                build_kwargs["bot_task"] = extra_body["bot_task"]
            if tokenizer is not None:
                # Feed segment-tokenized prompt_token_ids so AR matches HF
                # apply_chat_template byte-for-byte (engine BPE would merge
                # across template boundaries, e.g. "。\n\n" -> single id).
                result = build_prompt_tokens(prompt, tokenizer, **build_kwargs)
                prompt_token_ids = result.token_ids
                system_prompt_type = result.system_prompt_type
            else:
                prompt = build_prompt(prompt, **build_kwargs)
            if reference_images and len(reference_images) == 1:
                engine_prompt_data = {"image": reference_images[0]}
                modalities = ["image"]

            ar_task = "it2i" if reference_images else "t2i"
            # ar_image_size: None -> need_ratio=True (AR predicts ratio);
            # explicit size -> need_ratio=False (AR stops at terminator).
            ar_image_size: str | None = None
            if height is not None and width is not None:
                ar_image_size = f"{width}x{height}"
            ar_stop_token_ids = resolve_stop_token_ids(
                task=ar_task,
                bot_task=bot_task,
                tokenizer=tokenizer,
                image_size=ar_image_size,
            )

        engine_prompt: OmniTextPrompt = {"prompt": prompt}
        if prompt_token_ids is not None:
            engine_prompt["prompt_token_ids"] = prompt_token_ids
        if system_prompt_type is not None:
            engine_prompt["use_system_prompt"] = system_prompt_type
        # DiT's get_system_prompt(use_system_prompt, "image", system_prompt) reads
        # this; omitting it makes sys_type=custom yield an empty DiT prefix.
        if custom_system_prompt is not None:
            engine_prompt["system_prompt"] = custom_system_prompt
        engine_prompt["modalities"] = modalities
        if negative_prompt is not None:
            engine_prompt["negative_prompt"] = negative_prompt

        mm_processor_kwargs: dict[str, Any] = {}
        if height is not None:
            mm_processor_kwargs["target_h"] = height
        if width is not None:
            mm_processor_kwargs["target_w"] = width
        if seed is not None and engine_prompt_data is not None:
            mm_processor_kwargs["vae_generator_seed"] = int(seed)
        if mm_processor_kwargs:
            engine_prompt["mm_processor_kwargs"] = mm_processor_kwargs
        if engine_prompt_data is not None:
            engine_prompt["multi_modal_data"] = engine_prompt_data
            # Provide multi_modal_uuids so that newer vLLM versions can
            # validate multi_modal_data / multi_modal_uuids consistency.
            # Generate one uuid per image when the value is a list (multi-image inputs).
            engine_prompt["multi_modal_uuids"] = {
                k: [f"img-{k}-{i}" for i in range(len(v))] if isinstance(v, list) else [f"img-{k}-0"]
                for k, v in engine_prompt_data.items()
            }

        comprehension_idx = None
        for idx, stage in enumerate(stage_configs):
            if getattr(stage, "is_comprehension", False):
                comprehension_idx = idx
                break

        sampling_params_list = build_stage_sampling_params_list(
            stage_configs,
            default_params_list,
            diffusion_params=gen_params,
        )
        for idx, stage_cfg in enumerate(stage_configs):
            stage_type = get_stage_type(stage_cfg)
            default_stage_params = sampling_params_list[idx]

            # AR stop tokens: use stage_type=="llm" instead of comprehension_idx
            # (None for DictConfig where is_comprehension is nested in engine_args).
            if stage_type == "llm" and ar_stop_token_ids is not None:
                default_stage_params.stop_token_ids = ar_stop_token_ids

            if (
                comprehension_idx is not None
                and idx == comprehension_idx
                and seed is not None
                and hasattr(default_stage_params, "seed")
            ):
                default_stage_params.seed = seed

            # Inject target_h/w into AR stage for M-RoPE position pre-computation
            # (e.g. GLM-Image). max_tokens comes from deploy YAML.
            if comprehension_idx is not None and idx == comprehension_idx and height is not None and width is not None:
                extra_args = getattr(default_stage_params, "extra_args", None)
                if extra_args is None:
                    extra_args = {}
                    default_stage_params.extra_args = extra_args
                extra_args["target_h"] = int(height)
                extra_args["target_w"] = int(width)

            if stage_type == "diffusion":
                self._set_if_supported(
                    default_stage_params,
                    height=height,
                    width=width,
                    seed=seed,
                    generator_device=generator_device,
                    num_outputs_per_prompt=num_outputs_per_prompt,
                    quality=quality,
                    num_inference_steps=num_inference_steps,
                    guidance_scale=guidance_scale,
                    true_cfg_scale=true_cfg_scale,
                    num_frames=num_frames,
                    guidance_scale_2=guidance_scale_2,
                    layers=layers,
                    resolution=resolution,
                )
                apply_declared_extra_args(
                    default_stage_params,
                    self._get_diffusion_extra_body_params(),
                    extra_body,
                )
                if lora_body and isinstance(lora_body, dict):
                    try:
                        lora_req, lora_scale = parse_lora_request(lora_body)
                        if lora_req is not None:
                            default_stage_params.lora_request = lora_req
                            if lora_scale is not None:
                                default_stage_params.lora_scale = lora_scale
                    except Exception as e:  # pragma: no cover - safeguard
                        logger.warning("Failed to parse LoRA request: %s", e)

        return engine_prompt, sampling_params_list

    def _prepare_diffusion_image_request(
        self,
        *,
        prompt: str,
        extra_body: dict[str, Any] | None = None,
        reference_images: list[str] | None = None,
    ) -> tuple[Any, OmniTextPrompt, OmniDiffusionSamplingParams, list[Image.Image]] | ErrorResponse:
        if extra_body is None:
            extra_body = {}
        if reference_images is None:
            reference_images = []

        engine = self._diffusion_engine if self._diffusion_engine is not None else self.engine_client

        height, width = self._resolve_height_width_from_extra_body(extra_body)

        seed = extra_body.get("seed")
        generator_device = extra_body.get("generator_device")
        negative_prompt = extra_body.get("negative_prompt")
        num_outputs_per_prompt = extra_body.get("num_outputs_per_prompt", 1)
        lora_body = extra_body.get("lora")

        pil_images: list[Image.Image] = []
        for img_b64 in reference_images:
            try:
                img_bytes = base64.b64decode(img_b64)
                pil_images.append(Image.open(BytesIO(img_bytes)))
            except Exception as e:
                logger.warning("Failed to decode reference image: %s", e)

        gen_params = OmniDiffusionSamplingParams(
            height=height,
            width=width,
            num_outputs_per_prompt=num_outputs_per_prompt,
            seed=seed,
        )
        self._set_if_supported(
            gen_params,
            generator_device=generator_device,
            quality=extra_body.get("quality"),
            num_inference_steps=extra_body.get("num_inference_steps"),
            guidance_scale=extra_body.get("guidance_scale"),
            true_cfg_scale=extra_body.get("true_cfg_scale") or extra_body.get("cfg_scale"),
            num_frames=extra_body.get("num_frames"),
            guidance_scale_2=extra_body.get("guidance_scale_2"),
            layers=extra_body.get("layers"),
            resolution=extra_body.get("resolution"),
            strength=extra_body.get("strength"),
        )

        if lora_body and isinstance(lora_body, dict):
            try:
                lora_req, lora_scale = parse_lora_request(lora_body)
                if lora_req is not None:
                    gen_params.lora_request = lora_req
                    if lora_scale is not None:
                        gen_params.lora_scale = lora_scale
            except Exception as e:  # pragma: no cover - safeguard
                logger.warning("Failed to parse LoRA request: %s", e)

        gen_prompt: OmniTextPrompt = {
            "prompt": prompt,
            "negative_prompt": negative_prompt,
            "modalities": ["image"],
        }
        if pil_images:
            if len(pil_images) == 1:
                gen_prompt["multi_modal_data"] = {"image": pil_images[0]}
            else:
                od_config = getattr(engine, "od_config", None)
                supports_multimodal_inputs = getattr(od_config, "supports_multimodal_inputs", False)
                if od_config is None:
                    supports_multimodal_inputs = True
                if supports_multimodal_inputs:
                    gen_prompt["multi_modal_data"] = {"image": pil_images}
                else:
                    return self._create_error_response(
                        "Multiple input images are not supported by the current diffusion model. "
                        "For multi-image editing, start the server with Qwen-Image-Edit-2509 "
                        "and send multiple images in the user message content.",
                        status_code=400,
                    )

        return engine, gen_prompt, gen_params, pil_images

    async def generate_diffusion_images(
        self,
        *,
        prompt: str,
        extra_body: dict[str, Any] | None = None,
        reference_images: list[str] | None = None,
        request_id: str | None = None,
        arrival_time: float | None = None,
        stream: bool = False,
        model: str | None = None,
        output_format: str = "png",
        output_compression: int = 100,
        size: str = "auto",
        raw_request: Request | None = None,
    ) -> tuple[list[Image.Image], dict[str, Any], float, str | None] | ErrorResponse | AsyncIterator[str]:
        """Generate diffusion images and return raw images plus generation stats."""
        if request_id is None:
            request_id = f"chatcmpl-{uuid.uuid4().hex[:16]}"

        prepared = self._prepare_diffusion_image_request(
            prompt=prompt,
            extra_body=extra_body,
            reference_images=reference_images,
        )
        if isinstance(prepared, ErrorResponse):
            return prepared
        engine, gen_prompt, gen_params, pil_images = prepared
        if extra_body is None:
            extra_body = {}
        return_stage_metrics = self._truthy_extra_body_flag(extra_body, "return_stage_metrics")

        if isinstance(engine, AsyncOmni):
            diffusion_engine = cast(AsyncOmni, engine)
            stage_configs = getattr(diffusion_engine, "stage_configs", None) or []
            if stream and len(stage_configs) <= 1:
                return self._create_error_response(
                    "stream=true is only supported for multi-stage image editing pipelines.",
                    status_code=400,
                )
            if len(stage_configs) > 1:
                # Pull tokenizer from the comprehension (AR) stage so we can
                # build HF byte-for-byte prompt_token_ids in the helper. If
                # the engine doesn"t expose one, fall back to the legacy
                # string-prompt path (engine re-tokenizes).
                tokenizer = None
                get_tok = getattr(diffusion_engine, "get_tokenizer", None)
                if get_tok is not None:
                    try:
                        tokenizer = await get_tok()
                    except Exception as exc:
                        logger.warning("get_tokenizer failed; falling back to string prompt path: %s", exc)
                engine_prompt, sampling_params_list = self._build_multistage_generation_inputs(
                    engine=diffusion_engine,
                    prompt=prompt,
                    extra_body=extra_body,
                    reference_images=pil_images,
                    gen_params=gen_params,
                    tokenizer=tokenizer,
                )
            else:
                engine_prompt = gen_prompt
                sampling_params_list = [gen_params]

            sampling_params_list = coerce_param_message_types(sampling_params_list, stream)
            result_generator = diffusion_engine.generate(
                prompt=engine_prompt,
                sampling_params_list=sampling_params_list,
                request_id=request_id,
                arrival_time=arrival_time,
            )
            if stream:
                return self._stream_diffusion_image_chunks(
                    result_generator,
                    model=model or self._diffusion_model_name,
                    output_format=output_format,
                    output_compression=output_compression,
                    size=size,
                    return_stage_metrics=return_stage_metrics,
                    raw_request=raw_request,
                )

            result = None
            async for output in result_generator:
                result = output
            if result is None:
                return self._create_error_response("No output generated from AsyncOmni", status_code=500)
        elif stream:
            return self._create_error_response(
                "Streaming image edits require a multi-stage AsyncOmni engine.",
                status_code=400,
            )
        else:
            result = await engine.generate(
                prompt=gen_prompt,
                sampling_params=gen_params,
                request_id=request_id,
            )

        images = getattr(result, "images", [])
        stage_durations = result.stage_durations
        peak_memory_mb = result.peak_memory_mb
        cot_output = None

        req_out = result
        if req_out:
            prompt_obj = getattr(req_out, "prompt", None)
            if isinstance(prompt_obj, dict):
                extra = prompt_obj.get("extra", {})
                if isinstance(extra, dict):
                    ar_text = extra.get("ar_generated_text")
                    if isinstance(ar_text, str) and ar_text.strip():
                        cot_output = ar_text

        req_out = result
        if req_out:
            prompt_obj = getattr(req_out, "prompt", None)
            if isinstance(prompt_obj, dict):
                extra = prompt_obj.get("extra", {})
                if isinstance(extra, dict):
                    ar_text = extra.get("ar_generated_text")
                    if isinstance(ar_text, str) and ar_text.strip():
                        cot_output = ar_text

        req_out = result
        if req_out:
            prompt_obj = getattr(req_out, "prompt", None)
            if isinstance(prompt_obj, dict):
                extra = prompt_obj.get("extra", {})
                if isinstance(extra, dict):
                    ar_text = extra.get("ar_generated_text")
                    if isinstance(ar_text, str) and ar_text.strip():
                        cot_output = ar_text

        return self._flatten_diffusion_images(images), stage_durations, peak_memory_mb, cot_output

    async def _stream_diffusion_image_chunks(
        self,
        result_generator: AsyncIterator[OmniRequestOutput],
        *,
        model: str,
        output_format: str,
        output_compression: int,
        size: str,
        return_stage_metrics: bool,
        raw_request: Request | None = None,
    ) -> AsyncIterator[str]:
        """Yield image edit SSE chunks from multi-stage diffusion outputs."""
        created = int(time.time())
        emitted_image = False
        try:
            async for output in result_generator:
                final_output_type = getattr(output, "final_output_type", None)
                stage_id = getattr(output, "stage_id", None)
                metrics = getattr(output, "metrics", None) if return_stage_metrics else None
                if final_output_type == "text" and stage_id == 0:
                    for completion in output.outputs:
                        text = completion.text or ""
                        if not text:
                            continue
                        chunk = ImageEditARDeltaChunk(
                            delta=text,
                            index=completion.index,
                            created=created,
                            model=model,
                            metrics=metrics,
                        )
                        yield f"data: {chunk.model_dump_json()}\n\n"
                elif final_output_type == "image":
                    images = self._flatten_diffusion_images(getattr(output, "images", []))
                    if not images:
                        raise RuntimeError("Streaming image edit produced an empty final image output.")
                    image_data = [
                        ImageData(
                            b64_json=encode_image_base64_with_compression(
                                img,
                                format=output_format,
                                output_compression=output_compression,
                            ),
                            revised_prompt=None,
                        )
                        for img in images
                    ]
                    chunk = ImageEditImageChunk(
                        data=image_data,
                        output_format=output_format,
                        size=size,
                        created=created,
                        model=model,
                        metrics=metrics,
                    )
                    yield f"data: {chunk.model_dump_json()}\n\n"
                    emitted_image = True
            if not emitted_image:
                raise RuntimeError("Streaming image edit completed without a final image output.")
        except EngineDeadError as exc:
            logger.error("EngineDeadError during streaming image edit: %s", exc)
            data = self.create_streaming_error_response(exc)
            yield f"data: {data}\n\n"
            if raw_request is not None:
                terminate_if_errored(
                    server=raw_request.app.state.server,
                    engine=self.engine_client,
                )
            else:
                logger.warning(
                    "[OmniOpenAIServingChat] Engine dead during streaming image edit, "
                    "but cannot terminate server (no raw_request context)."
                )
        except OmniClientError as exc:
            logger.info("Client error during streaming image edit: %s", exc)
            chunk = ImageEditStreamError(
                created=created,
                model=model,
                error={
                    "message": exc.message,
                    "type": exc.error_type,
                    "code": exc.status_code,
                },
            )
            yield f"data: {chunk.model_dump_json()}\n\n"
        except Exception as exc:
            logger.exception("Streaming image edit failed: %s", exc)
            chunk = ImageEditStreamError(
                created=created,
                model=model,
                error={
                    "message": str(exc),
                    "type": "server_error",
                    "code": 500,
                },
            )
            yield f"data: {chunk.model_dump_json()}\n\n"
        yield "data: [DONE]\n\n"

    @staticmethod
    def _flatten_diffusion_images(images: Any) -> list[Image.Image]:
        flat_images: list[Image.Image] = []
        for item in images or []:
            if isinstance(item, list):
                flat_images.extend(item)
            else:
                flat_images.append(item)
        return flat_images

    async def _create_diffusion_chat_completion(
        self,
        request: ChatCompletionRequest,
        normalized_extra_args: dict[str, object],
        diffusion_request_args: dict[str, Any],
        raw_request: Request | None = None,
    ) -> ChatCompletionResponse | ErrorResponse:
        """Generate images via chat completion interface for diffusion models.

        Args:
            request: Chat completion request
            normalized_extra_args: Validated model-specific request arguments
            diffusion_request_args: Compatibility arguments plus validated common root fields
            raw_request: Raw FastAPI request object

        Returns:
            ChatCompletionResponse with generated images or ErrorResponse
        """
        try:
            request_id = f"chatcmpl-{uuid.uuid4().hex[:16]}"
            created_time = int(time.time())

            # Convert messages to dict format
            messages = []
            for msg in request.messages:
                if hasattr(msg, "model_dump"):
                    messages.append(msg.model_dump())
                elif isinstance(msg, dict):
                    messages.append(msg)
                else:
                    messages.append({"role": getattr(msg, "role", "user"), "content": getattr(msg, "content", "")})

            # Extract prompt and multimodal inputs from messages
            prompt, reference_images, reference_videos, reference_audios = self._extract_diffusion_prompt_and_media(
                messages
            )

            # Reference: text_to_image.py and text_to_video.py for supported parameters
            extra_body = diffusion_request_args

            # Parse size if provided (supports "1024x1024" format)
            height, width = self._resolve_height_width_from_extra_body(extra_body)

            # Get request parameters from extra_body.
            # Avoid hardcoded defaults here — let each pipeline's forward()
            # method apply its own model-specific default when the user does
            # not provide a value.
            num_inference_steps = extra_body.get("num_inference_steps")
            quality = extra_body.get("quality")
            guidance_scale = extra_body.get("guidance_scale")
            true_cfg_scale = extra_body.get("true_cfg_scale")
            seed = extra_body.get("seed")
            if seed is None:
                seed = getattr(request, "seed", None)
            negative_prompt = extra_body.get("negative_prompt")
            num_outputs_per_prompt = extra_body.get("num_outputs_per_prompt", 1)

            # Text-to-video parameters (ref: text_to_video.py)
            num_frames = extra_body.get("num_frames")
            guidance_scale_2 = extra_body.get("guidance_scale_2")
            lora_body = extra_body.get("lora")

            # Qwen-Image-Layered parameters
            layers = extra_body.get("layers")
            resolution = extra_body.get("resolution")
            try:
                layers = validate_layered_layers(layers)
            except ValueError as exc:
                return self._create_error_response(str(exc), status_code=400)

            logger.info(
                "Diffusion chat request %s: ref_images=%d, params=%s",
                request_id,
                len(reference_images),
                {key: value for key, value in extra_body.items() if value is not None},
            )
            _rl = getattr(self, "request_logger", None)
            if _rl:
                base_len = len(f"Diffusion chat request {request_id}: prompt=")
                raw_max = getattr(_rl, "max_log_len", None)
                cap = raw_max if isinstance(raw_max, int) else 200
                text = prompt[: max(cap - base_len, 0)]
                logger.debug("Diffusion chat request %s: prompt=%r", request_id, text)

            # Decode reference images if provided
            pil_images: list[Image.Image] = []
            for img_b64 in reference_images:
                try:
                    img_bytes = base64.b64decode(img_b64)
                    pil_images.append(Image.open(BytesIO(img_bytes)))
                except Exception as e:
                    logger.warning("Failed to decode reference image: %s", e)

            # Build generation kwargs
            gen_prompt: OmniTextPrompt = {
                "prompt": prompt,
                "negative_prompt": negative_prompt,
                "modalities": ["image"],
            }
            gen_params = OmniDiffusionSamplingParams(
                height=height,
                width=width,
                num_outputs_per_prompt=num_outputs_per_prompt,
                seed=seed,
            )

            # Only override defaults when the user explicitly provides values
            if num_inference_steps is not None:
                gen_params.num_inference_steps = num_inference_steps
            if quality is not None:
                gen_params.quality = quality
            if guidance_scale is not None:
                gen_params.guidance_scale = guidance_scale
            if true_cfg_scale is not None:
                gen_params.true_cfg_scale = true_cfg_scale
            apply_normalized_diffusion_request_extra_args(gen_params, normalized_extra_args)
            if num_frames is not None:
                gen_params.num_frames = num_frames
            if guidance_scale_2 is not None:
                gen_params.guidance_scale_2 = guidance_scale_2
            if layers is not None:
                gen_params.layers = layers
            if resolution is not None:
                gen_params.resolution = resolution

            # Parse per-request LoRA.
            if lora_body and isinstance(lora_body, dict):
                try:
                    lora_req, lora_scale = parse_lora_request(lora_body)
                    if lora_req is not None:
                        gen_params.lora_request = lora_req
                        if lora_scale is not None:
                            gen_params.lora_scale = lora_scale
                except Exception as exc:  # pragma: no cover - safeguard
                    logger.warning("Failed to parse LoRA request: %s", exc)

            # Route text modality for single-stage diffusion (img2text / text2text)
            requested_modalities = extra_body.get("modalities") or []
            is_text_request = "text" in requested_modalities

            if is_text_request:
                gen_prompt["modalities"] = ["text"]

            # Add reference image if provided (from messages content)
            if pil_images:
                if len(pil_images) == 1:
                    gen_prompt["multi_modal_data"] = {}
                    gen_prompt["multi_modal_data"]["image"] = pil_images[0]
                else:
                    od_config = getattr(self._diffusion_engine, "od_config", None)
                    supports_multimodal_inputs = getattr(od_config, "supports_multimodal_inputs", False)
                    if od_config is None:
                        # TODO: entry is asyncOmni. We hack the od config here.
                        supports_multimodal_inputs = True
                    if supports_multimodal_inputs:
                        gen_prompt["multi_modal_data"] = {}
                        gen_prompt["multi_modal_data"]["image"] = pil_images
                    else:
                        return self._create_error_response(
                            "Multiple input images are not supported by the current diffusion model. "
                            "For multi-image editing, start the server with Qwen-Image-Edit-2509 "
                            "and send multiple images in the user message content.",
                            status_code=400,
                        )

            if reference_videos:
                gen_params.extra_args["video_path"] = reference_videos[0]
            if reference_audios:
                gen_params.extra_args["audio_path"] = reference_audios[0]

            # Generate image or audio (e.g. AudioX) via AsyncOmni
            diffusion_engine = cast(AsyncOmni, self._diffusion_engine)
            stage_configs = list(getattr(diffusion_engine, "stage_configs", []) or [])
            default_sampling_params_list = get_default_sampling_params_list(diffusion_engine)
            sampling_params_list = build_stage_sampling_params_list(
                stage_configs,
                default_sampling_params_list,
                diffusion_params=gen_params,
                replace_diffusion_params=True,
            )
            for index, (stage_config, sampling_params) in enumerate(zip(stage_configs, sampling_params_list)):
                if get_stage_type(stage_config) != "diffusion" or index >= len(default_sampling_params_list):
                    continue
                default_extra_args = getattr(default_sampling_params_list[index], "extra_args", None)
                if default_extra_args:
                    sampling_params.extra_args = {
                        **default_extra_args,
                        **(getattr(sampling_params, "extra_args", None) or {}),
                    }
            if not sampling_params_list:
                sampling_params_list = [gen_params]

            result = None
            async for output in diffusion_engine.generate(
                prompt=gen_prompt,
                sampling_params_list=sampling_params_list,
                request_id=request_id,
            ):
                result = output
            if result is None:
                return self._create_error_response("No output generated from AsyncOmni")

            # Text output path (img2text / text2text)
            if is_text_request and result.final_output_type == "text":
                text_body = self._get_diffusion_text_output(result)
                message = ChatMessage(role="assistant", content=text_body)
                choice = ChatCompletionResponseChoice(
                    index=0,
                    message=message,
                    finish_reason="stop",
                    logprobs=None,
                    stop_reason=None,
                )
                response = OmniChatCompletionResponse(
                    id=request_id,
                    created=created_time,
                    model=self._diffusion_model_name,
                    choices=[choice],
                    usage=UsageInfo(
                        prompt_tokens=len(prompt.split()),
                        completion_tokens=len(text_body.split()),
                        total_tokens=len(prompt.split()) + len(text_body.split()),
                    ),
                    metrics=self._get_diffusion_extra_output_params(result),
                )
                logger.info(
                    "Diffusion chat completed for request %s: text output (%d chars)",
                    request_id,
                    len(text_body),
                )
                return response

            # Image output path (text2img / img2img)
            final_output_type = getattr(result, "final_output_type", "image")
            images = getattr(result, "images", [])
            multimodal_output = getattr(result, "multimodal_output", {}) or {}
            stage_durations = result.stage_durations
            peak_memory_mb = result.peak_memory_mb

            if final_output_type == "audio":
                sample_rate = 48000
                for key in ("audio_sample_rate", "sample_rate", "sampling_rate", "sr"):
                    raw_rate = multimodal_output.get(key)
                    try:
                        if raw_rate is not None:
                            sample_rate = int(raw_rate)
                            break
                    except (TypeError, ValueError):
                        pass

                audio_payload = multimodal_output.get("audio")
                if isinstance(audio_payload, list):
                    if len(audio_payload) == 0:
                        audio_payload = None
                    elif len(audio_payload) == 1:
                        audio_payload = audio_payload[0]
                if audio_payload is None:
                    return self._create_error_response("Audio generation completed but no audio was produced.")

                if isinstance(audio_payload, torch.Tensor):
                    audio_tensor = audio_payload.detach().cpu().float()
                else:
                    audio_tensor = torch.as_tensor(audio_payload).detach().cpu().float()
                # Pipelines deliver audio as (C, T), (T,), or (B, C, T) in channels-first
                # convention (torch default). Drop a leading batch dim, then transpose to
                # (T, C) for soundfile / CreateAudio. Flattening here would corrupt stereo.
                if audio_tensor.ndim == 3:
                    audio_tensor = audio_tensor[0]
                if audio_tensor.ndim == 2:
                    audio_tensor = audio_tensor.transpose(0, 1).contiguous()
                elif audio_tensor.ndim > 3:
                    raise ValueError(f"Unexpected audio tensor rank {audio_tensor.ndim}; expected 1-3 dims.")
                audio_array = audio_tensor.numpy()

                audio_format = self._resolve_audio_format(request)
                if isinstance(audio_format, ErrorResponse):
                    return audio_format

                audio_obj = CreateAudio(
                    audio_tensor=audio_array,
                    sample_rate=sample_rate,
                    response_format=audio_format,
                    speed=1.0,
                    base64_encode=True,
                )
                audio_response: AudioResponse = self.create_audio(audio_obj)
                audio_base64 = audio_response.audio_data
                audio_id = f"audio-{uuid.uuid4().hex[:16]}"
                expires_at = int((datetime.now(timezone.utc) + timedelta(hours=24)).timestamp())
                message = ChatMessage(
                    role="assistant",
                    audio=OpenAIChatCompletionAudio(
                        id=audio_id,
                        data=audio_base64,
                        expires_at=expires_at,
                        transcript="",
                    ),
                )
            else:
                # Convert images to base64 content
                image_contents: list[dict[str, Any]] = []
                flat_images = []
                for item in images:
                    if isinstance(item, list):
                        flat_images.extend(item)
                    else:
                        flat_images.append(item)

                for img in flat_images:
                    with BytesIO() as buffer:
                        img.save(buffer, format="PNG")
                        img_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
                    image_contents.append(
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/png;base64,{img_base64}",
                            },
                            "stage_durations": stage_durations,
                            "peak_memory_mb": peak_memory_mb,
                        }
                    )

                # Build response
                if not image_contents:
                    content = "Image generation completed but no images were produced."
                else:
                    content = image_contents

                # Use model_construct to bypass validation for multimodal content
                # (ChatMessage.content only accepts str, but we need list for images)
                # Then use object.__setattr__ to directly set the field, bypassing Pydantic's type checking
                import warnings as warnings_module

                with warnings_module.catch_warnings():
                    warnings_module.filterwarnings("ignore", category=UserWarning, module="pydantic")
                    message = ChatMessage.model_construct(role="assistant")
                    object.__setattr__(message, "content", content)
                    # Mark content as set in fields_set to ensure proper serialization
                    if hasattr(message, "__pydantic_fields_set__"):
                        message.__pydantic_fields_set__.add("content")
            choice = ChatCompletionResponseChoice.model_construct(
                index=0,
                message=message,
                finish_reason="stop",
                logprobs=None,
                stop_reason=None,
            )

            response = OmniChatCompletionResponse(
                id=request_id,
                created=created_time,
                model=self._diffusion_model_name,
                choices=[choice],
                usage=UsageInfo(
                    prompt_tokens=len(prompt.split()),
                    completion_tokens=1,
                    total_tokens=len(prompt.split()) + 1,
                ),
                metrics=self._get_diffusion_extra_output_params(result),
            )

            logger.info(
                "Diffusion chat completed for request %s: output_type=%s, image_count=%d",
                request_id,
                final_output_type,
                len(images),
            )

            return response

        except OmniClientError as e:
            logger.info("Client error during diffusion chat completion: %s", e)
            return self._create_error_response(
                e.message,
                err_type=e.error_type,
                status_code=e.status_code,
            )
        except Exception as e:
            logger.exception("Diffusion chat completion failed: %s", e)
            return self._create_error_response(
                f"Image generation failed: {str(e)}",
                status_code=500,
            )

    def _extract_diffusion_prompt_and_media(
        self,
        messages: list[dict[str, Any]],
    ) -> tuple[str, list[str], list[str], list[str]]:
        """Extract text prompt and multimodal inputs from chat messages.

        Args:
            messages: List of chat messages

        Returns:
            Tuple of (prompt_text, list_of_base64_images, list_of_video_urls, list_of_audio_urls)
        """
        prompt_parts: list[str] = []
        images: list[str] = []
        videos: list[str] = []
        audios: list[str] = []

        for message in messages:
            role = message.get("role", "")
            if role != "user":
                continue

            content = message.get("content", "")

            # String content
            if isinstance(content, str):
                prompt_parts.append(content)
                continue

            # List of content items
            if isinstance(content, list):
                for item in content:
                    if isinstance(item, str):
                        prompt_parts.append(item)
                    elif isinstance(item, dict):
                        # Handle {"type": "text", "text": "..."} format
                        if item.get("type") == "text":
                            prompt_parts.append(item.get("text", ""))
                        # Handle {"text": "..."} format
                        elif "text" in item and "type" not in item:
                            prompt_parts.append(item["text"])
                        # Handle {"type": "image_url", "image_url": {"url": "..."}}
                        elif item.get("type") == "image_url":
                            url = item.get("image_url", {}).get("url", "")
                            if url.startswith("data:image"):
                                try:
                                    _, b64_data = url.split(",", 1)
                                    images.append(b64_data)
                                except ValueError:
                                    logger.warning("Invalid data URL format")
                        elif item.get("type") == "video_url":
                            url = item.get("video_url", {}).get("url", "")
                            if isinstance(url, str) and url:
                                videos.append(url)
                        elif item.get("type") == "audio_url":
                            url = item.get("audio_url", {}).get("url", "")
                            if isinstance(url, str) and url:
                                audios.append(url)
                        # Handle {"image": "base64..."} format
                        elif "image" in item:
                            images.append(item["image"])
                        elif "video" in item and isinstance(item["video"], str):
                            videos.append(item["video"])
                        elif "audio" in item and isinstance(item["audio"], str):
                            audios.append(item["audio"])

        prompt = " ".join(prompt_parts).strip()
        return prompt, images, videos, audios

    def _extract_diffusion_prompt_and_images_from_messages(
        self,
        messages: list[Any],
    ) -> tuple[str, list[str]]:
        """Normalize mixed message types and extract prompt + reference images once."""
        prompt, images, _videos, _audios = self._extract_diffusion_prompt_and_media(self._messages_to_dicts(messages))
        return prompt, images

    @staticmethod
    def _messages_to_dicts(messages: list[Any]) -> list[dict[str, Any]]:
        """Normalize request messages to plain dicts."""
        out: list[dict[str, Any]] = []
        for msg in messages:
            if hasattr(msg, "model_dump"):
                out.append(msg.model_dump())
            elif isinstance(msg, dict):
                out.append(msg)
            else:
                out.append(
                    {
                        "role": getattr(msg, "role", "user"),
                        "content": getattr(msg, "content", ""),
                    }
                )
        return out

    @staticmethod
    def _resolve_height_width_from_extra_body(extra_body: dict[str, Any]) -> tuple[Any, Any]:
        """Extract generation height/width with optional size string fallback."""
        height = extra_body.get("height")
        width = extra_body.get("width")

        if "size" in extra_body and (height is None or width is None):
            try:
                size_str = extra_body["size"]
                if isinstance(size_str, str) and "x" in size_str.lower():
                    w, h = size_str.lower().split("x")
                    width, height = int(w), int(h)
            except ValueError:
                logger.warning("Invalid size format: %s", extra_body.get("size"))

        return height, width

    def _resolve_audio_format(self, request: ChatCompletionRequest) -> str | ErrorResponse:
        """Extract and validate the audio output format from a chat completion request."""
        audio_params = getattr(request, "audio", None)
        if isinstance(audio_params, dict):
            audio_format = audio_params.get("format", DEFAULT_AUDIO_FORMAT)
        else:
            audio_format = DEFAULT_AUDIO_FORMAT
        if audio_format not in SUPPORTED_CHAT_AUDIO_FORMATS:
            return self._create_error_response(
                f"Invalid audio format '{audio_format}'. Supported formats: {sorted(SUPPORTED_CHAT_AUDIO_FORMATS)}",
            )
        if audio_format == "pcm16":
            audio_format = "pcm"
        return audio_format

    def _create_error_response(
        self,
        message: str,
        err_type: str = "BadRequestError",
        status_code: int = 400,
    ) -> ErrorResponse:
        """Create an error response following OpenAI error format."""
        return ErrorResponse(
            error=ErrorInfo(
                message=message,
                type=err_type,
                code=status_code,
            )
        )