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.diffusion_engine import get_extra_body_params, get_extra_output_params
from vllm_omni.entrypoints.async_omni import AsyncOmni
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,
)

try:
    import soundfile
except ImportError:
    soundfile = None


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
from vllm.entrypoints.openai.engine.protocol import (
    DeltaFunctionCall,
    DeltaMessage,
    DeltaToolCall,
    ErrorInfo,
    ErrorResponse,
    FunctionCall,
    FunctionDefinition,
    PromptTokenUsageInfo,
    RequestResponseMetadata,
    ToolCall,
    UsageInfo,
)
from vllm.entrypoints.openai.engine.serving import ChatLikeRequest, clamp_prompt_logprobs
from vllm.entrypoints.openai.parser.harmony_utils import (
    get_streamable_parser_for_assistant,
    parse_chat_output,
)
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
from vllm.entrypoints.openai.utils import maybe_filter_parallel_tool_calls
from vllm.entrypoints.utils import should_include_usage
from vllm.inputs import PromptType
from vllm.logger import init_logger
from vllm.outputs import RequestOutput
from vllm.reasoning import ReasoningParser
from vllm.renderers import BaseRenderer, merge_kwargs
from vllm.renderers.inputs import TokPrompt
from vllm.sampling_params import 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.mistral_tool_parser import MistralToolCall
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 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,
    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.utils.audio import audio_chunk_pcm_bytes, audio_chunk_sample_rate

