KillerKing93/Transformers-TextEngine-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 datetime import datetime27from typing import Any, Dict, List, Optional, Tuple, Deque, Literal28 29from fastapi import FastAPI, HTTPException, Request, Header, Query, UploadFile, File, BackgroundTasks, Depends30from fastapi.middleware.cors import CORSMiddleware31from pydantic import BaseModel, ConfigDict, Field32from starlette.responses import JSONResponse33from fastapi.responses import StreamingResponse, Response, FileResponse34from starlette.staticfiles import StaticFiles35import json36import yaml37import threading38import time39import uuid40import sqlite341from collections import deque42import subprocess43import sys44import shutil45import asyncio46from concurrent.futures import ThreadPoolExecutor47import functools48 49# Load env50try:51 from dotenv import load_dotenv52 load_dotenv()53except Exception:54 pass55 56# Ensure HF cache dirs are relative to this project by default57ROOT_DIR = os.path.dirname(os.path.abspath(__file__))58DEFAULT_HF_CACHE = os.path.join(ROOT_DIR, "hf-cache")59if not os.getenv("HF_HOME"):60 os.environ["HF_HOME"] = DEFAULT_HF_CACHE61# Remove deprecated TRANSFORMERS_CACHE to avoid warnings62if os.getenv("TRANSFORMERS_CACHE"):63 del os.environ["TRANSFORMERS_CACHE"]64# Create directory eagerly to avoid later mkdir races65try:66 os.makedirs(os.environ["HF_HOME"], exist_ok=True)67except Exception:68 pass69 70# Optional heavy deps are imported lazily inside Engine to improve startup UX71import requests72from PIL import Image73import numpy as np74from huggingface_hub import snapshot_download, list_repo_files, hf_hub_download, get_hf_file_metadata75 76# OCR import (Disabled - use image captioning instead)77# try:78# from rapidocr_onnxruntime import RapidOCR79# except ImportError:80# RapidOCR = None81RapidOCR = None # Disabled82 83# Server config84PORT = int(os.getenv("PORT", "3000"))85DEFAULT_MODEL_ID = os.getenv("MODEL_REPO_ID", "unsloth/Qwen3-4B-Instruct-2507")86HF_TOKEN = os.getenv("HF_TOKEN", "").strip() or None87# Default max tokens: honor env, fallback to 4096 as previously discussed88DEFAULT_MAX_TOKENS = int(os.getenv("MAX_TOKENS", "4096"))89DEFAULT_TEMPERATURE = float(os.getenv("TEMPERATURE", "0.7"))90MAX_VIDEO_FRAMES = int(os.getenv("MAX_VIDEO_FRAMES", "16"))91DEVICE_MAP = os.getenv("DEVICE_MAP", "cpu") # Force CPU for current deployment92TORCH_DTYPE = os.getenv("TORCH_DTYPE", "float32") # float32 is faster on CPU93 94# Quantization config (BitsAndBytes) - disabled for CPU deployment95LOAD_IN_4BIT = str(os.getenv("LOAD_IN_4BIT", "0")).lower() in ("1", "true", "yes", "y") # Disabled96BNB_4BIT_COMPUTE_DTYPE = os.getenv("BNB_4BIT_COMPUTE_DTYPE", "float16")97BNB_4BIT_USE_DOUBLE_QUANT = str(os.getenv("BNB_4BIT_USE_DOUBLE_QUANT", "1")).lower() in ("1", "true", "yes", "y")98BNB_4BIT_QUANT_TYPE = os.getenv("BNB_4BIT_QUANT_TYPE", "nf4")99 100# Concurrency config101MAX_WORKERS = int(os.getenv("MAX_WORKERS", "4")) # Thread pool for concurrent processing102OCR_TIMEOUT_SECONDS = int(os.getenv("OCR_TIMEOUT_SECONDS", "120")) # 2 minute timeout for OCR103 104# Persistent session store (SQLite)105PERSIST_SESSIONS = str(os.getenv("PERSIST_SESSIONS", "0")).lower() in ("1", "true", "yes", "y")106SESSIONS_DB_PATH = os.getenv("SESSIONS_DB_PATH", "sessions.db")107SESSIONS_TTL_SECONDS = int(os.getenv("SESSIONS_TTL_SECONDS", "600"))108# Auto-cancel if all clients disconnect for duration (seconds). 0 disables it.109CANCEL_AFTER_DISCONNECT_SECONDS = int(os.getenv("CANCEL_AFTER_DISCONNECT_SECONDS", "3600"))110 111# Auto compression settings112ENABLE_AUTO_COMPRESSION = str(os.getenv("ENABLE_AUTO_COMPRESSION", "1")).lower() in ("1", "true", "yes", "y")113CONTEXT_MAX_TOKENS_AUTO = int(os.getenv("CONTEXT_MAX_TOKENS_AUTO", "0")) # 0 -> infer from model/tokenizer114CONTEXT_SAFETY_MARGIN = int(os.getenv("CONTEXT_SAFETY_MARGIN", "256"))115COMPRESSION_STRATEGY = os.getenv("COMPRESSION_STRATEGY", "truncate") # truncate | summarize (future)116 117# Eager model loading (download/check at startup before serving traffic)118EAGER_LOAD_MODEL = str(os.getenv("EAGER_LOAD_MODEL", "1")).lower() in ("1", "true", "yes", "y")119 120# Global thread pool executor for concurrent processing121executor = ThreadPoolExecutor(max_workers=MAX_WORKERS, thread_name_prefix="inference")122 123# Global OCR engine (Disabled)124# _ocr_engine = None125 126# def get_ocr_engine():127# global _ocr_engine128# if _ocr_engine is None and RapidOCR is not None:129# try:130# _ocr_engine = RapidOCR()131# print("[OCR] RapidOCR engine initialized")132# except Exception as e:133# print(f"[OCR] Failed to initialize RapidOCR: {e}")134# _ocr_engine = None135# return _ocr_engine136def get_ocr_engine():137 """Placeholder function - OCR disabled"""138 return None139 140def _log(msg: str):141 # Consistent, flush-immediate startup logs142 print(f"[startup] {msg}", flush=True)143 144def prefetch_model_assets(repo_id: str, token: Optional[str]) -> Optional[str]:145 """146 Reproducible prefetch driven by huggingface-cli:147 - Downloads the ENTIRE repo using CLI (visible progress bar).148 - Returns the local directory path where the repo is mirrored.149 - If CLI is unavailable, falls back to verbose API prefetch.150 """151 try:152 # Enable accelerated transfer only if hf_transfer is installed; otherwise disable to avoid runtime errors on Spaces153 try:154 import importlib.util as _imputil155 if _imputil.find_spec("hf_transfer") is not None:156 os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")157 else:158 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"159 except Exception:160 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"161 # XET acceleration if available; harmless if missing162 os.environ.setdefault("HF_HUB_ENABLE_XET", "1")163 164 cache_dir = os.getenv("HF_HOME") or os.getenv("TRANSFORMERS_CACHE") or ""165 if cache_dir:166 os.makedirs(cache_dir, exist_ok=True)167 168 # Resolve huggingface-cli path (Windows-friendly) - try hf first, fallback to huggingface-cli169 cli_path = shutil.which("hf")170 if not cli_path:171 cli_path = shutil.which("huggingface-cli")172 if not cli_path:173 candidates = []174 appdata = os.getenv("APPDATA")175 if appdata:176 candidates.append(os.path.join(appdata, "Python", "Python312", "Scripts", "hf.exe"))177 candidates.append(os.path.join(appdata, "Python", "Python312", "Scripts", "huggingface-cli.exe"))178 candidates.append(os.path.join(os.path.dirname(sys.executable), "Scripts", "hf.exe"))179 candidates.append(os.path.join(os.path.dirname(sys.executable), "Scripts", "huggingface-cli.exe"))180 cli_path = next((p for p in candidates if os.path.exists(p)), None)181 182 # Preferred: one-shot CLI download for the whole repo (shows live progress)183 if cli_path:184 local_root = os.path.join(cache_dir if cache_dir else ".", repo_id.replace("/", "_"))185 os.makedirs(local_root, exist_ok=True)186 _log(f"Using hf download to download entire repo -> '{local_root}'")187 cmd = [188 cli_path,189 "download",190 "--repo-type",191 "model",192 "--local-dir",193 local_root,194 repo_id,195 ]196 if token:197 cmd += ["--token", token]198 # Inherit stdio; users will see a proper progress bar199 subprocess.run(cmd, check=False)200 # Verify we have the essential files201 if os.path.exists(os.path.join(local_root, "config.json")) or os.path.exists(os.path.join(local_root, "model.safetensors")):202 _log("CLI prefetch completed")203 return local_root204 else:205 _log("CLI prefetch finished but essential files not found; will fallback to API mirroring")206 207 # Fallback: verbose API-driven prefetch with per-file logging208 _log(f"Prefetching (API) repo={repo_id} to cache='{cache_dir}'")209 try:210 files = list_repo_files(repo_id, repo_type="model", token=token)211 except Exception as e:212 _log(f"list_repo_files failed ({type(e).__name__}: {e}); falling back to snapshot_download")213 snapshot_download(repo_id, token=token, local_files_only=False)214 _log("Prefetch completed (snapshot)")215 return None216 217 total = len(files)218 _log(f"Found {total} files to ensure cached (API)")219 for i, fn in enumerate(files, start=1):220 try:221 meta = get_hf_file_metadata(repo_id, fn, repo_type="model", token=token)222 size_bytes = meta.size or 0223 except Exception:224 size_bytes = 0225 size_mb = size_bytes / (1024 * 1024) if size_bytes else 0.0226 _log(f"[{i}/{total}] fetching '{fn}' (~{size_mb:.2f} MB)")227 _ = hf_hub_download(228 repo_id=repo_id,229 filename=fn,230 repo_type="model",231 token=token,232 local_files_only=False,233 resume_download=True,234 )235 _log(f"[{i}/{total}] done '{fn}'")236 _log("Prefetch completed (API)")237 return None238 except Exception as e:239 _log(f"Prefetch skipped: {type(e).__name__}: {e}")240 return None241 242def is_data_url(url: str) -> bool:243 return url.startswith("data:") and ";base64," in url244 245 246def is_http_url(url: str) -> bool:247 return url.startswith("http://") or url.startswith("https://")248 249 250def decode_base64_to_bytes(b64: str) -> bytes:251 # strip possible "data:*;base64," prefix252 if "base64," in b64:253 b64 = b64.split("base64,", 1)[1]254 return base64.b64decode(b64, validate=False)255 256 257def fetch_bytes(url: str, headers: Optional[Dict[str, str]] = None, timeout: int = 60) -> bytes:258 if not is_http_url(url):259 raise ValueError(f"Only http(s) URLs supported for fetch, got: {url}")260 resp = requests.get(url, headers=headers or {}, timeout=timeout, stream=True)261 resp.raise_for_status()262 return resp.content263 264 265def load_image_from_any(src: Dict[str, Any]) -> Image.Image:266 """267 src can be:268 - { "url": "http(s)://..." } (also supports data URL)269 - { "b64_json": "<base64>" }270 - { "path": "local_path" } (optional)271 """272 if "b64_json" in src and src["b64_json"]:273 data = decode_base64_to_bytes(str(src["b64_json"]))274 return Image.open(io.BytesIO(data)).convert("RGB")275 276 if "url" in src and src["url"]:277 url = str(src["url"])278 if is_data_url(url):279 data = decode_base64_to_bytes(url)280 return Image.open(io.BytesIO(data)).convert("RGB")281 if is_http_url(url):282 data = fetch_bytes(url)283 return Image.open(io.BytesIO(data)).convert("RGB")284 # treat as local path285 if os.path.exists(url):286 with open(url, "rb") as f:287 return Image.open(io.BytesIO(f.read())).convert("RGB")288 raise ValueError(f"Invalid image url/path: {url}")289 290 if "path" in src and src["path"]:291 p = str(src["path"])292 if os.path.exists(p):293 with open(p, "rb") as f:294 return Image.open(io.BytesIO(f.read())).convert("RGB")295 raise ValueError(f"Image path not found: {p}")296 297 raise ValueError("Unsupported image source payload")298 299 300def write_bytes_tempfile(data: bytes, suffix: str) -> str:301 tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)302 with tmp as f:303 f.write(data)304 return tmp.name305 306 307def load_video_frames_from_any(src: Dict[str, Any], max_frames: int = MAX_VIDEO_FRAMES) -> List[Image.Image]:308 """309 Returns a list of PIL.Image frames (RGB) sampled up to max_frames.310 src can be:311 - { "url": "http(s)://..." } (mp4/mov/webm/etc.)312 - { "b64_json": "<base64 of a video file>" }313 - { "path": "local_path" }314 """315 # Prefer imageio.v3 if present, fallback to OpenCV316 # We load all frames then uniform sample if too many.317 def _load_all_frames(path: str) -> List[Image.Image]:318 frames: List[Image.Image] = []319 with contextlib.suppress(ImportError):320 import imageio.v3 as iio321 arr_iter = iio.imiter(path) # yields numpy arrays HxWxC322 for arr in arr_iter:323 if arr is None:324 continue325 if arr.ndim == 2:326 arr = np.stack([arr, arr, arr], axis=-1)327 if arr.shape[-1] == 4:328 arr = arr[..., :3]329 frames.append(Image.fromarray(arr).convert("RGB"))330 return frames331 332 # Fallback to OpenCV333 import cv2 # type: ignore334 cap = cv2.VideoCapture(path)335 ok, frame = cap.read()336 while ok:337 frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)338 frames.append(Image.fromarray(frame))339 ok, frame = cap.read()340 cap.release()341 return frames342 343 # Resolve to a local path344 local_path = None345 if "b64_json" in src and src["b64_json"]:346 data = decode_base64_to_bytes(str(src["b64_json"]))347 local_path = write_bytes_tempfile(data, suffix=".mp4")348 elif "url" in src and src["url"]:349 url = str(src["url"])350 if is_data_url(url):351 data = decode_base64_to_bytes(url)352 local_path = write_bytes_tempfile(data, suffix=".mp4")353 elif is_http_url(url):354 data = fetch_bytes(url)355 local_path = write_bytes_tempfile(data, suffix=".mp4")356 elif os.path.exists(url):357 local_path = url358 else:359 raise ValueError(f"Invalid video url/path: {url}")360 elif "path" in src and src["path"]:361 p = str(src["path"])362 if os.path.exists(p):363 local_path = p364 else:365 raise ValueError(f"Video path not found: {p}")366 else:367 raise ValueError("Unsupported video source payload")368 369 frames = _load_all_frames(local_path)370 # Uniform sample if too many frames371 if len(frames) > max_frames and max_frames > 0:372 idxs = np.linspace(0, len(frames) - 1, max_frames).astype(int).tolist()373 frames = [frames[i] for i in idxs]374 return frames375 376 377class ChatRequest(BaseModel):378 """OpenAI-compatible Chat Completions request body."""379 model: Optional[str] = Field(default=None, description="Model id (defaults to env MODEL_REPO_ID).")380 messages: List[Dict[str, Any]] = Field(description="OpenAI-style messages array. Supports text, image_url/input_image, video_url/input_video parts.")381 max_tokens: Optional[int] = Field(default=None, description="Max new tokens to generate.")382 temperature: Optional[float] = Field(default=None, description="Sampling temperature.")383 stream: Optional[bool] = Field(default=None, description="When true, returns Server-Sent Events stream.")384 session_id: Optional[str] = Field(default=None, description="Optional session id for resumable SSE.")385 # Pydantic v2 schema extras with rich examples386 model_config = ConfigDict(387 json_schema_extra={388 "examples": [389 {390 "summary": "Text-only",391 "value": {392 "messages": [393 {"role": "user", "content": "Hello, summarize the benefits of multimodal LLMs."}394 ],395 "max_tokens": 128396 }397 },398 {399 "summary": "Image by URL",400 "value": {401 "messages": [402 {403 "role": "user",404 "content": [405 {"type": "text", "text": "What is in this image?"},406 {"type": "image_url", "image_url": {"url": "https://example.com/cat.jpg"}}407 ]408 }409 ],410 "max_tokens": 128411 }412 },413 {414 "summary": "Video by URL (streaming SSE)",415 "value": {416 "messages": [417 {418 "role": "user",419 "content": [420 {"type": "text", "text": "Describe this clip briefly."},421 {"type": "video_url", "video_url": {"url": "https://example.com/clip.mp4"}}422 ]423 }424 ],425 "stream": True,426 "max_tokens": 128427 }428 }429 ]430 }431 )432 433class MessageModel(BaseModel):434 role: Literal["system", "user", "assistant"]435 content: str436 437class ChoiceModel(BaseModel):438 index: int439 message: MessageModel440 finish_reason: Optional[str] = None441 442class UsageModel(BaseModel):443 prompt_tokens: int444 completion_tokens: int445 total_tokens: int446 447class ChatCompletionResponse(BaseModel):448 """Non-streaming Chat Completions response (when stream=false)."""449 id: str450 object: str451 created: int452 model: str453 choices: List[ChoiceModel]454 usage: UsageModel455 context: Dict[str, Any] = {}456 457class HealthResponse(BaseModel):458 ok: bool459 modelReady: bool460 modelId: str461 error: Optional[str] = None462 context: Optional[Dict[str, Any]] = None463 464class CancelResponse(BaseModel):465 ok: bool466 session_id: str467 468 469class Engine:470 def __init__(self, model_id: str, hf_token: Optional[str] = None):471 # Lazy import heavy deps472 from transformers import AutoProcessor, AutoModelForCausalLM, AutoModelForVision2Seq, AutoModel, BitsAndBytesConfig473 # AutoModelForImageTextToText is the v5+ replacement for Vision2Seq in Transformers474 try:475 from transformers import AutoModelForImageTextToText # type: ignore476 except Exception:477 AutoModelForImageTextToText = None # type: ignore478 479 # Resolve device map to avoid 'meta' device on CPU Spaces480 # If DEVICE_MAP is "auto" but no CUDA is available, force "cpu" and disable low_cpu_mem_usage481 model_kwargs: Dict[str, Any] = {482 "trust_remote_code": True,483 }484 if hf_token:485 # Only pass 'token' (use_auth_token is deprecated and causes conflicts)486 model_kwargs["token"] = hf_token487 488 # Add quantization config if enabled489 if LOAD_IN_4BIT:490 try:491 import torch492 compute_dtype = getattr(torch, BNB_4BIT_COMPUTE_DTYPE, torch.float16)493 quant_config = BitsAndBytesConfig(494 load_in_4bit=True,495 bnb_4bit_compute_dtype=compute_dtype,496 bnb_4bit_use_double_quant=BNB_4BIT_USE_DOUBLE_QUANT,497 bnb_4bit_quant_type=BNB_4BIT_QUANT_TYPE,498 )499 model_kwargs["quantization_config"] = quant_config500 _log(f"Using 4-bit quantization: {BNB_4BIT_QUANT_TYPE}, compute_dtype={BNB_4BIT_COMPUTE_DTYPE}, double_quant={BNB_4BIT_USE_DOUBLE_QUANT}")501 except Exception as e:502 _log(f"BitsAndBytes quantization failed: {e}; falling back to full precision")503 504 # Device and dtype resolution505 try:506 import torch # local import to avoid heavy import at module load507 has_cuda = bool(getattr(torch, "cuda", None) and torch.cuda.is_available())508 except Exception:509 has_cuda = False510 511 resolved_device_map = DEVICE_MAP512 if str(DEVICE_MAP).lower() == "auto" and not has_cuda:513 resolved_device_map = "cpu"514 515 model_kwargs["device_map"] = resolved_device_map516 # Explicitly disable low_cpu_mem_usage on pure CPU to fully materialize weights (avoids meta tensors)517 if resolved_device_map == "cpu":518 model_kwargs["low_cpu_mem_usage"] = False519 # dtype - use 'dtype' instead of deprecated 'torch_dtype'520 if TORCH_DTYPE != "auto":521 try:522 import torch523 model_kwargs["dtype"] = getattr(torch, TORCH_DTYPE, TORCH_DTYPE)524 except Exception:525 model_kwargs["dtype"] = TORCH_DTYPE526 else:527 model_kwargs["dtype"] = "auto"528 # store for later529 self._resolved_device_map = resolved_device_map530 531 # Processor (handles text + images/videos)532 proc_kwargs: Dict[str, Any] = {"trust_remote_code": True}533 if hf_token:534 proc_kwargs["token"] = hf_token535 self.processor = AutoProcessor.from_pretrained(536 model_id,537 **proc_kwargs,538 ) # pragma: no cover539 540 # Prefer ImageTextToText (Transformers v5 path), then Vision2Seq, then CausalLM as a last resort541 model = None542 if 'AutoModelForImageTextToText' in globals() and AutoModelForImageTextToText is not None:543 try:544 model = AutoModelForImageTextToText.from_pretrained(model_id, **model_kwargs) # pragma: no cover545 except Exception:546 model = None547 if model is None:548 try:549 # AutoModelForVision2Seq is deprecated, but try it for compatibility550 model = AutoModelForVision2Seq.from_pretrained(model_id, **model_kwargs) # pragma: no cover551 except Exception:552 model = None553 if model is None:554 try:555 model = AutoModelForCausalLM.from_pretrained(model_id, **model_kwargs) # pragma: no cover556 except Exception:557 model = None558 if model is None:559 # Generic AutoModel as last-resort with trust_remote_code to load custom architectures560 model = AutoModel.from_pretrained(model_id, **model_kwargs) # pragma: no cover561 self.model = model.eval() # pragma: no cover562 # Ensure model is fully on CPU when resolved device_map is cpu (prevents meta device mix during inference)563 try:564 if str(getattr(self, "_resolved_device_map", "")).lower() == "cpu":565 _ = self.model.to("cpu")566 except Exception:567 pass568 # Ensure model is on CPU when resolved device_map is cpu (prevents meta device mix during inference)569 try:570 if getattr(self, "_resolved_device_map", None) == "cpu":571 _ = self.model.to("cpu")572 except Exception:573 pass574 575 self.model_id = model_id576 self.tokenizer = getattr(self.processor, "tokenizer", None)577 self.last_context_info: Dict[str, Any] = {}578 579 def _model_max_context(self) -> int:580 try:581 cfg = getattr(self.model, "config", None)582 if cfg is not None:583 v = getattr(cfg, "max_position_embeddings", None)584 if isinstance(v, int) and v > 0 and v < 10_000_000:585 return v586 except Exception:587 pass588 try:589 mx = int(getattr(self.tokenizer, "model_max_length", 0) or 0)590 if mx > 0 and mx < 10_000_000_000:591 return mx592 except Exception:593 pass594 return 32768595 596 def _count_prompt_tokens(self, text: str) -> int:597 try:598 if self.tokenizer is not None:599 enc = self.tokenizer([text], add_special_tokens=False, return_attention_mask=False)600 ids = enc["input_ids"][0]601 return len(ids)602 except Exception:603 pass604 return max(1, int(len(text.split()) * 1.3))605 606 def _auto_compress_if_needed(607 self, mm_messages: List[Dict[str, Any]], max_new_tokens: int608 ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]:609 info: Dict[str, Any] = {}610 # Build once to measure611 text0 = self.processor.apply_chat_template(mm_messages, tokenize=False, add_generation_prompt=True)612 prompt_tokens = self._count_prompt_tokens(text0)613 max_ctx = CONTEXT_MAX_TOKENS_AUTO if CONTEXT_MAX_TOKENS_AUTO > 0 else self._model_max_context()614 budget = max(1024, max_ctx - CONTEXT_SAFETY_MARGIN - int(max_new_tokens))615 if not ENABLE_AUTO_COMPRESSION or prompt_tokens <= budget:616 info = {617 "compressed": False,618 "prompt_tokens": int(prompt_tokens),619 "max_context": int(max_ctx),620 "budget": int(budget),621 "strategy": COMPRESSION_STRATEGY,622 "dropped_messages": 0,623 }624 return mm_messages, info625 626 # Truncate earliest non-system messages until within budget627 msgs = list(mm_messages)628 dropped = 0629 guard = 0630 while True:631 text = self.processor.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)632 prompt_tokens = self._count_prompt_tokens(text)633 if prompt_tokens <= budget or len(msgs) <= 1:634 break635 # drop earliest non-system636 drop_idx = None637 for j, m in enumerate(msgs):638 if (m.get("role") or "user") != "system":639 drop_idx = j640 break641 if drop_idx is None:642 break643 msgs.pop(drop_idx)644 dropped += 1645 guard += 1646 if guard > 10000:647 break648 649 info = {650 "compressed": True,651 "prompt_tokens": int(prompt_tokens),652 "max_context": int(max_ctx),653 "budget": int(budget),654 "strategy": "truncate",655 "dropped_messages": int(dropped),656 }657 return msgs, info658 659 def get_context_report(self) -> Dict[str, Any]:660 try:661 tk_max = int(getattr(self.tokenizer, "model_max_length", 0) or 0)662 except Exception:663 tk_max = 0664 return {665 "compressionEnabled": ENABLE_AUTO_COMPRESSION,666 "strategy": COMPRESSION_STRATEGY,667 "safetyMargin": CONTEXT_SAFETY_MARGIN,668 "modelMaxContext": self._model_max_context(),669 "tokenizerModelMaxLength": tk_max,670 "last": self.last_context_info or {},671 }672 673 def build_mm_messages(674 self, openai_messages: List[Dict[str, Any]]675 ) -> Tuple[List[Dict[str, Any]], List[Image.Image], List[List[Image.Image]]]:676 """677 Convert OpenAI-style messages to Qwen multimodal messages.678 Returns:679 - messages for apply_chat_template680 - flat list of images in encounter order681 - list of videos (each is list of PIL frames)682 """683 mm_msgs: List[Dict[str, Any]] = []684 images: List[Image.Image] = []685 videos: List[List[Image.Image]] = []686 687 for msg in openai_messages:688 role = msg.get("role", "user")689 content = msg.get("content", "")690 691 parts: List[Dict[str, Any]] = []692 693 if isinstance(content, str):694 if content:695 parts.append({"type": "text", "text": content})696 elif isinstance(content, list):697 for p in content:698 ptype = p.get("type")699 if ptype == "text":700 txt = p.get("text", "")701 if txt:702 parts.append({"type": "text", "text": txt})703 elif ptype in ("image_url", "input_image"):704 src: Dict[str, Any] = {}705 if ptype == "image_url":706 u = (p.get("image_url") or {}).get("url") if isinstance(p.get("image_url"), dict) else p.get("image_url")707 src["url"] = u708 else:709 b64 = p.get("image") or p.get("b64_json") or p.get("data") or (p.get("image_url") or {}).get("url")710 if b64:711 src["b64_json"] = b64712 try:713 img = load_image_from_any(src)714 images.append(img)715 parts.append({"type": "image", "image": img})716 except Exception as e:717 raise ValueError(f"Failed to parse image part: {e}") from e718 elif ptype in ("video_url", "input_video"):719 src = {}720 if ptype == "video_url":721 u = (p.get("video_url") or {}).get("url") if isinstance(p.get("video_url"), dict) else p.get("video_url")722 src["url"] = u723 else:724 b64 = p.get("video") or p.get("b64_json") or p.get("data")725 if b64:726 src["b64_json"] = b64727 try:728 frames = load_video_frames_from_any(src, max_frames=MAX_VIDEO_FRAMES)729 videos.append(frames)730 parts.append({"type": "video", "video": frames})731 except Exception as e:732 raise ValueError(f"Failed to parse video part: {e}") from e733 else:734 if isinstance(p, dict):735 txt = p.get("text")736 if isinstance(txt, str) and txt:737 parts.append({"type": "text", "text": txt})738 else:739 if content:740 parts.append({"type": "text", "text": str(content)})741 742 mm_msgs.append({"role": role, "content": parts})743 744 return mm_msgs, images, videos745 746 def infer(self, messages: List[Dict[str, Any]], max_tokens: int, temperature: float) -> str:747 mm_messages, images, videos = self.build_mm_messages(messages)748 # Auto-compress if needed based on context budget749 mm_messages, ctx_info = self._auto_compress_if_needed(mm_messages, max_tokens)750 self.last_context_info = ctx_info751 752 # Build chat template753 text = self.processor.apply_chat_template(754 mm_messages,755 tokenize=False,756 add_generation_prompt=True,757 )758 759 proc_kwargs: Dict[str, Any] = {"text": [text], "return_tensors": "pt"}760 if images:761 proc_kwargs["images"] = images762 if videos:763 proc_kwargs["videos"] = videos764 765 inputs = self.processor(**proc_kwargs)766 # Move tensors to the correct device767 try:768 if str(getattr(self, "_resolved_device_map", "")).lower() == "cpu":769 # Explicit CPU placement avoids 'meta' device errors on Spaces770 inputs = {k: (v.to("cpu") if hasattr(v, "to") else v) for k, v in inputs.items()}771 else:772 device = getattr(self.model, "device", None) or next(self.model.parameters()).device773 inputs = {k: (v.to(device) if hasattr(v, "to") else v) for k, v in inputs.items()}774 except Exception:775 pass776 777 do_sample = temperature is not None and float(temperature) > 0.0778 779 # Safer on CPU: run without gradients to reduce memory pressure and avoid autograd hooks780 try:781 import torch782 with torch.no_grad():783 gen_ids = self.model.generate(784 **inputs,785 max_new_tokens=int(max_tokens),786 temperature=float(temperature),787 do_sample=do_sample,788 use_cache=True,789 )790 except Exception:791 # Fallback without no_grad if torch import fails (very unlikely)792 gen_ids = self.model.generate(793 **inputs,794 max_new_tokens=int(max_tokens),795 temperature=float(temperature),796 do_sample=do_sample,797 use_cache=True,798 )799 800 # Decode801 output = self.processor.batch_decode(802 gen_ids,803 skip_special_tokens=True,804 clean_up_tokenization_spaces=False,805 )[0]806 807 # Best-effort: return only the assistant reply after the last template marker if present808 parts = re.split(r"\n?assistant:\s*", output, flags=re.IGNORECASE)809 if len(parts) >= 2:810 return parts[-1].strip()811 return output.strip()812 813 def infer_stream(814 self,815 messages: List[Dict[str, Any]],816 max_tokens: int,817 temperature: float,818 cancel_event: Optional[threading.Event] = None,819 ):820 """Streaming implementation that simulates streaming using the working infer() method"""821 print(f"[STREAM] Starting simulated streaming with max_tokens={max_tokens}, temperature={temperature}")822 print(f"[STREAM] Input messages: {messages}")823 824 # For Qwen3, use the working infer() method and simulate streaming825 # by yielding chunks of the generated text826 try:827 # Generate the complete response using the working infer() method828 print(f"[STREAM] Generating complete response using infer() method...")829 full_response = self.infer(messages, max_tokens, temperature)830 print(f"[STREAM] Generated full response: '{full_response}'")831 print(f"[STREAM] Response length: {len(full_response)} characters")832 833 if not full_response.strip():834 print("[STREAM] No response generated - returning empty")835 return836 837 # Simulate streaming by yielding chunks of the response838 import time839 import re840 841 # Split response into natural chunks (words, phrases, sentences)842 # This makes the streaming feel more natural than character-by-character843 words = re.findall(r'\S+|\s+', full_response) # Keep spaces as separate tokens844 845 chunk_size = 2 # Yield 2-3 words at a time for natural flow846 current_text = ""847 piece_count = 0848 849 print(f"[STREAM] Starting simulated streaming with {len(words)} word tokens...")850 start_time = time.time()851 timeout_seconds = 30 # 30 second timeout for streaming852 853 for i in range(0, len(words), chunk_size):854 # Check timeout855 if time.time() - start_time > timeout_seconds:856 print(f"[STREAM] Streaming timeout after {piece_count} pieces")857 break858 859 # Check cancellation860 if cancel_event and cancel_event.is_set():861 print(f"[STREAM] Streaming cancelled after {piece_count} pieces")862 break863 864 # Get next chunk865 chunk_words = words[i:i + chunk_size]866 chunk = ''.join(chunk_words)867 868 if not chunk.strip(): # Skip empty whitespace chunks869 continue870 871 current_text += chunk872 piece_count += 1873 874 print(f"[STREAM] Yielding chunk #{piece_count}: '{chunk}'")875 print(f"[STREAM] Total so far: '{current_text[-100:]}'")876 877 yield chunk878 879 # Small delay to simulate real-time generation880 time.sleep(0.1)881 882 print(f"[STREAM] Simulated streaming completed")883 print(f"[STREAM] Total chunks yielded: {piece_count}")884 print(f"[STREAM] Final text: '{current_text}'")885 886 except Exception as e:887 print(f"[STREAM] Simulated streaming error: {e}")888 import traceback889 traceback.print_exc()890 891 892# Simple in-memory resumable SSE session store + optional SQLite persistence893class _SSESession:894 def __init__(self, maxlen: int = 2048, ttl_seconds: int = 600):895 self.buffer: Deque[Tuple[int, str]] = deque(maxlen=maxlen) # (idx, sse_line_block)896 self.last_idx: int = -1897 self.created: float = time.time()898 self.finished: bool = False899 self.cond = threading.Condition()900 self.thread: Optional[threading.Thread] = None901 self.ttl_seconds = ttl_seconds902 # Cancellation + client tracking903 self.cancel_event = threading.Event()904 self.listeners: int = 0905 self.cancel_timer = None # type: ignore906 907 908class _SessionStore:909 def __init__(self, ttl_seconds: int = 600, max_sessions: int = 256):910 self._sessions: Dict[str, _SSESession] = {}911 self._lock = threading.Lock()912 self._ttl = ttl_seconds913 self._max_sessions = max_sessions914 915 def get_or_create(self, sid: str) -> _SSESession:916 with self._lock:917 sess = self._sessions.get(sid)918 if sess is None:919 sess = _SSESession(ttl_seconds=self._ttl)920 self._sessions[sid] = sess921 return sess922 923 def get(self, sid: str) -> Optional[_SSESession]:924 with self._lock:925 return self._sessions.get(sid)926 927 def gc(self):928 now = time.time()929 with self._lock:930 # remove expired931 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)]932 for k in expired:933 self._sessions.pop(k, None)934 # bound session count935 if len(self._sessions) > self._max_sessions:936 for k, _ in sorted(self._sessions.items(), key=lambda kv: kv[1].created)[: max(0, len(self._sessions) - self._max_sessions)]:937 self._sessions.pop(k, None)938 939 940class _SQLiteStore:941 def __init__(self, db_path: str):942 self.db_path = db_path943 self._lock = threading.Lock()944 self._conn = sqlite3.connect(self.db_path, check_same_thread=False)945 self._conn.execute("PRAGMA journal_mode=WAL;")946 self._conn.execute("PRAGMA synchronous=NORMAL;")947 self._ensure_schema()948 949 def _ensure_schema(self):950 cur = self._conn.cursor()951 cur.execute(952 "CREATE TABLE IF NOT EXISTS sessions (session_id TEXT PRIMARY KEY, created REAL, finished INTEGER DEFAULT 0)"953 )954 cur.execute(955 "CREATE TABLE IF NOT EXISTS events (session_id TEXT, idx INTEGER, data TEXT, created REAL, PRIMARY KEY(session_id, idx))"956 )957 cur.execute("CREATE INDEX IF NOT EXISTS idx_events_session ON events(session_id, idx)")958 self._conn.commit()959 960 def ensure_session(self, session_id: str, created: int):961 with self._lock:962 self._conn.execute(963 "INSERT OR IGNORE INTO sessions(session_id, created, finished) VALUES (?, ?, 0)",964 (session_id, float(created)),965 )966 self._conn.commit()967 968 def append_event(self, session_id: str, idx: int, payload: Dict[str, Any]):969 data = json.dumps(payload, ensure_ascii=False)970 with self._lock:971 self._conn.execute(972 "INSERT OR REPLACE INTO events(session_id, idx, data, created) VALUES (?, ?, ?, ?)",973 (session_id, idx, data, time.time()),974 )975 self._conn.commit()976 977 def get_events_after(self, session_id: str, last_idx: int) -> List[Tuple[int, str]]:978 with self._lock:979 cur = self._conn.execute(980 "SELECT idx, data FROM events WHERE session_id=? AND idx>? ORDER BY idx ASC", (session_id, last_idx)981 )982 return [(int(r[0]), str(r[1])) for r in cur.fetchall()]983 984 def mark_finished(self, session_id: str):985 with self._lock:986 self._conn.execute("UPDATE sessions SET finished=1 WHERE session_id=?", (session_id,))987 self._conn.commit()988 989 def session_meta(self, session_id: str) -> Tuple[bool, int]:990 with self._lock:991 row = self._conn.execute("SELECT finished FROM sessions WHERE session_id=?", (session_id,)).fetchone()992 finished = bool(row[0]) if row else False993 row2 = self._conn.execute("SELECT MAX(idx) FROM events WHERE session_id=?", (session_id,)).fetchone()994 last_idx = int(row2[0]) if row2 and row2[0] is not None else -1995 return finished, last_idx996 997 def gc(self, ttl_seconds: int):998 cutoff = time.time() - float(ttl_seconds)999 with self._lock:1000 cur = self._conn.execute("SELECT session_id FROM sessions WHERE finished=1 AND created<?", (cutoff,))1001 ids = [r[0] for r in cur.fetchall()]1002 for sid in ids:1003 self._conn.execute("DELETE FROM events WHERE session_id=?", (sid,))1004 self._conn.execute("DELETE FROM sessions WHERE session_id=?", (sid,))1005 self._conn.commit()1006 1007 1008def _sse_event(session_id: str, idx: int, payload: Dict[str, Any]) -> str:1009 # Include SSE id line so clients can send Last-Event-ID to resume.1010 return f"id: {session_id}:{idx}\n" + f"data: {json.dumps(payload, ensure_ascii=False)}\n\n"1011 1012 1013_STORE = _SessionStore()1014_DB_STORE = _SQLiteStore(SESSIONS_DB_PATH) if PERSIST_SESSIONS else None1015 1016# FastAPI app and OpenAPI tags1017tags_metadata = [1018 {"name": "meta", "description": "Service metadata and OpenAPI schema"},1019 {"name": "health", "description": "Readiness and runtime info including context window report"},1020 {"name": "chat", "description": "OpenAI-compatible chat completions (non-stream and streaming SSE)"},1021 {"name": "ocr", "description": "Optical Character Recognition endpoints"},1022]1023 1024app = FastAPI(1025 title="Qwen3-VL Inference Server",1026 version="1.0.0",1027 description="OpenAI-compatible inference server for Qwen3-VL with multimodal support, streaming SSE with resume, context auto-compression, and optional SQLite persistence.",1028 openapi_tags=tags_metadata,1029)1030 1031# Initialize startup time for uptime tracking1032import time1033app.state.start_time = time.time()1034app.add_middleware(1035 CORSMiddleware,1036 allow_origins=["*"],1037 allow_credentials=True,1038 allow_methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"],1039 allow_headers=["*"],1040 expose_headers=["*"],1041)1042 1043# Startup hook is defined after get_engine() so globals are initialized first.1044# Serve static web UI if present1045_WEB_DIR = os.path.join(ROOT_DIR, "web")1046if os.path.isdir(_WEB_DIR):1047 try:1048 app.mount("/web", StaticFiles(directory=_WEB_DIR, html=True), name="web")1049 except Exception:1050 pass1051 1052# Engine singletons1053_engine: Optional[Engine] = None1054_engine_error: Optional[str] = None1055 1056 1057def get_engine() -> Engine:1058 global _engine, _engine_error1059 if _engine is not None:1060 return _engine1061 try:1062 model_id = DEFAULT_MODEL_ID1063 _log(f"Preparing model '{model_id}' (HF_HOME={os.getenv('HF_HOME')}, cache={os.getenv('TRANSFORMERS_CACHE')})")1064 local_repo_dir = prefetch_model_assets(model_id, HF_TOKEN)1065 load_id = local_repo_dir if (local_repo_dir and os.path.exists(os.path.join(local_repo_dir, 'config.json'))) else model_id1066 _log(f"Loading processor and model from: {load_id}")1067 _engine = Engine(model_id=load_id, hf_token=HF_TOKEN)1068 _engine_error = None1069 _log(f"Model ready: {_engine.model_id}")1070 return _engine1071 except Exception as e:1072 _engine_error = f"{type(e).__name__}: {e}"1073 _log(f"Engine init failed: {_engine_error}")1074 raise1075 1076# Eager-load model at startup after definitions so it downloads/checks before serving traffic.1077@app.on_event("startup")1078def _startup_load_model():1079 # Initialize marketplace database1080 try:1081 from database import init_db1082 init_db()1083 print("[startup] Marketplace database initialized")1084 except Exception as e:1085 print(f"[startup] Database initialization failed: {e}")1086 1087 if EAGER_LOAD_MODEL:1088 print("[startup] EAGER_LOAD_MODEL=1: initializing model...")1089 print("[startup] OCR engine disabled - using image captioning instead")1090 try:1091 # Then initialize the model1092 _ = get_engine()1093 print("[startup] Model loaded:", _engine.model_id if _engine else "unknown")1094 except Exception as e:1095 # Log error but don't fail - allow server to start without model1096 print("[startup] Initialization failed:", e)1097 print("[startup] Server will start without full initialization")1098 else:1099 print("[startup] EAGER_LOAD_MODEL=0: skipping initialization")1100 1101 1102@app.get("/", tags=["meta"], include_in_schema=False)1103def root():1104 """1105 Redirect to the simple chat interface to avoid CORS issues with external scripts.1106 The main interface loads external dependencies that get blocked on Hugging Face Spaces.1107 """1108 from fastapi.responses import RedirectResponse1109 return RedirectResponse(url="/simple", status_code=302)1110 1111 1112@app.get("/debug", tags=["meta"], include_in_schema=False)1113def debug_interface():1114 """1115 Serve the debug interface for comprehensive testing and monitoring1116 """1117 debug_path = os.path.join(ROOT_DIR, "web", "debug_interface.html")1118 if os.path.exists(debug_path):1119 return FileResponse(debug_path, media_type="text/html; charset=utf-8")1120 1121 # Fallback if file doesn't exist1122 html = """<!doctype html><html><head><meta charset='utf-8'><title>Debug Interface</title></head>1123 <body style="font-family:system-ui,Segoe UI,Roboto;padding:24px;background:#1a202c;color:#e2e8f0">1124 <h2>Debug Interface Not Found</h2>1125 <p>The debug interface file was not found. Please ensure debug_interface.html is in the web directory.</p>1126 <p><a href="/" style="color:#93c5fd">Return to Main Interface</a></p>1127 </body></html>"""1128 return Response(html, media_type="text/html; charset=utf-8")1129 1130 1131@app.get("/simple", tags=["meta"], include_in_schema=False)1132def simple_chat_interface():1133 """1134 Serve the simple chat interface that works reliably1135 """1136 simple_path = os.path.join(ROOT_DIR, "web", "simple_chat.html")1137 if os.path.exists(simple_path):1138 return FileResponse(simple_path, media_type="text/html; charset=utf-8")1139 1140 # Fallback if file doesn't exist1141 html = """<!doctype html><html><head><meta charset='utf-8'><title>Simple Chat</title></head>1142 <body style="font-family:system-ui,Segoe UI,Roboto;padding:24px;background:#1a202c;color:#e2e8f0">1143 <h2>Simple Chat Interface Not Found</h2>1144 <p>The simple chat interface file was not found. Please ensure simple_chat.html is in the web directory.</p>1145 <p><a href="/" style="color:#93c5fd">Return to Main Interface</a></p>1146 </body></html>"""1147 return Response(html, media_type="text/html; charset=utf-8")1148 1149 1150@app.get("/openapi.yaml", tags=["meta"])1151def openapi_yaml():1152 """Serve OpenAPI schema as YAML for tooling compatibility."""1153 schema = app.openapi()1154 yml = yaml.safe_dump(schema, sort_keys=False)1155 return Response(yml, media_type="application/yaml")1156 1157 1158@app.get("/health", tags=["health"], response_model=HealthResponse)1159def health():1160 ready = False1161 err = None1162 model_id = DEFAULT_MODEL_ID1163 global _engine, _engine_error1164 if _engine is not None:1165 ready = True1166 model_id = _engine.model_id1167 elif _engine_error:1168 err = _engine_error1169 ctx = None1170 try:1171 if _engine is not None:1172 ctx = _engine.get_context_report()1173 except Exception:1174 ctx = None1175 return JSONResponse({"ok": True, "modelReady": ready, "modelId": model_id, "error": err, "context": ctx})1176 1177 1178@app.post(1179 "/v1/chat/completions",1180 tags=["chat"],1181 response_model=ChatCompletionResponse,1182 responses={1183 200: {1184 "description": "When stream=true, the response is text/event-stream (SSE). When stream=false, JSON body matches ChatCompletionResponse.",1185 "content": {1186 "text/event-stream": {1187 "schema": {"type": "string"},1188 "examples": {1189 "sse": {1190 "summary": "SSE stream example",1191 "value": "id: sess-123:0\ndata: {\"id\":\"sess-123\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\"}}]}\n\n"1192 }1193 }1194 }1195 },1196 }1197 },1198)1199 1200@app.options("/v1/chat/completions", tags=["chat"])