KillerKing93/Transformers-InferenceServer-OpenAPI
0
1#!/usr/bin/env python2# -*- coding: utf-8 -*-3"""4FastAPI Inference Server (OpenAI-compatible) for Qwen3-VL multimodal model.5 6- Default model: unsloth/Qwen3-4B-Instruct-25077- Endpoints:8 * GET /openapi.yaml (OpenAPI schema in YAML)9 * GET /health (readiness + context report)10 * POST /v1/chat/completions (non-stream and streaming SSE)11 * POST /v1/cancel/{session_id} (custom cancel endpoint)12 13Notes:14- Uses Hugging Face Transformers with trust_remote_code=True.15- Supports OpenAI-style chat messages with text, image_url/input_image, video_url/input_video.16- Streaming SSE supports resume (session_id + Last-Event-ID) with optional SQLite persistence.17- Auto prompt compression prevents context overflow with a simple truncate strategy.18"""19 20import os21import io22import re23import base6424import tempfile25import contextlib26from typing import Any, Dict, List, Optional, Tuple, Deque, Literal27 28from fastapi import FastAPI, HTTPException, Request, Header, Query, UploadFile, File, BackgroundTasks29from fastapi.middleware.cors import CORSMiddleware30from pydantic import BaseModel, ConfigDict, Field31from starlette.responses import JSONResponse32from fastapi.responses import StreamingResponse, Response, FileResponse33from starlette.staticfiles import StaticFiles34import json35import yaml36import threading37import time38import uuid39import sqlite340from collections import deque41import subprocess42import sys43import shutil44import asyncio45from concurrent.futures import ThreadPoolExecutor46import functools47 48# Load env49try:50 from dotenv import load_dotenv51 load_dotenv()52except Exception:53 pass54 55# Ensure HF cache dirs are relative to this project by default56ROOT_DIR = os.path.dirname(os.path.abspath(__file__))57DEFAULT_HF_CACHE = os.path.join(ROOT_DIR, "hf-cache")58if not os.getenv("HF_HOME"):59 os.environ["HF_HOME"] = DEFAULT_HF_CACHE60# Remove deprecated TRANSFORMERS_CACHE to avoid warnings61if os.getenv("TRANSFORMERS_CACHE"):62 del os.environ["TRANSFORMERS_CACHE"]63# Create directory eagerly to avoid later mkdir races64try:65 os.makedirs(os.environ["HF_HOME"], exist_ok=True)66except Exception:67 pass68 69# Optional heavy deps are imported lazily inside Engine to improve startup UX70import requests71from PIL import Image72import numpy as np73from huggingface_hub import snapshot_download, list_repo_files, hf_hub_download, get_hf_file_metadata74 75# OCR import76try:77 from rapidocr_onnxruntime import RapidOCR78except ImportError:79 RapidOCR = None80 81# Server config82PORT = int(os.getenv("PORT", "3000"))83DEFAULT_MODEL_ID = os.getenv("MODEL_REPO_ID", "unsloth/Qwen3-4B-Instruct-2507")84HF_TOKEN = os.getenv("HF_TOKEN", "").strip() or None85# Default max tokens: honor env, fallback to 4096 as previously discussed86DEFAULT_MAX_TOKENS = int(os.getenv("MAX_TOKENS", "4096"))87DEFAULT_TEMPERATURE = float(os.getenv("TEMPERATURE", "0.7"))88MAX_VIDEO_FRAMES = int(os.getenv("MAX_VIDEO_FRAMES", "16"))89DEVICE_MAP = os.getenv("DEVICE_MAP", "cpu") # Force CPU for current deployment90TORCH_DTYPE = os.getenv("TORCH_DTYPE", "float32") # float32 is faster on CPU91 92# Quantization config (BitsAndBytes) - disabled for CPU deployment93LOAD_IN_4BIT = str(os.getenv("LOAD_IN_4BIT", "0")).lower() in ("1", "true", "yes", "y") # Disabled94BNB_4BIT_COMPUTE_DTYPE = os.getenv("BNB_4BIT_COMPUTE_DTYPE", "float16")95BNB_4BIT_USE_DOUBLE_QUANT = str(os.getenv("BNB_4BIT_USE_DOUBLE_QUANT", "1")).lower() in ("1", "true", "yes", "y")96BNB_4BIT_QUANT_TYPE = os.getenv("BNB_4BIT_QUANT_TYPE", "nf4")97 98# Concurrency config99MAX_WORKERS = int(os.getenv("MAX_WORKERS", "4")) # Thread pool for concurrent processing100OCR_TIMEOUT_SECONDS = int(os.getenv("OCR_TIMEOUT_SECONDS", "120")) # 2 minute timeout for OCR101 102# Persistent session store (SQLite)103PERSIST_SESSIONS = str(os.getenv("PERSIST_SESSIONS", "0")).lower() in ("1", "true", "yes", "y")104SESSIONS_DB_PATH = os.getenv("SESSIONS_DB_PATH", "sessions.db")105SESSIONS_TTL_SECONDS = int(os.getenv("SESSIONS_TTL_SECONDS", "600"))106# Auto-cancel if all clients disconnect for duration (seconds). 0 disables it.107CANCEL_AFTER_DISCONNECT_SECONDS = int(os.getenv("CANCEL_AFTER_DISCONNECT_SECONDS", "3600"))108 109# Auto compression settings110ENABLE_AUTO_COMPRESSION = str(os.getenv("ENABLE_AUTO_COMPRESSION", "1")).lower() in ("1", "true", "yes", "y")111CONTEXT_MAX_TOKENS_AUTO = int(os.getenv("CONTEXT_MAX_TOKENS_AUTO", "0")) # 0 -> infer from model/tokenizer112CONTEXT_SAFETY_MARGIN = int(os.getenv("CONTEXT_SAFETY_MARGIN", "256"))113COMPRESSION_STRATEGY = os.getenv("COMPRESSION_STRATEGY", "truncate") # truncate | summarize (future)114 115# Eager model loading (download/check at startup before serving traffic)116EAGER_LOAD_MODEL = str(os.getenv("EAGER_LOAD_MODEL", "1")).lower() in ("1", "true", "yes", "y")117 118# Global thread pool executor for concurrent processing119executor = ThreadPoolExecutor(max_workers=MAX_WORKERS, thread_name_prefix="inference")120 121# Global OCR engine122_ocr_engine = None123 124def get_ocr_engine():125 global _ocr_engine126 if _ocr_engine is None and RapidOCR is not None:127 try:128 _ocr_engine = RapidOCR()129 print("[OCR] RapidOCR engine initialized")130 except Exception as e:131 print(f"[OCR] Failed to initialize RapidOCR: {e}")132 _ocr_engine = None133 return _ocr_engine134 135def _log(msg: str):136 # Consistent, flush-immediate startup logs137 print(f"[startup] {msg}", flush=True)138 139def prefetch_model_assets(repo_id: str, token: Optional[str]) -> Optional[str]:140 """141 Reproducible prefetch driven by huggingface-cli:142 - Downloads the ENTIRE repo using CLI (visible progress bar).143 - Returns the local directory path where the repo is mirrored.144 - If CLI is unavailable, falls back to verbose API prefetch.145 """146 try:147 # Enable accelerated transfer only if hf_transfer is installed; otherwise disable to avoid runtime errors on Spaces148 try:149 import importlib.util as _imputil150 if _imputil.find_spec("hf_transfer") is not None:151 os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")152 else:153 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"154 except Exception:155 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"156 # XET acceleration if available; harmless if missing157 os.environ.setdefault("HF_HUB_ENABLE_XET", "1")158 159 cache_dir = os.getenv("HF_HOME") or os.getenv("TRANSFORMERS_CACHE") or ""160 if cache_dir:161 os.makedirs(cache_dir, exist_ok=True)162 163 # Resolve huggingface-cli path (Windows-friendly) - try hf first, fallback to huggingface-cli164 cli_path = shutil.which("hf")165 if not cli_path:166 cli_path = shutil.which("huggingface-cli")167 if not cli_path:168 candidates = []169 appdata = os.getenv("APPDATA")170 if appdata:171 candidates.append(os.path.join(appdata, "Python", "Python312", "Scripts", "hf.exe"))172 candidates.append(os.path.join(appdata, "Python", "Python312", "Scripts", "huggingface-cli.exe"))173 candidates.append(os.path.join(os.path.dirname(sys.executable), "Scripts", "hf.exe"))174 candidates.append(os.path.join(os.path.dirname(sys.executable), "Scripts", "huggingface-cli.exe"))175 cli_path = next((p for p in candidates if os.path.exists(p)), None)176 177 # Preferred: one-shot CLI download for the whole repo (shows live progress)178 if cli_path:179 local_root = os.path.join(cache_dir if cache_dir else ".", repo_id.replace("/", "_"))180 os.makedirs(local_root, exist_ok=True)181 _log(f"Using hf download to download entire repo -> '{local_root}'")182 cmd = [183 cli_path,184 "download",185 "--repo-type",186 "model",187 "--local-dir",188 local_root,189 repo_id,190 ]191 if token:192 cmd += ["--token", token]193 # Inherit stdio; users will see a proper progress bar194 subprocess.run(cmd, check=False)195 # Verify we have the essential files196 if os.path.exists(os.path.join(local_root, "config.json")) or os.path.exists(os.path.join(local_root, "model.safetensors")):197 _log("CLI prefetch completed")198 return local_root199 else:200 _log("CLI prefetch finished but essential files not found; will fallback to API mirroring")201 202 # Fallback: verbose API-driven prefetch with per-file logging203 _log(f"Prefetching (API) repo={repo_id} to cache='{cache_dir}'")204 try:205 files = list_repo_files(repo_id, repo_type="model", token=token)206 except Exception as e:207 _log(f"list_repo_files failed ({type(e).__name__}: {e}); falling back to snapshot_download")208 snapshot_download(repo_id, token=token, local_files_only=False)209 _log("Prefetch completed (snapshot)")210 return None211 212 total = len(files)213 _log(f"Found {total} files to ensure cached (API)")214 for i, fn in enumerate(files, start=1):215 try:216 meta = get_hf_file_metadata(repo_id, fn, repo_type="model", token=token)217 size_bytes = meta.size or 0218 except Exception:219 size_bytes = 0220 size_mb = size_bytes / (1024 * 1024) if size_bytes else 0.0221 _log(f"[{i}/{total}] fetching '{fn}' (~{size_mb:.2f} MB)")222 _ = hf_hub_download(223 repo_id=repo_id,224 filename=fn,225 repo_type="model",226 token=token,227 local_files_only=False,228 resume_download=True,229 )230 _log(f"[{i}/{total}] done '{fn}'")231 _log("Prefetch completed (API)")232 return None233 except Exception as e:234 _log(f"Prefetch skipped: {type(e).__name__}: {e}")235 return None236 237def is_data_url(url: str) -> bool:238 return url.startswith("data:") and ";base64," in url239 240 241def is_http_url(url: str) -> bool:242 return url.startswith("http://") or url.startswith("https://")243 244 245def decode_base64_to_bytes(b64: str) -> bytes:246 # strip possible "data:*;base64," prefix247 if "base64," in b64:248 b64 = b64.split("base64,", 1)[1]249 return base64.b64decode(b64, validate=False)250 251 252def fetch_bytes(url: str, headers: Optional[Dict[str, str]] = None, timeout: int = 60) -> bytes:253 if not is_http_url(url):254 raise ValueError(f"Only http(s) URLs supported for fetch, got: {url}")255 resp = requests.get(url, headers=headers or {}, timeout=timeout, stream=True)256 resp.raise_for_status()257 return resp.content258 259 260def load_image_from_any(src: Dict[str, Any]) -> Image.Image:261 """262 src can be:263 - { "url": "http(s)://..." } (also supports data URL)264 - { "b64_json": "<base64>" }265 - { "path": "local_path" } (optional)266 """267 if "b64_json" in src and src["b64_json"]:268 data = decode_base64_to_bytes(str(src["b64_json"]))269 return Image.open(io.BytesIO(data)).convert("RGB")270 271 if "url" in src and src["url"]:272 url = str(src["url"])273 if is_data_url(url):274 data = decode_base64_to_bytes(url)275 return Image.open(io.BytesIO(data)).convert("RGB")276 if is_http_url(url):277 data = fetch_bytes(url)278 return Image.open(io.BytesIO(data)).convert("RGB")279 # treat as local path280 if os.path.exists(url):281 with open(url, "rb") as f:282 return Image.open(io.BytesIO(f.read())).convert("RGB")283 raise ValueError(f"Invalid image url/path: {url}")284 285 if "path" in src and src["path"]:286 p = str(src["path"])287 if os.path.exists(p):288 with open(p, "rb") as f:289 return Image.open(io.BytesIO(f.read())).convert("RGB")290 raise ValueError(f"Image path not found: {p}")291 292 raise ValueError("Unsupported image source payload")293 294 295def write_bytes_tempfile(data: bytes, suffix: str) -> str:296 tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)297 with tmp as f:298 f.write(data)299 return tmp.name300 301 302def load_video_frames_from_any(src: Dict[str, Any], max_frames: int = MAX_VIDEO_FRAMES) -> List[Image.Image]:303 """304 Returns a list of PIL.Image frames (RGB) sampled up to max_frames.305 src can be:306 - { "url": "http(s)://..." } (mp4/mov/webm/etc.)307 - { "b64_json": "<base64 of a video file>" }308 - { "path": "local_path" }309 """310 # Prefer imageio.v3 if present, fallback to OpenCV311 # We load all frames then uniform sample if too many.312 def _load_all_frames(path: str) -> List[Image.Image]:313 frames: List[Image.Image] = []314 with contextlib.suppress(ImportError):315 import imageio.v3 as iio316 arr_iter = iio.imiter(path) # yields numpy arrays HxWxC317 for arr in arr_iter:318 if arr is None:319 continue320 if arr.ndim == 2:321 arr = np.stack([arr, arr, arr], axis=-1)322 if arr.shape[-1] == 4:323 arr = arr[..., :3]324 frames.append(Image.fromarray(arr).convert("RGB"))325 return frames326 327 # Fallback to OpenCV328 import cv2 # type: ignore329 cap = cv2.VideoCapture(path)330 ok, frame = cap.read()331 while ok:332 frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)333 frames.append(Image.fromarray(frame))334 ok, frame = cap.read()335 cap.release()336 return frames337 338 # Resolve to a local path339 local_path = None340 if "b64_json" in src and src["b64_json"]:341 data = decode_base64_to_bytes(str(src["b64_json"]))342 local_path = write_bytes_tempfile(data, suffix=".mp4")343 elif "url" in src and src["url"]:344 url = str(src["url"])345 if is_data_url(url):346 data = decode_base64_to_bytes(url)347 local_path = write_bytes_tempfile(data, suffix=".mp4")348 elif is_http_url(url):349 data = fetch_bytes(url)350 local_path = write_bytes_tempfile(data, suffix=".mp4")351 elif os.path.exists(url):352 local_path = url353 else:354 raise ValueError(f"Invalid video url/path: {url}")355 elif "path" in src and src["path"]:356 p = str(src["path"])357 if os.path.exists(p):358 local_path = p359 else:360 raise ValueError(f"Video path not found: {p}")361 else:362 raise ValueError("Unsupported video source payload")363 364 frames = _load_all_frames(local_path)365 # Uniform sample if too many frames366 if len(frames) > max_frames and max_frames > 0:367 idxs = np.linspace(0, len(frames) - 1, max_frames).astype(int).tolist()368 frames = [frames[i] for i in idxs]369 return frames370 371 372class ChatRequest(BaseModel):373 """OpenAI-compatible Chat Completions request body."""374 model: Optional[str] = Field(default=None, description="Model id (defaults to env MODEL_REPO_ID).")375 messages: List[Dict[str, Any]] = Field(description="OpenAI-style messages array. Supports text, image_url/input_image, video_url/input_video parts.")376 max_tokens: Optional[int] = Field(default=None, description="Max new tokens to generate.")377 temperature: Optional[float] = Field(default=None, description="Sampling temperature.")378 stream: Optional[bool] = Field(default=None, description="When true, returns Server-Sent Events stream.")379 session_id: Optional[str] = Field(default=None, description="Optional session id for resumable SSE.")380 # Pydantic v2 schema extras with rich examples381 model_config = ConfigDict(382 json_schema_extra={383 "examples": [384 {385 "summary": "Text-only",386 "value": {387 "messages": [388 {"role": "user", "content": "Hello, summarize the benefits of multimodal LLMs."}389 ],390 "max_tokens": 128391 }392 },393 {394 "summary": "Image by URL",395 "value": {396 "messages": [397 {398 "role": "user",399 "content": [400 {"type": "text", "text": "What is in this image?"},401 {"type": "image_url", "image_url": {"url": "https://example.com/cat.jpg"}}402 ]403 }404 ],405 "max_tokens": 128406 }407 },408 {409 "summary": "Video by URL (streaming SSE)",410 "value": {411 "messages": [412 {413 "role": "user",414 "content": [415 {"type": "text", "text": "Describe this clip briefly."},416 {"type": "video_url", "video_url": {"url": "https://example.com/clip.mp4"}}417 ]418 }419 ],420 "stream": True,421 "max_tokens": 128422 }423 }424 ]425 }426 )427 428class MessageModel(BaseModel):429 role: Literal["system", "user", "assistant"]430 content: str431 432class ChoiceModel(BaseModel):433 index: int434 message: MessageModel435 finish_reason: Optional[str] = None436 437class UsageModel(BaseModel):438 prompt_tokens: int439 completion_tokens: int440 total_tokens: int441 442class ChatCompletionResponse(BaseModel):443 """Non-streaming Chat Completions response (when stream=false)."""444 id: str445 object: str446 created: int447 model: str448 choices: List[ChoiceModel]449 usage: UsageModel450 context: Dict[str, Any] = {}451 452class HealthResponse(BaseModel):453 ok: bool454 modelReady: bool455 modelId: str456 error: Optional[str] = None457 context: Optional[Dict[str, Any]] = None458 459class CancelResponse(BaseModel):460 ok: bool461 session_id: str462 463 464class Engine:465 def __init__(self, model_id: str, hf_token: Optional[str] = None):466 # Lazy import heavy deps467 from transformers import AutoProcessor, AutoModelForCausalLM, AutoModelForVision2Seq, AutoModel, BitsAndBytesConfig468 # AutoModelForImageTextToText is the v5+ replacement for Vision2Seq in Transformers469 try:470 from transformers import AutoModelForImageTextToText # type: ignore471 except Exception:472 AutoModelForImageTextToText = None # type: ignore473 474 # Resolve device map to avoid 'meta' device on CPU Spaces475 # If DEVICE_MAP is "auto" but no CUDA is available, force "cpu" and disable low_cpu_mem_usage476 model_kwargs: Dict[str, Any] = {477 "trust_remote_code": True,478 }479 if hf_token:480 # Only pass 'token' (use_auth_token is deprecated and causes conflicts)481 model_kwargs["token"] = hf_token482 483 # Add quantization config if enabled484 if LOAD_IN_4BIT:485 try:486 import torch487 compute_dtype = getattr(torch, BNB_4BIT_COMPUTE_DTYPE, torch.float16)488 quant_config = BitsAndBytesConfig(489 load_in_4bit=True,490 bnb_4bit_compute_dtype=compute_dtype,491 bnb_4bit_use_double_quant=BNB_4BIT_USE_DOUBLE_QUANT,492 bnb_4bit_quant_type=BNB_4BIT_QUANT_TYPE,493 )494 model_kwargs["quantization_config"] = quant_config495 _log(f"Using 4-bit quantization: {BNB_4BIT_QUANT_TYPE}, compute_dtype={BNB_4BIT_COMPUTE_DTYPE}, double_quant={BNB_4BIT_USE_DOUBLE_QUANT}")496 except Exception as e:497 _log(f"BitsAndBytes quantization failed: {e}; falling back to full precision")498 499 # Device and dtype resolution500 try:501 import torch # local import to avoid heavy import at module load502 has_cuda = bool(getattr(torch, "cuda", None) and torch.cuda.is_available())503 except Exception:504 has_cuda = False505 506 resolved_device_map = DEVICE_MAP507 if str(DEVICE_MAP).lower() == "auto" and not has_cuda:508 resolved_device_map = "cpu"509 510 model_kwargs["device_map"] = resolved_device_map511 # Explicitly disable low_cpu_mem_usage on pure CPU to fully materialize weights (avoids meta tensors)512 if resolved_device_map == "cpu":513 model_kwargs["low_cpu_mem_usage"] = False514 # dtype - use 'dtype' instead of deprecated 'torch_dtype'515 if TORCH_DTYPE != "auto":516 try:517 import torch518 model_kwargs["dtype"] = getattr(torch, TORCH_DTYPE, TORCH_DTYPE)519 except Exception:520 model_kwargs["dtype"] = TORCH_DTYPE521 else:522 model_kwargs["dtype"] = "auto"523 # store for later524 self._resolved_device_map = resolved_device_map525 526 # Processor (handles text + images/videos)527 proc_kwargs: Dict[str, Any] = {"trust_remote_code": True}528 if hf_token:529 proc_kwargs["token"] = hf_token530 self.processor = AutoProcessor.from_pretrained(531 model_id,532 **proc_kwargs,533 ) # pragma: no cover534 535 # Prefer ImageTextToText (Transformers v5 path), then Vision2Seq, then CausalLM as a last resort536 model = None537 if 'AutoModelForImageTextToText' in globals() and AutoModelForImageTextToText is not None:538 try:539 model = AutoModelForImageTextToText.from_pretrained(model_id, **model_kwargs) # pragma: no cover540 except Exception:541 model = None542 if model is None:543 try:544 # AutoModelForVision2Seq is deprecated, but try it for compatibility545 model = AutoModelForVision2Seq.from_pretrained(model_id, **model_kwargs) # pragma: no cover546 except Exception:547 model = None548 if model is None:549 try:550 model = AutoModelForCausalLM.from_pretrained(model_id, **model_kwargs) # pragma: no cover551 except Exception:552 model = None553 if model is None:554 # Generic AutoModel as last-resort with trust_remote_code to load custom architectures555 model = AutoModel.from_pretrained(model_id, **model_kwargs) # pragma: no cover556 self.model = model.eval() # pragma: no cover557 # Ensure model is fully on CPU when resolved device_map is cpu (prevents meta device mix during inference)558 try:559 if str(getattr(self, "_resolved_device_map", "")).lower() == "cpu":560 _ = self.model.to("cpu")561 except Exception:562 pass563 # Ensure model is on CPU when resolved device_map is cpu (prevents meta device mix during inference)564 try:565 if getattr(self, "_resolved_device_map", None) == "cpu":566 _ = self.model.to("cpu")567 except Exception:568 pass569 570 self.model_id = model_id571 self.tokenizer = getattr(self.processor, "tokenizer", None)572 self.last_context_info: Dict[str, Any] = {}573 574 def _model_max_context(self) -> int:575 try:576 cfg = getattr(self.model, "config", None)577 if cfg is not None:578 v = getattr(cfg, "max_position_embeddings", None)579 if isinstance(v, int) and v > 0 and v < 10_000_000:580 return v581 except Exception:582 pass583 try:584 mx = int(getattr(self.tokenizer, "model_max_length", 0) or 0)585 if mx > 0 and mx < 10_000_000_000:586 return mx587 except Exception:588 pass589 return 32768590 591 def _count_prompt_tokens(self, text: str) -> int:592 try:593 if self.tokenizer is not None:594 enc = self.tokenizer([text], add_special_tokens=False, return_attention_mask=False)595 ids = enc["input_ids"][0]596 return len(ids)597 except Exception:598 pass599 return max(1, int(len(text.split()) * 1.3))600 601 def _auto_compress_if_needed(602 self, mm_messages: List[Dict[str, Any]], max_new_tokens: int603 ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]:604 info: Dict[str, Any] = {}605 # Build once to measure606 text0 = self.processor.apply_chat_template(mm_messages, tokenize=False, add_generation_prompt=True)607 prompt_tokens = self._count_prompt_tokens(text0)608 max_ctx = CONTEXT_MAX_TOKENS_AUTO if CONTEXT_MAX_TOKENS_AUTO > 0 else self._model_max_context()609 budget = max(1024, max_ctx - CONTEXT_SAFETY_MARGIN - int(max_new_tokens))610 if not ENABLE_AUTO_COMPRESSION or prompt_tokens <= budget:611 info = {612 "compressed": False,613 "prompt_tokens": int(prompt_tokens),614 "max_context": int(max_ctx),615 "budget": int(budget),616 "strategy": COMPRESSION_STRATEGY,617 "dropped_messages": 0,618 }619 return mm_messages, info620 621 # Truncate earliest non-system messages until within budget622 msgs = list(mm_messages)623 dropped = 0624 guard = 0625 while True:626 text = self.processor.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)627 prompt_tokens = self._count_prompt_tokens(text)628 if prompt_tokens <= budget or len(msgs) <= 1:629 break630 # drop earliest non-system631 drop_idx = None632 for j, m in enumerate(msgs):633 if (m.get("role") or "user") != "system":634 drop_idx = j635 break636 if drop_idx is None:637 break638 msgs.pop(drop_idx)639 dropped += 1640 guard += 1641 if guard > 10000:642 break643 644 info = {645 "compressed": True,646 "prompt_tokens": int(prompt_tokens),647 "max_context": int(max_ctx),648 "budget": int(budget),649 "strategy": "truncate",650 "dropped_messages": int(dropped),651 }652 return msgs, info653 654 def get_context_report(self) -> Dict[str, Any]:655 try:656 tk_max = int(getattr(self.tokenizer, "model_max_length", 0) or 0)657 except Exception:658 tk_max = 0659 return {660 "compressionEnabled": ENABLE_AUTO_COMPRESSION,661 "strategy": COMPRESSION_STRATEGY,662 "safetyMargin": CONTEXT_SAFETY_MARGIN,663 "modelMaxContext": self._model_max_context(),664 "tokenizerModelMaxLength": tk_max,665 "last": self.last_context_info or {},666 }667 668 def build_mm_messages(669 self, openai_messages: List[Dict[str, Any]]670 ) -> Tuple[List[Dict[str, Any]], List[Image.Image], List[List[Image.Image]]]:671 """672 Convert OpenAI-style messages to Qwen multimodal messages.673 Returns:674 - messages for apply_chat_template675 - flat list of images in encounter order676 - list of videos (each is list of PIL frames)677 """678 mm_msgs: List[Dict[str, Any]] = []679 images: List[Image.Image] = []680 videos: List[List[Image.Image]] = []681 682 for msg in openai_messages:683 role = msg.get("role", "user")684 content = msg.get("content", "")685 686 parts: List[Dict[str, Any]] = []687 688 if isinstance(content, str):689 if content:690 parts.append({"type": "text", "text": content})691 elif isinstance(content, list):692 for p in content:693 ptype = p.get("type")694 if ptype == "text":695 txt = p.get("text", "")696 if txt:697 parts.append({"type": "text", "text": txt})698 elif ptype in ("image_url", "input_image"):699 src: Dict[str, Any] = {}700 if ptype == "image_url":701 u = (p.get("image_url") or {}).get("url") if isinstance(p.get("image_url"), dict) else p.get("image_url")702 src["url"] = u703 else:704 b64 = p.get("image") or p.get("b64_json") or p.get("data") or (p.get("image_url") or {}).get("url")705 if b64:706 src["b64_json"] = b64707 try:708 img = load_image_from_any(src)709 images.append(img)710 parts.append({"type": "image", "image": img})711 except Exception as e:712 raise ValueError(f"Failed to parse image part: {e}") from e713 elif ptype in ("video_url", "input_video"):714 src = {}715 if ptype == "video_url":716 u = (p.get("video_url") or {}).get("url") if isinstance(p.get("video_url"), dict) else p.get("video_url")717 src["url"] = u718 else:719 b64 = p.get("video") or p.get("b64_json") or p.get("data")720 if b64:721 src["b64_json"] = b64722 try:723 frames = load_video_frames_from_any(src, max_frames=MAX_VIDEO_FRAMES)724 videos.append(frames)725 parts.append({"type": "video", "video": frames})726 except Exception as e:727 raise ValueError(f"Failed to parse video part: {e}") from e728 else:729 if isinstance(p, dict):730 txt = p.get("text")731 if isinstance(txt, str) and txt:732 parts.append({"type": "text", "text": txt})733 else:734 if content:735 parts.append({"type": "text", "text": str(content)})736 737 mm_msgs.append({"role": role, "content": parts})738 739 return mm_msgs, images, videos740 741 def infer(self, messages: List[Dict[str, Any]], max_tokens: int, temperature: float) -> str:742 mm_messages, images, videos = self.build_mm_messages(messages)743 # Auto-compress if needed based on context budget744 mm_messages, ctx_info = self._auto_compress_if_needed(mm_messages, max_tokens)745 self.last_context_info = ctx_info746 747 # Build chat template748 text = self.processor.apply_chat_template(749 mm_messages,750 tokenize=False,751 add_generation_prompt=True,752 )753 754 proc_kwargs: Dict[str, Any] = {"text": [text], "return_tensors": "pt"}755 if images:756 proc_kwargs["images"] = images757 if videos:758 proc_kwargs["videos"] = videos759 760 inputs = self.processor(**proc_kwargs)761 # Move tensors to the correct device762 try:763 if str(getattr(self, "_resolved_device_map", "")).lower() == "cpu":764 # Explicit CPU placement avoids 'meta' device errors on Spaces765 inputs = {k: (v.to("cpu") if hasattr(v, "to") else v) for k, v in inputs.items()}766 else:767 device = getattr(self.model, "device", None) or next(self.model.parameters()).device768 inputs = {k: (v.to(device) if hasattr(v, "to") else v) for k, v in inputs.items()}769 except Exception:770 pass771 772 do_sample = temperature is not None and float(temperature) > 0.0773 774 # Safer on CPU: run without gradients to reduce memory pressure and avoid autograd hooks775 try:776 import torch777 with torch.no_grad():778 gen_ids = self.model.generate(779 **inputs,780 max_new_tokens=int(max_tokens),781 temperature=float(temperature),782 do_sample=do_sample,783 use_cache=True,784 )785 except Exception:786 # Fallback without no_grad if torch import fails (very unlikely)787 gen_ids = self.model.generate(788 **inputs,789 max_new_tokens=int(max_tokens),790 temperature=float(temperature),791 do_sample=do_sample,792 use_cache=True,793 )794 795 # Decode796 output = self.processor.batch_decode(797 gen_ids,798 skip_special_tokens=True,799 clean_up_tokenization_spaces=False,800 )[0]801 802 # Best-effort: return only the assistant reply after the last template marker if present803 parts = re.split(r"\n?assistant:\s*", output, flags=re.IGNORECASE)804 if len(parts) >= 2:805 return parts[-1].strip()806 return output.strip()807 808 def infer_stream(809 self,810 messages: List[Dict[str, Any]],811 max_tokens: int,812 temperature: float,813 cancel_event: Optional[threading.Event] = None,814 ):815 from transformers import TextIteratorStreamer, StoppingCriteria, StoppingCriteriaList816 817 mm_messages, images, videos = self.build_mm_messages(messages)818 # Auto-compress if needed based on context budget819 mm_messages, ctx_info = self._auto_compress_if_needed(mm_messages, max_tokens)820 self.last_context_info = ctx_info821 822 text = self.processor.apply_chat_template(823 mm_messages,824 tokenize=False,825 add_generation_prompt=True,826 )827 828 proc_kwargs: Dict[str, Any] = {"text": [text], "return_tensors": "pt"}829 if images:830 proc_kwargs["images"] = images831 if videos:832 proc_kwargs["videos"] = videos833 834 inputs = self.processor(**proc_kwargs)835 try:836 if str(getattr(self, "_resolved_device_map", "")).lower() == "cpu":837 inputs = {k: (v.to("cpu") if hasattr(v, "to") else v) for k, v in inputs.items()}838 else:839 device = getattr(self.model, "device", None) or next(self.model.parameters()).device840 inputs = {k: (v.to(device) if hasattr(v, "to") else v) for k, v in inputs.items()}841 except Exception:842 pass843 844 do_sample = temperature is not None and float(temperature) > 0.0845 846 streamer = TextIteratorStreamer(847 getattr(self.processor, "tokenizer", None),848 skip_prompt=True,849 skip_special_tokens=True,850 )851 852 gen_kwargs = dict(853 **inputs,854 max_new_tokens=int(max_tokens),855 temperature=float(temperature),856 do_sample=do_sample,857 use_cache=True,858 streamer=streamer,859 )860 861 # Optional cooperative cancellation via StoppingCriteria862 if cancel_event is not None:863 class _CancelCrit(StoppingCriteria):864 def __init__(self, ev: threading.Event):865 self.ev = ev866 867 def __call__(self, input_ids, scores, **kwargs):868 return bool(self.ev.is_set())869 870 gen_kwargs["stopping_criteria"] = StoppingCriteriaList([_CancelCrit(cancel_event)])871 872 # Wrap generation with torch.no_grad() to avoid autograd overhead on CPU and reduce failure surface873 def _runner():874 try:875 import torch876 with torch.no_grad():877 self.model.generate(**gen_kwargs)878 except Exception:879 # Let streamer finish gracefully even if generation throws880 pass881 882 th = threading.Thread(target=_runner)883 th.start()884 885 for piece in streamer:886 if piece:887 yield piece888 889 890# Simple in-memory resumable SSE session store + optional SQLite persistence891class _SSESession:892 def __init__(self, maxlen: int = 2048, ttl_seconds: int = 600):893 self.buffer: Deque[Tuple[int, str]] = deque(maxlen=maxlen) # (idx, sse_line_block)894 self.last_idx: int = -1895 self.created: float = time.time()896 self.finished: bool = False897 self.cond = threading.Condition()898 self.thread: Optional[threading.Thread] = None899 self.ttl_seconds = ttl_seconds900 # Cancellation + client tracking901 self.cancel_event = threading.Event()902 self.listeners: int = 0903 self.cancel_timer = None # type: ignore904 905 906class _SessionStore:907 def __init__(self, ttl_seconds: int = 600, max_sessions: int = 256):908 self._sessions: Dict[str, _SSESession] = {}909 self._lock = threading.Lock()910 self._ttl = ttl_seconds911 self._max_sessions = max_sessions912 913 def get_or_create(self, sid: str) -> _SSESession:914 with self._lock:915 sess = self._sessions.get(sid)916 if sess is None:917 sess = _SSESession(ttl_seconds=self._ttl)918 self._sessions[sid] = sess919 return sess920 921 def get(self, sid: str) -> Optional[_SSESession]:922 with self._lock:923 return self._sessions.get(sid)924 925 def gc(self):926 now = time.time()927 with self._lock:928 # remove expired929 expired = [k for k, v in self._sessions.items() if (now - v.created) > self._ttl or (v.finished and (now - v.created) > self._ttl / 4)]930 for k in expired:931 self._sessions.pop(k, None)932 # bound session count933 if len(self._sessions) > self._max_sessions:934 for k, _ in sorted(self._sessions.items(), key=lambda kv: kv[1].created)[: max(0, len(self._sessions) - self._max_sessions)]:935 self._sessions.pop(k, None)936 937 938class _SQLiteStore:939 def __init__(self, db_path: str):940 self.db_path = db_path941 self._lock = threading.Lock()942 self._conn = sqlite3.connect(self.db_path, check_same_thread=False)943 self._conn.execute("PRAGMA journal_mode=WAL;")944 self._conn.execute("PRAGMA synchronous=NORMAL;")945 self._ensure_schema()946 947 def _ensure_schema(self):948 cur = self._conn.cursor()949 cur.execute(950 "CREATE TABLE IF NOT EXISTS sessions (session_id TEXT PRIMARY KEY, created REAL, finished INTEGER DEFAULT 0)"951 )952 cur.execute(953 "CREATE TABLE IF NOT EXISTS events (session_id TEXT, idx INTEGER, data TEXT, created REAL, PRIMARY KEY(session_id, idx))"954 )955 cur.execute("CREATE INDEX IF NOT EXISTS idx_events_session ON events(session_id, idx)")956 self._conn.commit()957 958 def ensure_session(self, session_id: str, created: int):959 with self._lock:960 self._conn.execute(961 "INSERT OR IGNORE INTO sessions(session_id, created, finished) VALUES (?, ?, 0)",962 (session_id, float(created)),963 )964 self._conn.commit()965 966 def append_event(self, session_id: str, idx: int, payload: Dict[str, Any]):967 data = json.dumps(payload, ensure_ascii=False)968 with self._lock:969 self._conn.execute(970 "INSERT OR REPLACE INTO events(session_id, idx, data, created) VALUES (?, ?, ?, ?)",971 (session_id, idx, data, time.time()),972 )973 self._conn.commit()974 975 def get_events_after(self, session_id: str, last_idx: int) -> List[Tuple[int, str]]:976 with self._lock:977 cur = self._conn.execute(978 "SELECT idx, data FROM events WHERE session_id=? AND idx>? ORDER BY idx ASC", (session_id, last_idx)979 )980 return [(int(r[0]), str(r[1])) for r in cur.fetchall()]981 982 def mark_finished(self, session_id: str):983 with self._lock:984 self._conn.execute("UPDATE sessions SET finished=1 WHERE session_id=?", (session_id,))985 self._conn.commit()986 987 def session_meta(self, session_id: str) -> Tuple[bool, int]:988 with self._lock:989 row = self._conn.execute("SELECT finished FROM sessions WHERE session_id=?", (session_id,)).fetchone()990 finished = bool(row[0]) if row else False991 row2 = self._conn.execute("SELECT MAX(idx) FROM events WHERE session_id=?", (session_id,)).fetchone()992 last_idx = int(row2[0]) if row2 and row2[0] is not None else -1993 return finished, last_idx994 995 def gc(self, ttl_seconds: int):996 cutoff = time.time() - float(ttl_seconds)997 with self._lock:998 cur = self._conn.execute("SELECT session_id FROM sessions WHERE finished=1 AND created<?", (cutoff,))999 ids = [r[0] for r in cur.fetchall()]1000 for sid in ids:1001 self._conn.execute("DELETE FROM events WHERE session_id=?", (sid,))1002 self._conn.execute("DELETE FROM sessions WHERE session_id=?", (sid,))1003 self._conn.commit()1004 1005 1006def _sse_event(session_id: str, idx: int, payload: Dict[str, Any]) -> str:1007 # Include SSE id line so clients can send Last-Event-ID to resume.1008 return f"id: {session_id}:{idx}\n" + f"data: {json.dumps(payload, ensure_ascii=False)}\n\n"1009 1010 1011_STORE = _SessionStore()1012_DB_STORE = _SQLiteStore(SESSIONS_DB_PATH) if PERSIST_SESSIONS else None1013 1014# FastAPI app and OpenAPI tags1015tags_metadata = [1016 {"name": "meta", "description": "Service metadata and OpenAPI schema"},1017 {"name": "health", "description": "Readiness and runtime info including context window report"},1018 {"name": "chat", "description": "OpenAI-compatible chat completions (non-stream and streaming SSE)"},1019 {"name": "ocr", "description": "Optical Character Recognition endpoints"},1020]1021 1022app = FastAPI(1023 title="Qwen3-VL Inference Server",1024 version="1.0.0",1025 description="OpenAI-compatible inference server for Qwen3-VL with multimodal support, streaming SSE with resume, context auto-compression, and optional SQLite persistence.",1026 openapi_tags=tags_metadata,1027)1028app.add_middleware(1029 CORSMiddleware,1030 allow_origins=["*"],1031 allow_methods=["*"],1032 allow_headers=["*"],1033)1034 1035# Startup hook is defined after get_engine() so globals are initialized first.1036# Serve static web UI if present1037_WEB_DIR = os.path.join(ROOT_DIR, "web")1038if os.path.isdir(_WEB_DIR):1039 try:1040 app.mount("/web", StaticFiles(directory=_WEB_DIR, html=True), name="web")1041 except Exception:1042 pass1043 1044# Engine singletons1045_engine: Optional[Engine] = None1046_engine_error: Optional[str] = None1047 1048 1049def get_engine() -> Engine:1050 global _engine, _engine_error1051 if _engine is not None:1052 return _engine1053 try:1054 model_id = DEFAULT_MODEL_ID1055 _log(f"Preparing model '{model_id}' (HF_HOME={os.getenv('HF_HOME')}, cache={os.getenv('TRANSFORMERS_CACHE')})")1056 local_repo_dir = prefetch_model_assets(model_id, HF_TOKEN)1057 load_id = local_repo_dir if (local_repo_dir and os.path.exists(os.path.join(local_repo_dir, 'config.json'))) else model_id1058 _log(f"Loading processor and model from: {load_id}")1059 _engine = Engine(model_id=load_id, hf_token=HF_TOKEN)1060 _engine_error = None1061 _log(f"Model ready: {_engine.model_id}")1062 return _engine1063 except Exception as e:1064 _engine_error = f"{type(e).__name__}: {e}"1065 _log(f"Engine init failed: {_engine_error}")1066 raise1067 1068# Eager-load model at startup after definitions so it downloads/checks before serving traffic.1069@app.on_event("startup")1070def _startup_load_model():1071 # Initialize marketplace database1072 try:1073 from database import init_db1074 init_db()1075 print("[startup] Marketplace database initialized")1076 except Exception as e:1077 print(f"[startup] Database initialization failed: {e}")1078 1079 if EAGER_LOAD_MODEL:1080 print("[startup] EAGER_LOAD_MODEL=1: initializing model and OCR engine...")1081 try:1082 # Initialize OCR engine first1083 _ = get_ocr_engine()1084 print("[startup] OCR engine initialized")1085 1086 # Then initialize the model1087 _ = get_engine()1088 print("[startup] Model loaded:", _engine.model_id if _engine else "unknown")1089 except Exception as e:1090 # Log error but don't fail - allow server to start without model1091 print("[startup] Initialization failed:", e)1092 print("[startup] Server will start without full initialization")1093 else:1094 print("[startup] EAGER_LOAD_MODEL=0: skipping initialization")1095 1096 1097@app.get("/", tags=["meta"], include_in_schema=False)1098def root():1099 """1100 Serve the client web UI. The UI calls an external Hugging Face Space API1101 (default is KillerKing93/Transformers-InferenceServer-OpenAPI) and does NOT1102 use internal server endpoints for chat. You can change the base via the input1103 field or ?api= query string in the page.1104 """1105 index_path = os.path.join(ROOT_DIR, "web", "index.html")1106 if os.path.exists(index_path):1107 return FileResponse(index_path, media_type="text/html; charset=utf-8")1108 # Inline minimal fallback to make root return an HTML page even if COPY failed1109 html = """<!doctype html><html><head><meta charset='utf-8'><title>Qwen3‑VL Chat</title></head>1110 <body style="font-family:system-ui,Segoe UI,Roboto;padding:24px;background:#0f172a;color:#e2e8f0">1111 <h2>Qwen3‑VL Chat UI</h2>1112 <p>The static UI was not found inside the container. This page is a fallback.</p>1113 <p>Try pulling the latest image or rebuilding the Space so that <code>/app/web/index.html</code> is present.</p>1114 <p>Once copied, this URL will serve the full UI. For now you can open the raw UI file from the repo or call the API directly.</p>1115 <ul>1116 <li><a href="./docs" style="color:#93c5fd">Swagger UI</a></li>1117 <li><a href="./openapi.yaml" style="color:#93c5fd">OpenAPI YAML</a></li>1118 </ul>1119 </body></html>"""1120 return Response(html, media_type="text/html; charset=utf-8")1121 1122 1123@app.get("/openapi.yaml", tags=["meta"])1124def openapi_yaml():1125 """Serve OpenAPI schema as YAML for tooling compatibility."""1126 schema = app.openapi()1127 yml = yaml.safe_dump(schema, sort_keys=False)1128 return Response(yml, media_type="application/yaml")1129 1130 1131@app.get("/health", tags=["health"], response_model=HealthResponse)1132def health():1133 ready = False1134 err = None1135 model_id = DEFAULT_MODEL_ID1136 global _engine, _engine_error1137 if _engine is not None:1138 ready = True1139 model_id = _engine.model_id1140 elif _engine_error:1141 err = _engine_error1142 ctx = None1143 try:1144 if _engine is not None:1145 ctx = _engine.get_context_report()1146 except Exception:1147 ctx = None1148 return JSONResponse({"ok": True, "modelReady": ready, "modelId": model_id, "error": err, "context": ctx})1149 1150 1151@app.post(1152 "/v1/chat/completions",1153 tags=["chat"],1154 response_model=ChatCompletionResponse,1155 responses={1156 200: {1157 "description": "When stream=true, the response is text/event-stream (SSE). When stream=false, JSON body matches ChatCompletionResponse.",1158 "content": {1159 "text/event-stream": {1160 "schema": {"type": "string"},1161 "examples": {1162 "sse": {1163 "summary": "SSE stream example",1164 "value": "id: sess-123:0\ndata: {\"id\":\"sess-123\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\"}}]}\n\n"1165 }1166 }1167 }1168 },1169 }1170 },1171)1172async def chat_completions(1173 request: Request,1174 body: ChatRequest,1175 last_event_id: Optional[str] = Query(default=None, alias="last_event_id", description="Resume SSE from this id: 'session_id:index'"),1176 last_event_id_header: Optional[str] = Header(default=None, alias="Last-Event-ID", convert_underscores=False, description="SSE resume id 'session_id:index'"),1177):1178 # Ensure engine is loaded1179 try:1180 engine = get_engine()1181 except Exception as e:1182 raise HTTPException(status_code=503, detail=f"Model not ready: {e}")1183 1184 if not body or not isinstance(body.messages, list) or len(body.messages) == 0:1185 raise HTTPException(status_code=400, detail="messages must be a non-empty array")1186 1187 max_tokens = int(body.max_tokens) if isinstance(body.max_tokens, int) else DEFAULT_MAX_TOKENS1188 temperature = float(body.temperature) if body.temperature is not None else DEFAULT_TEMPERATURE1189 do_stream = bool(body.stream)1190 1191 # Parse Last-Event-ID (header or ?last_event_id) and derive/align session_id1192 le_id = last_event_id_header or last_event_id1193 sid_from_header: Optional[str] = None1194 last_idx_from_header: int = -11195 if le_id:1196 try:1197 sid_from_header, idx_str = le_id.split(":", 1)1198 last_idx_from_header = int(idx_str)1199 except Exception:1200 sid_from_header = None