live_server.py

#!/usr/bin/env python3
"""Persistent GPU inference API for the AdaSpatial-MLLM interactive demo."""

from __future__ import annotations

import io
import os
import re
import threading
import time
from contextlib import asynccontextmanager
from typing import Annotated, Any, Literal

from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.concurrency import run_in_threadpool
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image, ImageOps, UnidentifiedImageError

from thinker_r.confidence import score_generated_choice
from thinker_r.routing import (
    escalation_reasons,
    is_directional,
    majority_vote,
    mirror_direction_text,
    resolve_task_type,
    routed_condition,
)

PROMPT_SUFFIX = (
    "\nFirst imagine the mental 3D scene, think about the reasoning process in the mind "
    "and then provide the user with the answer. The reasoning process and answer are "
    "enclosed within <think> </think> and <answer> </answer> tags, respectively. Special "
    "tokens should be used to represent mental 3D scene at the beginning of your response, "
    "i.e., mental 3D scene here <think> reasoning process here </think><answer> answer here "
    "</answer>."
)
IMAGE_PAD = "<|vision_start|><|image_pad|><|vision_end|>"
MAX_FILE_BYTES = 10 * 1024 * 1024
MAX_IMAGE_SIDE = int(os.environ.get("THINKER_MAX_IMAGE_SIDE", "640"))
MODEL_PATH = os.environ.get("THINKER_MODEL_PATH", "/root/autodl-tmp/models/3DThinker-Mindcube")
ALLOWED_ORIGINS = [
    origin.strip()
    for origin in os.environ.get(
        "THINKER_ALLOWED_ORIGINS",
        "http://127.0.0.1:7860,http://localhost:7860",
    ).split(",")
    if origin.strip()
]


def extract_tag(text: str, tag: str) -> str:
    match = re.search(rf"<{tag}>\s*(.*?)\s*</{tag}>", text, flags=re.IGNORECASE | re.DOTALL)
    return match.group(1).strip() if match else ""


def extract_choice(text: str) -> str | None:
    candidate = extract_tag(text, "answer") or text
    patterns = (
        r"(?:answer|choice|option)\s*(?:is|:)?\s*\(?([A-D])\)?\b",
        r"^\s*\(?([A-D])\)?(?:[.)\s]|$)",
        r"\b([A-D])\b",
    )
    for pattern in patterns:
        matches = re.findall(pattern, candidate, flags=re.IGNORECASE | re.MULTILINE)
        if matches:
            return matches[-1].upper()
    return None


def scalarize(value: Any) -> Any:
    if isinstance(value, dict):
        return {key: scalarize(item) for key, item in value.items()}
    if isinstance(value, (list, tuple)):
        return [scalarize(item) for item in value]
    if hasattr(value, "item"):
        return value.item()
    return value


def compose_question(question: str, options: list[str]) -> str:
    clean_question = question.strip()
    if not clean_question:
        raise HTTPException(status_code=422, detail="Question cannot be empty")
    clean_options = [option.strip() for option in options]
    if any(not option for option in clean_options):
        raise HTTPException(status_code=422, detail="All four answer options are required")
    return clean_question + " " + " ".join(
        f"{letter}. {option}" for letter, option in zip("ABCD", clean_options, strict=True)
    )


async def decode_images(files: list[UploadFile]) -> list[Image.Image]:
    if not 2 <= len(files) <= 4:
        raise HTTPException(status_code=422, detail="Upload between 2 and 4 viewpoint images")
    images = []
    for upload in files:
        if upload.content_type not in {"image/jpeg", "image/png", "image/webp"}:
            raise HTTPException(status_code=415, detail=f"Unsupported image type: {upload.content_type}")
        payload = await upload.read(MAX_FILE_BYTES + 1)
        if len(payload) > MAX_FILE_BYTES:
            raise HTTPException(status_code=413, detail=f"{upload.filename} exceeds 10 MB")
        try:
            with Image.open(io.BytesIO(payload)) as image:
                converted = image.convert("RGB")
                converted.thumbnail((MAX_IMAGE_SIDE, MAX_IMAGE_SIDE), Image.Resampling.LANCZOS)
                images.append(converted.copy())
        except (UnidentifiedImageError, OSError) as error:
            raise HTTPException(status_code=422, detail=f"Invalid image: {upload.filename}") from error
    return images