logger = init_logger(__name__)


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

    @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.engine_client = None
        instance.has_kv_connector = False
        return instance

    def _get_diffusion_extra_body_params(self) -> frozenset[str]:
        """Return the pipeline-declared ``EXTRA_BODY_PARAMS`` (cached)."""
        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 EXTRA_BODY_PARAMS from pipeline: %s", e)

        self._diffusion_extra_body_params = params
        return params

    def _get_diffusion_extra_output_params(
        self,
        custom_output: dict[str, Any] | None,
    ) -> dict[str, Any] | None:
        """Pick pipeline-declared ``EXTRA_OUTPUT_PARAMS`` keys from *custom_output*."""
        if not custom_output:
            return None

        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 EXTRA_OUTPUT_PARAMS from pipeline: %s", e)
            self._diffusion_extra_output_params = params

        if not self._diffusion_extra_output_params:
            return None
        out = {k: custom_output[k] for k in self._diffusion_extra_output_params if k in custom_output}
        return out or None

    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:
        # Handle diffusion mode
        if self._diffusion_mode:
            return await self._create_diffusion_chat_completion(request, 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.reasoning_parser_cls:
                chat_template_kwargs = self._prepare_extra_chat_template_kwargs(
                    request.chat_template_kwargs,
                    self.default_chat_template_kwargs,
                )
                reasoning_parser = self.reasoning_parser_cls(
                    tokenizer,
                    chat_template_kwargs=chat_template_kwargs,  # type: ignore[call-arg]
                )

            tool_parser = self.tool_parser

            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._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

                chat_template_kwargs = request.chat_template_kwargs or {}
                chat_template_kwargs.update(reasoning_effort=request.reasoning_effort)

                # Merge chat_template_kwargs with defaults
                merged_template_kwargs = self._prepare_extra_chat_template_kwargs(
                    chat_template_kwargs,
                    self.default_chat_template_kwargs,
                )
                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.openai_serving_render._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

        num_inference_steps = None
        cfg_text_scale = None
        cfg_img_scale = None
        # 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")

                # [NOTE] When sending request via openai client Python library,
                #   `extra_body` is flattented and merged into the payload's root.
                #   These extra fields are accessible via `model_extra` property (from Pydantic base class).
                #   When sending raw request with curl, no flattening happens. Directly read the `extra_body` dict.
                extra_body = getattr(request, "extra_body", None)
                if not extra_body:
                    extra_body = request.model_extra or {}

                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")
                cfg_text_scale = extra_body.get("cfg_text_scale")
                cfg_img_scale = extra_body.get("cfg_img_scale")

                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()
                    }

                # Forward image_gen extra_args from diffusion-stage sampling
                # params into the prompt dict so that thinker2imagegen can
                # access them (e.g. byte5_text for ByT5 glyph rendering).
                if hasattr(request, "sampling_params_list") and request.sampling_params_list:
                    for sp_raw in request.sampling_params_list:
                        ea = (
                            (sp_raw.get("extra_args") or {})
                            if isinstance(sp_raw, dict)
                            else (getattr(sp_raw, "extra_args", None) or {})
                        )
                        ig = ea.get("image_gen") or {}
                        if ig:
                            tprompt["image_gen_extra_args"] = ig
                            break

                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]] = []
        try:
            for i, engine_prompt in enumerate(engine_prompts):
                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)

                # Apply user-specified overrides to diffusion stage(s) for image generation
                for idx, sp in enumerate(sampling_params_list):
                    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 hasattr(sp, "extra_args") and sp.extra_args is not None:
                        if cfg_text_scale is not None:
                            sp.extra_args["cfg_text_scale"] = cfg_text_scale
                        if cfg_img_scale is not None:
                            sp.extra_args["cfg_img_scale"] = cfg_img_scale
                        for _key in self._get_diffusion_extra_body_params():
                            _val = extra_body.get(_key)
                            if _val is not None:
                                sp.extra_args[_key] = _val

                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,
                )

                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,
                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,
            )
        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),
        )

        (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

        speaker = getattr(request, "voice", None) or getattr(request, "speaker", None)
        normalized = validate_requested_speaker(speaker, self._get_supported_speakers())
        if normalized is not None:
            if "additional_information" not in engine_prompt or engine_prompt["additional_information"] is None:
                engine_prompt["additional_information"] = {}
            engine_prompt["additional_information"]["speaker"] = [normalized]

        language = getattr(request, "language", None)
        if language is not None and isinstance(language, str) and language.strip():
            if "additional_information" not in engine_prompt or engine_prompt["additional_information"] is None:
                engine_prompt["additional_information"] = {}
            engine_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():
            if "additional_information" not in engine_prompt or engine_prompt["additional_information"] is None:
                engine_prompt["additional_information"] = {}
            engine_prompt["additional_information"]["instruction"] = instructions.strip()

        return conversation, [engine_prompt]

    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``.
        """
        stage_configs = list(getattr(self.engine_client, "stage_configs", []) 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}")
        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.

        For models needing spatial metadata (e.g. GLM-Image), target_h/w is
        injected into extra_args so the runner can build M-RoPE position grids.
        max_tokens is NOT computed dynamically — it uses the deploy YAML default.

        Args:
            default_params: Default SamplingParams from stage config YAML.
            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)

        # For GLM-Image: compute max_tokens from height/width with mode-aware
        # budgeting (t2i vs i2i).
        extra_body = getattr(request, "extra_body", {}) or {}
        height, width = self._resolve_height_width_from_extra_body(extra_body)

        if height is not None and width is not None:
            # Keep target size in stage-0 sampling params so runner/model can
            # build deterministic M-RoPE grids for t2i (no MM features).
            extra_args = dict(getattr(params, "extra_args", {}) or {})
            extra_args["target_h"] = int(height)
            extra_args["target_w"] = int(width)
            params.extra_args = extra_args

        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 _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,
        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
        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:
            if tool_choice_auto and self.tool_parser:
                tool_parsers: list[ToolParser | None] = [self.tool_parser(tokenizer, request.tools)] * 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.request_output
                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
                    # 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] = self.extract_tool_call_required_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,
                                )
                                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
                        stage_pools = getattr(self.engine_client.engine, "stage_pools", None)
                        # The orchestrator binds requests by their internal id,
                        # not the user-visible external id, so look up the
                        # replica with req_state.request_id (internal).
                        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)
                    # 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,
                )
                if self.enable_prompt_tokens_details and num_cached_tokens:
                    final_usage.prompt_tokens_details = PromptTokenUsageInfo(cached_tokens=num_cached_tokens)

                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,
    ) -> 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 omni_outputs.request_output is not None and not getattr(omni_outputs.request_output, "finished", False):
                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.request_output is not None:
                    (
                        choices_data,
                        usage,
                        prompt_logprobs,
                        prompt_token_ids,
                        kv_transfer_params,
                    ) = self._create_text_choice(
                        request,
                        omni_outputs,
                        tokenizer,
                        conversation,
                        role,
                        reasoning_parser,
                    )
                    final_res = omni_outputs.request_output
                else:
                    # Diffusion pipeline text output (e.g. single-stage
                    # img2text / text2text) — no AR request_output, so build
                    # a simple text choice from custom_output.
                    text_body = (omni_outputs.custom_output or {}).get("text_output", "")
                    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)
            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.custom_output)
                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,
    ):
        final_res = omni_outputs.request_output
        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:
                reasoning, content, _ = parse_chat_output(token_ids)
                if not request.include_reasoning:
                    reasoning = None

                if self.tool_parser is not None:
                    tool_parser = self.tool_parser(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 not self.tool_parser) 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.tool_parser
            ):
                try:
                    tool_parser = self.tool_parser(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,
        )
        if self.enable_prompt_tokens_details and final_res.num_cached_tokens:
            usage.prompt_tokens_details = PromptTokenUsageInfo(cached_tokens=final_res.num_cached_tokens)

        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
    ):
        choices: list[ChatCompletionResponseChoice] = []
        final_res = omni_outputs.request_output
        # OMNI: Access multimodal_output from CompletionOutput (outputs[0]), not from RequestOutput
        # Reference: examples/offline_inference/qwen3_omni/end2end.py line 421
        mm_output = final_res.outputs[0].multimodal_output
        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:
            return self._create_error_response("Audio generation completed but no audio was produced.")
        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_obj = CreateAudio(
            audio_tensor=audio_tensor,
            sample_rate=sample_rate,
            response_format="wav",
            speed=1.0,
            stream_format="audio",
            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] = []
        final_res = omni_outputs.request_output

        # 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 request_output for pipeline mode
        # OMNI: Access multimodal_output from CompletionOutput (outputs[0]), not from RequestOutput
        elif final_res is not None and final_res.outputs:
            completion_output = final_res.outputs[0]
            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")
        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,
                    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,
                )
                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,
            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.request_output, "images", [])
        stage_durations = result.stage_durations
        peak_memory_mb = result.peak_memory_mb
        cot_output = None

        req_out = getattr(result, "request_output", None)
        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 = getattr(result, "request_output", None)
        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:
                    request_output = output.request_output
                    for completion in request_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.request_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,
        raw_request: Request | None = None,
    ) -> ChatCompletionResponse | ErrorResponse:
        """Generate images via chat completion interface for diffusion models.

        Args:
            request: Chat completion request
            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
            )

            # Extract generation parameters from extra_body (preferred)
            # Reference: text_to_image.py and text_to_video.py for supported parameters
            # [NOTE] When sending request via openai client Python library,
            #   `extra_body` is flattented and merged into the payload's root.
            #   These extra fields are accessible via `model_extra` property (from Pydantic base class).
            #   When sending raw request with curl, no flattening happens. Directly read the `extra_body` dict.
            extra_body = getattr(request, "extra_body", None)
            if not extra_body:
                extra_body = request.model_extra or {}

            # 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")
            guidance_scale = extra_body.get("guidance_scale")
            true_cfg_scale = extra_body.get("true_cfg_scale") or extra_body.get("cfg_scale")
            cfg_text_scale = extra_body.get("cfg_text_scale")
            cfg_img_scale = extra_body.get("cfg_img_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 e:
                return self._create_error_response(str(e), status_code=400)

            logger.info(
                "Diffusion chat request %s: prompt=%r, ref_images=%d, params=%s",
                request_id,
                prompt[:50] + "..." if len(prompt) > 50 else prompt,
                len(reference_images),
                {k: v for k, v in extra_body.items() if v is not None},
            )

            # 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 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
            if cfg_text_scale is not None:
                gen_params.extra_args["cfg_text_scale"] = cfg_text_scale
            if cfg_img_scale is not None:
                gen_params.extra_args["cfg_img_scale"] = cfg_img_scale
            for _key in self._get_diffusion_extra_body_params():
                _val = extra_body.get(_key)
                if _val is not None:
                    gen_params.extra_args[_key] = _val
            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

            # Pipeline-agnostic escape hatch (mirrors ``extra_params`` on the /v1/videos
            # endpoint in ``serving_video.py``): a single reserved ``extra_args`` dict in
            # ``extra_body`` flows straight into ``gen_params.extra_args``, with no keys
            # hardcoded here.
            extra_args_body = extra_body.get("extra_args")
            if isinstance(extra_args_body, dict):
                gen_params.extra_args.update(extra_args_body)

            # 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 e:  # pragma: no cover - safeguard
                    logger.warning("Failed to parse LoRA request: %s", e)

            # 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 [])
            sampling_params_list = build_stage_sampling_params_list(
                stage_configs,
                get_default_sampling_params_list(diffusion_engine),
                diffusion_params=gen_params,
                replace_diffusion_params=True,
            )

            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 = (result.custom_output or {}).get("text_output", "")
                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.custom_output),
                )
                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")
            # Handle nested OmniRequestOutput structure where images might be in request_output
            images = getattr(result.request_output, "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_obj = CreateAudio(
                    audio_tensor=audio_array,
                    sample_rate=sample_rate,
                    response_format="wav",
                    speed=1.0,
                    stream_format="audio",
                    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.custom_output),
            )

            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 _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,
            )
        )