class ModelRuntime:
    """Lazy-loaded model runtime with a single GPU inference lane."""

    def __init__(self) -> None:
        self.processor = None
        self.model = None
        self.torch = None
        self.process_vision_info = None
        self.load_lock = threading.Lock()
        self.infer_lock = threading.Lock()
        self.loaded_at: float | None = None

    @property
    def loaded(self) -> bool:
        return self.model is not None and self.processor is not None

    def load(self) -> None:
        if self.loaded:
            return
        with self.load_lock:
            if self.loaded:
                return
            import torch
            from qwen_vl_utils import process_vision_info
            from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration

            torch.set_float32_matmul_precision("high")
            torch.backends.cuda.matmul.allow_tf32 = True
            processor = AutoProcessor.from_pretrained(
                MODEL_PATH,
                local_files_only=True,
                use_fast=False,
            )
            processor.tokenizer.padding_side = "left"
            model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
                MODEL_PATH,
                device_map="cuda:0",
                torch_dtype=torch.bfloat16,
                local_files_only=True,
                low_cpu_mem_usage=True,
            ).eval()
            model.generation_config.temperature = None
            model.generation_config.top_p = None
            self.processor = processor
            self.model = model
            self.torch = torch
            self.process_vision_info = process_vision_info
            self.loaded_at = time.time()

    def infer_once(
        self,
        images: list[Image.Image],
        prompt: str,
        *,
        condition: str,
        adaptive: bool,
        latent_steps: int,
    ) -> dict[str, Any]:
        self.load()
        torch = self.torch
        processor = self.processor
        model = self.model
        content = [{"type": "image", "image": image} for image in images]
        content.append(
            {
                "type": "text",
                "text": "<image>" * len(images) + prompt + PROMPT_SUFFIX,
            }
        )
        conversation = [{"role": "user", "content": content}]
        text = processor.tokenizer.apply_chat_template(conversation, tokenize=False)
        text = text.replace(IMAGE_PAD, "").replace("<image>", IMAGE_PAD)
        text += "<|im_start|>assistant"
        image_inputs, _ = self.process_vision_info(conversation)
        inputs = processor(text=[text], images=image_inputs, return_tensors="pt", padding=True)
        padded_input_tokens = int(inputs.input_ids.shape[-1])
        input_tokens = int(inputs.attention_mask.sum().item())
        inputs = inputs.to(model.device)

        torch.manual_seed(42)
        torch.cuda.manual_seed_all(42)
        model._thinker_inference = {
            "max_latent_steps": latent_steps,
            "force_latent_steps": False,
            "intervention": "normal",
            "adaptive": adaptive,
            "min_latent_steps": 4,
            "convergence_epsilon": 0.04,
            "convergence_patience": 2,
            "seed": 42,
        }
        torch.cuda.reset_peak_memory_stats()
        torch.cuda.synchronize()
        started = time.perf_counter()
        generation_output = None
        try:
            with torch.inference_mode():
                generation_output = model.generate(
                    **inputs,
                    max_new_tokens=1024,
                    do_sample=False,
                    tokenizer=processor.tokenizer,
                    return_dict_in_generate=True,
                    output_scores=True,
                )
            torch.cuda.synchronize()
            latency = time.perf_counter() - started
            peak_vram_gib = torch.cuda.max_memory_allocated() / 2**30
            generated = generation_output.sequences[0, padded_input_tokens:]
            eos_positions = (generated == processor.tokenizer.eos_token_id).nonzero(as_tuple=True)[0]
            if eos_positions.numel():
                generated = generated[: int(eos_positions[0].item()) + 1]
            response = processor.tokenizer.decode(generated, skip_special_tokens=False)
            confidence = score_generated_choice(
                processor.tokenizer,
                generated.tolist(),
                generation_output.scores,
            )
            if confidence:
                confidence.pop("candidate_token_ids", None)
            trace = scalarize(dict(getattr(model, "_thinker_last_trace", {})))
            return {
                "condition": condition,
                "prediction": extract_choice(response),
                "answer": extract_tag(response, "answer"),
                "reasoning": extract_tag(response, "think"),
                "response": response,
                "steps": trace.get("latent_pad_steps"),
                "compute_steps": trace.get("latent_pad_steps"),
                "forward_steps": trace.get("latent_forward_steps"),
                "exit_reason": trace.get("exit_reason"),
                "distances": [
                    float(values[0]) for values in trace.get("convergence_distances", [])
                ],
                "latency_s": latency,
                "peak_vram_gib": peak_vram_gib,
                "input_tokens": input_tokens,
                "output_tokens": int(generated.numel()),
                "choice_confidence": confidence,
                "trace": trace,
            }
        finally:
            del inputs
            if generation_output is not None:
                del generation_output
            torch.cuda.empty_cache()

    @staticmethod
    def _optimized_result(
        components: dict[str, dict[str, Any]],
        selected: str,
        task_type: str,
        route_source: str,
        reasons: list[str],
    ) -> dict[str, Any]:
        if reasons:
            vote = majority_vote(
                components,
                ("fixed4", "fixed8", "adaptive"),
                fallback="adaptive",
            )
            prediction = vote["prediction"]
            matching = [
                name
                for name in ("adaptive", "fixed8", "fixed4")
                if components[name].get("prediction") == prediction
            ]
            representative_name = matching[0] if matching else "adaptive"
            result = dict(components[representative_name])
            ordered_votes = sorted(vote["votes"].values(), reverse=True)
            vote_margin = (
                (ordered_votes[0] - ordered_votes[1]) / sum(ordered_votes)
                if len(ordered_votes) > 1
                else 1.0
            )
            result.update(
                {
                    "prediction": prediction,
                    "condition": "optimized",
                    "strategy": "task_aware_cascade_vote",
                    "exit_reason": "cascade_vote",
                    "compute_steps": sum(run["steps"] for run in components.values()),
                    "latency_s": sum(run["latency_s"] for run in components.values()),
                    "peak_vram_gib": max(run["peak_vram_gib"] for run in components.values()),
                    "choice_confidence": {
                        "source": "budget_vote",
                        "confidence": vote["agreement"],
                        "margin": vote_margin,
                        "probabilities": {},
                    },
                    "votes": vote["votes"],
                    "vote_tie": vote["tie"],
                }
            )
        else:
            result = dict(components[selected])
            result.update(
                {
                    "condition": "optimized",
                    "strategy": f"task_aware_{selected}",
                    "compute_steps": sum(run["steps"] for run in components.values()),
                    "latency_s": sum(run["latency_s"] for run in components.values()),
                    "peak_vram_gib": max(run["peak_vram_gib"] for run in components.values()),
                }
            )
        mirror_run = components.get("mirror_adaptive")
        result.update(
            {
                "task_type": task_type,
                "task_route_source": route_source,
                "escalated": bool(reasons),
                "escalation_reasons": reasons,
                "component_predictions": {
                    name: run.get("prediction") for name, run in components.items()
                },
                "component_steps": {name: run.get("steps") for name, run in components.items()},
                "run_count": len(components),
                "symmetry_checked": mirror_run is not None,
                "symmetry_prediction": mirror_run.get("prediction") if mirror_run else None,
                "symmetry_consistent": (
                    mirror_run.get("prediction") == components["adaptive"].get("prediction")
                    if mirror_run
                    else None
                ),
            }
        )
        return result

    def optimize(
        self,
        images: list[Image.Image],
        prompt: str,
        task_type: str,
        route_source: str,
        *,
        fixed12: dict[str, Any] | None = None,
        symmetry_check: bool = False,
    ) -> tuple[dict[str, Any], dict[str, dict[str, Any]]]:
        selected = routed_condition(task_type)
        if selected == "fixed12":
            fixed = fixed12 or self.infer_once(
                images,
                prompt,
                condition="fixed12",
                adaptive=False,
                latent_steps=12,
            )
            components = {"fixed12": fixed}
            return (
                self._optimized_result(components, selected, task_type, route_source, []),
                components,
            )

        adaptive = self.infer_once(
            images,
            prompt,
            condition="adaptive",
            adaptive=True,
            latent_steps=12,
        )
        components = {"adaptive": adaptive}
        reasons = escalation_reasons(task_type, adaptive)
        if symmetry_check:
            mirrored_images = [ImageOps.mirror(image) for image in images]
            try:
                components["mirror_adaptive"] = self.infer_once(
                    mirrored_images,
                    mirror_direction_text(prompt),
                    condition="mirror_adaptive",
                    adaptive=True,
                    latent_steps=12,
                )
            finally:
                for image in mirrored_images:
                    image.close()
            if components["mirror_adaptive"].get("prediction") != adaptive.get("prediction"):
                reasons.append("mirror_inconsistent")
        if reasons:
            components["fixed4"] = self.infer_once(
                images,
                prompt,
                condition="fixed4",
                adaptive=False,
                latent_steps=4,
            )
            components["fixed8"] = self.infer_once(
                images,
                prompt,
                condition="fixed8",
                adaptive=False,
                latent_steps=8,
            )
        return (
            self._optimized_result(components, selected, task_type, route_source, reasons),
            components,
        )

    def compare(
        self,
        images: list[Image.Image],
        prompt: str,
        task_type: str,
        route_source: str,
        symmetry_check: bool,
    ) -> dict[str, Any]:
        with self.infer_lock:
            fixed = self.infer_once(
                images,
                prompt,
                condition="fixed12",
                adaptive=False,
                latent_steps=12,
            )
            optimized, components = self.optimize(
                images,
                prompt,
                task_type,
                route_source,
                fixed12=fixed,
                symmetry_check=symmetry_check,
            )
        return {"fixed12": fixed, **components, "optimized": optimized}

    def run_optimized(
        self,
        images: list[Image.Image],
        prompt: str,
        task_type: str,
        route_source: str,
        symmetry_check: bool,
    ) -> dict[str, Any]:
        """Run only the deployed policy, without paying for a comparison baseline."""

        with self.infer_lock:
            optimized, components = self.optimize(
                images,
                prompt,
                task_type,
                route_source,
                symmetry_check=symmetry_check,
            )
        return {**components, "optimized": optimized}

    def recover_cuda(self) -> None:
        if self.torch is not None:
            self.torch.cuda.empty_cache()


runtime = ModelRuntime()


@asynccontextmanager
async def lifespan(_: FastAPI):
    yield


app = FastAPI(title="AdaSpatial-MLLM Live Inference", version="1.0.0", lifespan=lifespan)
app.add_middleware(
    CORSMiddleware,
    allow_origins=ALLOWED_ORIGINS,
    allow_credentials=False,
    allow_methods=["GET", "POST"],
    allow_headers=["Content-Type"],
)


@app.get("/health")
def health() -> dict[str, Any]:
    gpu = None
    available = False
    try:
        import torch

        available = torch.cuda.is_available()
        gpu = torch.cuda.get_device_name(0) if available else None
    except ImportError:
        pass
    return {
        "status": "ok" if available else "degraded",
        "gpu_available": available,
        "gpu": gpu,
        "model_loaded": runtime.loaded,
        "model_path": MODEL_PATH,
        "max_image_side": MAX_IMAGE_SIDE,
        "policy": "task-aware adaptive cascade",
    }


@app.post("/infer/compare")
async def infer_compare(
    images: Annotated[list[UploadFile], File()],
    question: Annotated[str, Form(min_length=5, max_length=2000)],
    option_a: Annotated[str, Form(min_length=1, max_length=300)],
    option_b: Annotated[str, Form(min_length=1, max_length=300)],
    option_c: Annotated[str, Form(min_length=1, max_length=300)],
    option_d: Annotated[str, Form(min_length=1, max_length=300)],
    task_type: Annotated[Literal["auto", "among", "rotation", "around"], Form()] = "auto",
    reference_answer: Annotated[Literal["", "A", "B", "C", "D"], Form()] = "",
    symmetry_check: Annotated[bool, Form()] = False,
) -> dict[str, Any]:
    decoded = await decode_images(images)
    options = [option_a, option_b, option_c, option_d]
    prompt = compose_question(question, options)
    resolved_task, route_source = resolve_task_type(question, task_type)
    started = time.perf_counter()
    try:
        runs = await run_in_threadpool(
            runtime.compare,
            decoded,
            prompt,
            resolved_task,
            route_source,
            symmetry_check,
        )
    except RuntimeError as error:
        runtime.recover_cuda()
        message = str(error)
        if "out of memory" in message.lower():
            raise HTTPException(
                status_code=503,
                detail=(
                    "GPU memory was exhausted by this input. The cache was cleared; "
                    "retry with fewer or smaller images."
                ),
            ) from error
        raise HTTPException(status_code=500, detail=f"Model inference failed: {message[:240]}") from error
    for run in runs.values():
        run["correct"] = None if not reference_answer else run["prediction"] == reference_answer
    return {
        "question": question.strip(),
        "options": [
            {"key": letter, "text": option.strip()}
            for letter, option in zip("ABCD", options, strict=True)
        ],
        "reference_answer": reference_answer or None,
        "task_type": resolved_task,
        "task_route_source": route_source,
        "directional": is_directional(question, options),
        "symmetry_check": symmetry_check,
        "runs": runs,
        "total_latency_s": time.perf_counter() - started,
    }


@app.post("/infer/optimized")
async def infer_optimized(
    images: Annotated[list[UploadFile], File()],
    question: Annotated[str, Form(min_length=5, max_length=2000)],
    option_a: Annotated[str, Form(min_length=1, max_length=300)],
    option_b: Annotated[str, Form(min_length=1, max_length=300)],
    option_c: Annotated[str, Form(min_length=1, max_length=300)],
    option_d: Annotated[str, Form(min_length=1, max_length=300)],
    task_type: Annotated[Literal["auto", "among", "rotation", "around"], Form()] = "auto",
    reference_answer: Annotated[Literal["", "A", "B", "C", "D"], Form()] = "",
    symmetry_check: Annotated[bool, Form()] = False,
) -> dict[str, Any]:
    """Run the training-free Adaptive+ deployment policy without a fixed-12 baseline."""

    decoded = await decode_images(images)
    options = [option_a, option_b, option_c, option_d]
    prompt = compose_question(question, options)
    resolved_task, route_source = resolve_task_type(question, task_type)
    started = time.perf_counter()
    try:
        runs = await run_in_threadpool(
            runtime.run_optimized,
            decoded,
            prompt,
            resolved_task,
            route_source,
            symmetry_check,
        )
    except RuntimeError as error:
        runtime.recover_cuda()
        message = str(error)
        if "out of memory" in message.lower():
            raise HTTPException(
                status_code=503,
                detail=(
                    "GPU memory was exhausted by this input. The cache was cleared; "
                    "retry with fewer or smaller images."
                ),
            ) from error
        raise HTTPException(status_code=500, detail=f"Model inference failed: {message[:240]}") from error
    for run in runs.values():
        run["correct"] = None if not reference_answer else run["prediction"] == reference_answer
    return {
        "question": question.strip(),
        "options": [
            {"key": letter, "text": option.strip()}
            for letter, option in zip("ABCD", options, strict=True)
        ],
        "reference_answer": reference_answer or None,
        "task_type": resolved_task,
        "task_route_source": route_source,
        "directional": is_directional(question, options),
        "symmetry_check": symmetry_check,
        "runs": runs,
        "total_latency_s": time.perf_counter() - started,
    }