import asyncio import base64 import io import json import logging import time import uuid import threading from contextlib import asynccontextmanager from typing import AsyncGenerator, Dict, List, Optional, Union from pathlib import Path import numpy as np import onnxruntime as ort from fastapi import FastAPI, HTTPException, Request, UploadFile, File from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse from huggingface_hub import hf_hub_download, list_repo_files from pydantic import BaseModel, Field from sse_starlette.sse import EventSourceResponse from transformers import AutoImageProcessor, PreTrainedTokenizerFast from PIL import Image import aiohttp from config import settings # Configure logging logging.basicConfig( level=getattr(logging, settings.log_level.upper()), format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" ) logger = logging.getLogger(__name__) # ============================================================================== # Pydantic Models for OpenAI-compatible API # ============================================================================== class ImageContent(BaseModel): type: str = "image" image_url: Optional[str] = None # data:image/jpeg;base64,... or URL class TextContent(BaseModel): type: str = "text" text: str class VisionMessage(BaseModel): role: str = Field(..., description="Role: 'system', 'user', or 'assistant'") content: Union[str, List[Union[ImageContent, TextContent, dict]]] = Field(..., description="Message content") class VisionCompletionRequest(BaseModel): model: str = Field(default="lfm-vision", description="Model identifier") messages: List[VisionMessage] = Field(..., description="Conversation messages") temperature: Optional[float] = Field(default=None, ge=0.0, le=2.0) top_p: Optional[float] = Field(default=None, ge=0.0, le=1.0) top_k: Optional[int] = Field(default=None, ge=0) max_tokens: Optional[int] = Field(default=None, ge=1) stream: bool = Field(default=False, description="Enable streaming response") stop: Optional[Union[str, List[str]]] = Field(default=None) class ChatMessage(BaseModel): role: str = Field(..., description="Role: 'system', 'user', or 'assistant'") content: str = Field(..., description="Message content") class ChatCompletionRequest(BaseModel): model: str = Field(default="lfm-vision", description="Model identifier") messages: List[ChatMessage] = Field(..., description="Conversation messages") temperature: Optional[float] = Field(default=None, ge=0.0, le=2.0) top_p: Optional[float] = Field(default=None, ge=0.0, le=1.0) top_k: Optional[int] = Field(default=None, ge=0) max_tokens: Optional[int] = Field(default=None, ge=1) stream: bool = Field(default=False, description="Enable streaming response") class ChatCompletionChoice(BaseModel): index: int message: ChatMessage finish_reason: Optional[str] = None class ChatCompletionResponse(BaseModel): id: str object: str = "chat.completion" created: int model: str choices: List[ChatCompletionChoice] usage: Dict[str, int] class ModelInfo(BaseModel): id: str object: str = "model" created: int owned_by: str = "liquid-ai" class ModelListResponse(BaseModel): object: str = "list" data: List[ModelInfo] # ============================================================================== # ONNX Vision Model Manager # ============================================================================== # ONNX dtype mapping ONNX_DTYPE = { "tensor(float)": np.float32, "tensor(float16)": np.float16, "tensor(int64)": np.int64 } class Lfm2VlProcessorWrapper: """ Custom processor wrapper that combines ImageProcessor + Tokenizer. This bypasses the AutoProcessor tokenizer auto-detection bug in LFM models. """ def __init__(self, image_processor, tokenizer): self.image_processor = image_processor self.tokenizer = tokenizer def apply_chat_template(self, messages, add_generation_prompt=True, tokenize=False, **kwargs): """ Apply chat template for vision-language model. Converts vision message format [{"type": "image"}, {"type": "text", "text": "..."}] to text with placeholders as expected by the tokenizer. """ # Transform vision messages to text format text_messages = [] for msg in messages: role = msg.get("role", "user") if isinstance(msg, dict) else getattr(msg, "role", "user") content = msg.get("content", "") if isinstance(msg, dict) else getattr(msg, "content", "") if isinstance(content, list): # Vision message format: [{"type": "image"}, {"type": "text", "text": "..."}] text_parts = [] for item in content: if isinstance(item, dict): item_type = item.get("type", "") if item_type == "image": text_parts.append("") elif item_type == "text": text_parts.append(item.get("text", "")) else: text_parts.append(str(item)) content = "".join(text_parts) text_messages.append({"role": role, "content": content}) return self.tokenizer.apply_chat_template( text_messages, add_generation_prompt=add_generation_prompt, tokenize=tokenize, **kwargs ) def __call__(self, images=None, text=None, **kwargs): """ Process images and text for the vision-language model. CRITICAL: The vision encoder produces N image embeddings (e.g., 256 for a 512x512 image). Each embedding needs its own token position in input_ids. This method: 1. Processes images FIRST to determine N (number of image tokens) 2. Expands single in text to N consecutive tokens 3. Tokenizes the expanded text Returns a dict with pixel_values, input_ids, attention_mask, etc. """ result = {} return_tensors = kwargs.pop('return_tensors', None) num_image_tokens = 0 # Step 1: Process images FIRST to get the number of image tokens if images is not None: image_outputs = self.image_processor(images=images, return_tensors=return_tensors) result.update(image_outputs) # Calculate number of image tokens from pixel_values shape # pixel_values shape: [batch, num_patches, hidden_dim] # The MLP projector in LFM2.5-VL reduces patches by factor of 4 # Reference: https://huggingface.co/LiquidAI/LFM2.5-VL-1.6B if 'pixel_values' in image_outputs: pv = image_outputs['pixel_values'] num_patches = pv.shape[1] if hasattr(pv, 'shape') else pv.size(1) # MLP projector reduces by factor of 4: 1024 patches → 256 tokens num_image_tokens = num_patches // 4 logger.debug(f"Image processing: {num_patches} patches → {num_image_tokens} image tokens") # Step 2: Expand placeholder(s) to match token count if text is not None: # Ensure text is a string if isinstance(text, list): text = text[0] if len(text) == 1 else " ".join(text) # Expand each placeholder to N tokens if num_image_tokens > 0 and "" in text: # Count existing placeholders image_count = text.count("") # Each placeholder represents one image, expand to num_image_tokens tokens_per_image = num_image_tokens // image_count if image_count > 0 else num_image_tokens expanded_image = "" * tokens_per_image text = text.replace("", expanded_image) logger.debug(f"Expanded {image_count} placeholder(s) to {tokens_per_image} tokens each") text_outputs = self.tokenizer( text, return_tensors=return_tensors, padding=kwargs.get('padding', False), truncation=kwargs.get('truncation', False), max_length=kwargs.get('max_length', None) ) result.update(text_outputs) return result class ONNXVisionModelManager: """Manages ONNX Vision-Language model with 3 sessions: embed_tokens, embed_images, decoder.""" def __init__(self): self._embed_tokens = None self._embed_images = None self._decoder = None self._processor = None self._cache_template = None self._lock = threading.Lock() @property def is_loaded(self) -> bool: return all([self._embed_tokens, self._embed_images, self._decoder]) def download_models(self) -> Dict[str, str]: """Download ONNX model files from HuggingFace.""" model_id = settings.model_id encoder_var = settings.encoder_variant decoder_var = settings.decoder_variant logger.info(f"Downloading model: {model_id}") logger.info(f" Encoder variant: {encoder_var}") logger.info(f" Decoder variant: {decoder_var}") paths = {} # Download embed_tokens (use same variant as encoder or fp16) embed_suffix = f"_fp16" if encoder_var in ["fp16", "q8", "q4"] else "" paths["embed_tokens"] = hf_hub_download(model_id, f"onnx/embed_tokens{embed_suffix}.onnx") # Download embed_images (vision encoder) img_suffix = f"_{encoder_var}" if encoder_var != "fp32" else "" paths["embed_images"] = hf_hub_download(model_id, f"onnx/embed_images{img_suffix}.onnx") # Download decoder dec_suffix = f"_{decoder_var}" if decoder_var != "fp32" else "" paths["decoder"] = hf_hub_download(model_id, f"onnx/decoder{dec_suffix}.onnx") # Download all data files - use exact prefix matching to avoid downloading wrong variants # Expected files for selected variants only (e.g., decoder_q8.onnx_data, not decoder.onnx_data) expected_prefixes = [ f"onnx/embed_tokens{embed_suffix}.onnx_data", f"onnx/embed_images{img_suffix}.onnx_data", f"onnx/decoder{dec_suffix}.onnx_data" ] for f in list_repo_files(model_id): if f.startswith("onnx/") and ".onnx_data" in f: # Check if this file STARTS WITH one of our expected prefixes # This handles split files like decoder_q8.onnx_data, decoder_q8.onnx_data_1, etc. if any(f.startswith(prefix) for prefix in expected_prefixes): logger.info(f"Downloading: {f}") hf_hub_download(model_id, f) return paths def load_model(self) -> None: """Load the ONNX models and processor.""" with self._lock: if self.is_loaded: return logger.info("=" * 60) logger.info("Loading LFM2.5-VL-1.6B Vision-Language ONNX model...") logger.info(f"Model: {settings.model_id}") logger.info(f"Encoder: {settings.encoder_variant} (Q8 = ~95% accuracy)") logger.info(f"Decoder: {settings.decoder_variant}") logger.info("=" * 60) start_time = time.time() # Download models paths = self.download_models() # Configure ONNX Runtime for CPU sess_options = ort.SessionOptions() sess_options.intra_op_num_threads = settings.num_threads sess_options.inter_op_num_threads = settings.num_threads sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL # Load ONNX sessions self._embed_tokens = ort.InferenceSession( paths["embed_tokens"], sess_options=sess_options, providers=['CPUExecutionProvider'] ) self._embed_images = ort.InferenceSession( paths["embed_images"], sess_options=sess_options, providers=['CPUExecutionProvider'] ) self._decoder = ort.InferenceSession( paths["decoder"], sess_options=sess_options, providers=['CPUExecutionProvider'] ) # Load processor components separately to bypass TokenizersBackend bug # LFM models incorrectly specify TokenizersBackend as tokenizer_class logger.info("Loading image processor...") image_processor = AutoImageProcessor.from_pretrained( settings.model_id, trust_remote_code=True ) logger.info("Loading tokenizer with PreTrainedTokenizerFast...") tokenizer = PreTrainedTokenizerFast.from_pretrained( settings.model_id, trust_remote_code=True ) # Create our custom processor wrapper self._processor = Lfm2VlProcessorWrapper( image_processor=image_processor, tokenizer=tokenizer ) logger.info(f"✓ Processor created: {type(self._processor).__name__}") # Initialize cache template for decoder self._init_cache_template() load_time = time.time() - start_time logger.info("=" * 60) logger.info(f"✓ Model loaded in {load_time:.2f}s") logger.info(f" Threads: {settings.num_threads}") logger.info(f" Provider: CPU") logger.info("=" * 60) def _init_cache_template(self) -> None: """Initialize KV cache template for decoder.""" self._cache_template = {} for inp in self._decoder.get_inputs(): if inp.name in {"inputs_embeds", "attention_mask", "position_ids"}: continue shape = [d if isinstance(d, int) else 1 for d in inp.shape] for i, d in enumerate(inp.shape): if isinstance(d, str) and "sequence" in d.lower(): shape[i] = 0 dtype = ONNX_DTYPE.get(inp.type, np.float32) self._cache_template[inp.name] = (shape, dtype) def _create_empty_cache(self) -> Dict[str, np.ndarray]: """Create a new empty KV cache.""" return { name: np.zeros(shape, dtype=dtype) for name, (shape, dtype) in self._cache_template.items() } @property def processor(self): if self._processor is None: raise RuntimeError("Processor not loaded") return self._processor def process_image(self, image: Image.Image) -> Dict[str, np.ndarray]: """Process image to embeddings.""" # Ensure RGB if image.mode != "RGB": image = image.convert("RGB") return image def generate( self, images: List[Image.Image], messages: List[dict], max_tokens: int = 512, temperature: float = 0.1, top_k: int = 50, top_p: float = 0.1, stop_tokens: Optional[List[int]] = None ) -> List[int]: """Generate tokens using ONNX Vision model.""" tokenizer = self._processor.tokenizer if stop_tokens is None: stop_tokens = [tokenizer.eos_token_id] # Process inputs through processor prompt = self._processor.apply_chat_template(messages, add_generation_prompt=True) inputs = self._processor( images=images if images else None, text=prompt, return_tensors="pt" ) # Convert to numpy with correct dtypes input_ids = inputs["input_ids"].numpy().astype(np.int64) # Get token embeddings token_outputs = self._embed_tokens.run(None, {"input_ids": input_ids}) token_embeds = token_outputs[0] # Process images if present if images and "pixel_values" in inputs: pixel_values = inputs["pixel_values"].numpy().astype(np.float32) pixel_attention_mask = inputs.get("pixel_attention_mask", None) spatial_shapes = inputs.get("spatial_shapes", None) image_feed = {"pixel_values": pixel_values} if pixel_attention_mask is not None: image_feed["pixel_attention_mask"] = pixel_attention_mask.numpy().astype(np.int64) if spatial_shapes is not None: image_feed["spatial_shapes"] = spatial_shapes.numpy().astype(np.int64) image_outputs = self._embed_images.run(None, image_feed) image_embeds = image_outputs[0] # Replace tokens with image embeddings image_token_id = tokenizer.convert_tokens_to_ids("") image_positions = np.where(input_ids[0] == image_token_id)[0] for i, pos in enumerate(image_positions): if i < len(image_embeds): token_embeds[0, pos] = image_embeds[i] # Initialize KV cache cache = self._create_empty_cache() seq_len = token_embeds.shape[1] generated_tokens = [] for step in range(max_tokens): if step == 0: embeds = token_embeds.astype(np.float32) else: last_token = np.array([[generated_tokens[-1]]], dtype=np.int64) embeds = self._embed_tokens.run(None, {"input_ids": last_token})[0].astype(np.float32) attn_mask = np.ones((1, seq_len + len(generated_tokens)), dtype=np.int64) feed = {"inputs_embeds": embeds, "attention_mask": attn_mask, **cache} outputs = self._decoder.run(None, feed) # Get logits and apply temperature logits = outputs[0][0, -1] if temperature > 0: logits = logits / temperature # Apply top-k if top_k > 0: indices_to_remove = np.argsort(logits)[:-top_k] logits[indices_to_remove] = -np.inf # Apply top-p (nucleus sampling) if top_p < 1.0: sorted_indices = np.argsort(logits)[::-1] sorted_logits = logits[sorted_indices] probs = np.exp(sorted_logits - np.max(sorted_logits)) probs = probs / probs.sum() cumulative_probs = np.cumsum(probs) sorted_indices_to_remove = cumulative_probs > top_p sorted_indices_to_remove[1:] = sorted_indices_to_remove[:-1].copy() sorted_indices_to_remove[0] = False indices_to_remove = sorted_indices[sorted_indices_to_remove] logits[indices_to_remove] = -np.inf # Sample probs = np.exp(logits - np.max(logits)) probs = probs / probs.sum() next_token = int(np.random.choice(len(probs), p=probs)) else: next_token = int(np.argmax(logits)) generated_tokens.append(next_token) # Update cache for i, out in enumerate(self._decoder.get_outputs()[1:], 1): name = out.name.replace("present_conv", "past_conv").replace("present.", "past_key_values.") if name in cache: cache[name] = outputs[i] if next_token in stop_tokens: break return generated_tokens def generate_stream( self, images: List[Image.Image], messages: List[dict], max_tokens: int = 2000, temperature: float = 0.1, top_k: int = 50, top_p: float = 0.1, stop_tokens: Optional[List[int]] = None ): """Streaming generation for Vision model.""" tokenizer = self._processor.tokenizer if stop_tokens is None: stop_tokens = [tokenizer.eos_token_id] # Process inputs through processor prompt = self._processor.apply_chat_template(messages, add_generation_prompt=True) inputs = self._processor( images=images if images else None, text=prompt, return_tensors="pt" ) # Convert to numpy with correct dtypes input_ids = inputs["input_ids"].numpy().astype(np.int64) # Get token embeddings token_outputs = self._embed_tokens.run(None, {"input_ids": input_ids}) token_embeds = token_outputs[0] # Process images if present if images and "pixel_values" in inputs: pixel_values = inputs["pixel_values"].numpy().astype(np.float32) pixel_attention_mask = inputs.get("pixel_attention_mask", None) spatial_shapes = inputs.get("spatial_shapes", None) image_feed = {"pixel_values": pixel_values} if pixel_attention_mask is not None: image_feed["pixel_attention_mask"] = pixel_attention_mask.numpy().astype(np.int64) if spatial_shapes is not None: image_feed["spatial_shapes"] = spatial_shapes.numpy().astype(np.int64) image_outputs = self._embed_images.run(None, image_feed) image_embeds = image_outputs[0] # Replace tokens with image embeddings image_token_id = tokenizer.convert_tokens_to_ids("") image_positions = np.where(input_ids[0] == image_token_id)[0] for i, pos in enumerate(image_positions): if i < len(image_embeds): token_embeds[0, pos] = image_embeds[i] # Initialize KV cache cache = self._create_empty_cache() seq_len = token_embeds.shape[1] generated_tokens = [] # Pre-allocate attention mask max_possible_len = seq_len + max_tokens attn_mask = np.ones((1, max_possible_len), dtype=np.int64) # Pre-compute flags use_temp = temperature > 0 use_top_k = top_k > 0 use_top_p = top_p < 1.0 feed = {} for step in range(max_tokens): current_len = seq_len + step if step == 0: embeds = token_embeds.astype(np.float32) else: last_token = np.array([[generated_tokens[-1]]], dtype=np.int64) embeds = self._embed_tokens.run(None, {"input_ids": last_token})[0].astype(np.float32) # Update Feed Dict feed.clear() feed["inputs_embeds"] = embeds feed["attention_mask"] = attn_mask[:, :current_len] feed.update(cache) # Inference outputs = self._decoder.run(None, feed) logits = outputs[0][0, -1] # Sampling if use_temp: logits /= temperature if use_top_k and top_k < len(logits): top_k_idx = np.argpartition(logits, -top_k)[-top_k:] mask = np.ones(logits.shape, dtype=bool) mask[top_k_idx] = False logits[mask] = -np.inf if use_top_p: valid_mask = logits > -np.inf if valid_mask.any(): valid_logits = logits[valid_mask] valid_indices = np.where(valid_mask)[0] sorted_indices = np.argsort(valid_logits)[::-1] sorted_logits = valid_logits[sorted_indices] exp_logits = np.exp(sorted_logits - np.max(sorted_logits)) probs = exp_logits / exp_logits.sum() cumulative = np.cumsum(probs) cutoff = np.searchsorted(cumulative, top_p) cutoff = min(cutoff + 1, len(sorted_logits)) accepted_indices = sorted_indices[:cutoff] accepted_probs = probs[:cutoff] accepted_probs /= accepted_probs.sum() sample_idx = np.searchsorted(np.cumsum(accepted_probs), np.random.rand()) next_token = int(valid_indices[accepted_indices[sample_idx]]) else: next_token = int(np.argmax(logits)) else: valid_mask = logits > -np.inf valid_logits = logits[valid_mask] valid_indices = np.where(valid_mask)[0] exp_logits = np.exp(valid_logits - np.max(valid_logits)) probs = exp_logits / exp_logits.sum() sample_idx = np.searchsorted(np.cumsum(probs), np.random.rand()) next_token = int(valid_indices[sample_idx]) else: next_token = int(np.argmax(logits)) generated_tokens.append(next_token) yield next_token if next_token in stop_tokens: break # Update Cache for i, out in enumerate(self._decoder.get_outputs()[1:], 1): name = out.name.replace("present_conv", "past_conv").replace("present.", "past_key_values.") if name in cache: cache[name] = outputs[i] def unload(self) -> None: """Unload models from memory.""" with self._lock: if self._embed_tokens is not None: del self._embed_tokens del self._embed_images del self._decoder del self._processor self._embed_tokens = None self._embed_images = None self._decoder = None self._processor = None logger.info("Models unloaded") # Global model manager model_manager = ONNXVisionModelManager() # ============================================================================== # Image Processing Utilities # ============================================================================== def resize_image_for_model(image: Image.Image, max_dim: int = 512) -> Image.Image: """ Resize image to max dimension while preserving aspect ratio. Uses LANCZOS (highest quality) resampling for best visual fidelity. This optimization ensures: - Consistent processing time (~3-4s) regardless of input size - Single-patch processing (256 tokens) instead of tiling - Reduced memory usage Args: image: PIL Image to resize max_dim: Maximum dimension (width or height), default 512 Returns: Resized PIL Image (or original if already small enough) """ width, height = image.size # Skip if already small enough if width <= max_dim and height <= max_dim: logger.debug(f"Image {width}x{height} already within {max_dim}px limit") return image # Calculate new dimensions (preserve aspect ratio) ratio = min(max_dim / width, max_dim / height) new_width = int(width * ratio) new_height = int(height * ratio) logger.info(f"Resizing image: {width}x{height} → {new_width}x{new_height} (LANCZOS)") # Resize with high-quality LANCZOS filter return image.resize((new_width, new_height), Image.Resampling.LANCZOS) async def load_image_from_url(url: str) -> Image.Image: """Load image from URL, convert to RGB, and resize for optimal processing.""" async with aiohttp.ClientSession() as session: async with session.get(url) as response: if response.status != 200: raise HTTPException(status_code=400, detail=f"Failed to fetch image from URL: {url}") data = await response.read() image = Image.open(io.BytesIO(data)) # Convert to RGB to ensure consistent channel format if image.mode != 'RGB': image = image.convert('RGB') # Resize for optimal model processing (max 512x512) image = resize_image_for_model(image) return image def load_image_from_base64(data_url: str) -> Image.Image: """Load image from base64 data URL, convert to RGB, and resize for optimal processing.""" # Format: data:image/jpeg;base64,/9j/4AAQ... if "," in data_url: header, encoded = data_url.split(",", 1) else: encoded = data_url image_data = base64.b64decode(encoded) image = Image.open(io.BytesIO(image_data)) # Convert to RGB to ensure consistent channel format if image.mode != 'RGB': image = image.convert('RGB') # Resize for optimal model processing (max 512x512) image = resize_image_for_model(image) return image async def process_image_content(content: Union[ImageContent, dict]) -> Optional[Image.Image]: """Process image content from request.""" if isinstance(content, dict): content = ImageContent(**content) if content.type != "image": return None if not content.image_url: return None url = content.image_url # Check if it's a base64 data URL if url.startswith("data:"): return load_image_from_base64(url) else: # It's a regular URL return await load_image_from_url(url) # ============================================================================== # Application Lifecycle # ============================================================================== @asynccontextmanager async def lifespan(app: FastAPI): """Application lifespan handler.""" logger.info("Starting LFM2.5-VL Vision API Server (ONNX Runtime)...") loop = asyncio.get_event_loop() await loop.run_in_executor(None, model_manager.load_model) yield logger.info("Shutting down...") model_manager.unload() # ============================================================================== # FastAPI Application # ============================================================================== app = FastAPI( title=settings.app_name, description="Fast CPU inference for LiquidAI LFM2.5-VL-1.6B Vision-Language model using ONNX Runtime", version=settings.app_version, lifespan=lifespan, docs_url="/docs", redoc_url="/redoc", ) origins = [ "http://127.0.0.1:5500", "http://127.0.0.1:5501", "http://localhost:5500", "http://localhost:5173", "https://toolboxesai.com" ] app.add_middleware( CORSMiddleware, allow_origins=origins, allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) @app.middleware("http") async def add_cors_for_null_origin(request: Request, call_next): """Handle CORS for null origin (when HTML is opened from file://).""" origin = request.headers.get("origin", "") response = await call_next(request) if origin == "null" or not origin: response.headers["Access-Control-Allow-Origin"] = "*" response.headers["Access-Control-Allow-Methods"] = "GET, POST, PUT, DELETE, OPTIONS" response.headers["Access-Control-Allow-Headers"] = "*" response.headers["Access-Control-Expose-Headers"] = "*" return response # ============================================================================== # Helper Functions # ============================================================================== def generate_id() -> str: return f"chatcmpl-{uuid.uuid4().hex[:12]}" async def extract_images_and_text(messages: List[VisionMessage]) -> tuple[List[Image.Image], List[dict]]: """Extract images and convert messages to processor format.""" images = [] processed_messages = [] for msg in messages: if isinstance(msg.content, str): # Simple text message processed_messages.append({ "role": msg.role, "content": msg.content }) else: # Mixed content (images + text) content_parts = [] for item in msg.content: if isinstance(item, dict): item_type = item.get("type", "") else: item_type = item.type if item_type == "image": image = await process_image_content(item) if image: images.append(image) content_parts.append({"type": "image"}) elif item_type == "text": text = item.get("text", "") if isinstance(item, dict) else item.text content_parts.append({"type": "text", "text": text}) processed_messages.append({ "role": msg.role, "content": content_parts }) return images, processed_messages async def stream_vision_completion(request: VisionCompletionRequest) -> AsyncGenerator[str, None]: """Streaming vision completion.""" request_id = generate_id() created = int(time.time()) loop = asyncio.get_running_loop() async_queue = asyncio.Queue() # Extract images and process messages images, processed_messages = await extract_images_and_text(request.messages) tokenizer = model_manager.processor.tokenizer # Config max_tokens = request.max_tokens or settings.max_tokens temperature = request.temperature if request.temperature is not None else settings.temperature top_k = request.top_k if request.top_k is not None else settings.top_k top_p = request.top_p if request.top_p is not None else settings.top_p # Prepare stop tokens stop_tokens = [tokenizer.eos_token_id] if request.stop: if isinstance(request.stop, str): encoded = tokenizer.encode(request.stop, add_special_tokens=False) if encoded: stop_tokens.append(encoded[0]) elif isinstance(request.stop, list): for stop_str in request.stop: encoded = tokenizer.encode(stop_str, add_special_tokens=False) if encoded: stop_tokens.append(encoded[0]) def generate_tokens(): try: for token in model_manager.generate_stream( images, processed_messages, max_tokens=max_tokens, temperature=temperature, top_k=top_k, top_p=top_p, stop_tokens=stop_tokens ): loop.call_soon_threadsafe(async_queue.put_nowait, ("token", token)) except Exception as e: logger.error(f"Stream generation error: {e}") loop.call_soon_threadsafe(async_queue.put_nowait, ("error", str(e))) finally: loop.call_soon_threadsafe(async_queue.put_nowait, ("done", None)) threading.Thread(target=generate_tokens, daemon=True).start() try: while True: msg_type, data = await async_queue.get() if msg_type == "token": text = tokenizer.decode([data], skip_special_tokens=True) if text: chunk = { "id": request_id, "object": "chat.completion.chunk", "created": created, "model": request.model, "choices": [{ "index": 0, "delta": {"content": text}, "finish_reason": None }] } yield {"data": json.dumps(chunk)} elif msg_type == "done": final = { "id": request_id, "object": "chat.completion.chunk", "created": created, "model": request.model, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}] } yield {"data": json.dumps(final)} yield {"data": "[DONE]"} break elif msg_type == "error": logger.error(f"Stream error: {data}") yield {"data": json.dumps({"error": {"message": data}})} break except asyncio.CancelledError: logger.info(f"Stream cancelled for request {request_id[:8]}") raise except Exception as e: logger.error(f"Streaming error: {e}") yield {"data": json.dumps({"error": {"message": str(e)}})} # ============================================================================== # API Endpoints # ============================================================================== @app.get("/health") async def health(): if not model_manager.is_loaded: raise HTTPException(status_code=503, detail="Model not loaded") return {"status": "healthy"} @app.get("/v1/models", response_model=ModelListResponse) async def list_models(): return ModelListResponse( data=[ ModelInfo(id="lfm-vision", created=int(time.time())), ModelInfo(id="lfm-2.5-vl-1.6b-onnx", created=int(time.time())) ] ) @app.post("/v1/vision/completions") async def vision_completions(request: VisionCompletionRequest): """Vision-language completion with image support.""" if not model_manager.is_loaded: raise HTTPException(status_code=503, detail="Model not loaded") if request.stream: return EventSourceResponse( stream_vision_completion(request), media_type="text/event-stream", ping=30000, ping_message_factory=lambda: '{"type": "ping"}' ) try: # Extract images and process messages images, processed_messages = await extract_images_and_text(request.messages) tokenizer = model_manager.processor.tokenizer max_tokens = request.max_tokens or settings.max_tokens temperature = request.temperature if request.temperature is not None else settings.temperature top_k = request.top_k if request.top_k is not None else settings.top_k top_p = request.top_p if request.top_p is not None else settings.top_p start_time = time.time() loop = asyncio.get_event_loop() tokens = await loop.run_in_executor( None, lambda: model_manager.generate( images, processed_messages, max_tokens=max_tokens, temperature=temperature, top_k=top_k, top_p=top_p ) ) response_text = tokenizer.decode(tokens, skip_special_tokens=True) gen_time = time.time() - start_time logger.debug(f"Generated {len(tokens)} tokens in {gen_time:.2f}s") return ChatCompletionResponse( id=generate_id(), created=int(time.time()), model=request.model, choices=[ ChatCompletionChoice( index=0, message=ChatMessage(role="assistant", content=response_text), finish_reason="stop" ) ], usage={ "prompt_tokens": 0, # Would need to track input tokens "completion_tokens": len(tokens), "total_tokens": len(tokens) } ) except Exception as e: logger.error(f"Vision completion error: {e}") raise HTTPException(status_code=500, detail=str(e)) @app.post("/v1/chat/completions") async def chat_completions(request: ChatCompletionRequest): """Text-only chat completion (for compatibility).""" if not model_manager.is_loaded: raise HTTPException(status_code=503, detail="Model not loaded") # Convert to vision request format (no images) vision_messages = [ VisionMessage(role=m.role, content=m.content) for m in request.messages ] vision_request = VisionCompletionRequest( model=request.model, messages=vision_messages, temperature=request.temperature, top_p=request.top_p, top_k=request.top_k, max_tokens=request.max_tokens, stream=request.stream ) return await vision_completions(vision_request) @app.post("/v1/vision/upload") async def upload_image( file: UploadFile = File(...), prompt: str = "What is in this image?" ): """Direct image upload endpoint.""" if not model_manager.is_loaded: raise HTTPException(status_code=503, detail="Model not loaded") # Validate file type content_type = file.content_type or "" file_ext = Path(file.filename or "").suffix.lower().lstrip(".") if file_ext not in settings.supported_formats and not any(fmt in content_type for fmt in settings.supported_formats): raise HTTPException( status_code=400, detail=f"Unsupported image format. Supported: {settings.supported_formats}" ) # Read and process image contents = await file.read() if len(contents) > settings.max_image_size_mb * 1024 * 1024: raise HTTPException( status_code=400, detail=f"Image too large. Max size: {settings.max_image_size_mb}MB" ) try: image = Image.open(io.BytesIO(contents)) except Exception as e: raise HTTPException(status_code=400, detail=f"Invalid image: {e}") # Create request messages = [{ "role": "user", "content": [ {"type": "image"}, {"type": "text", "text": prompt} ] }] tokenizer = model_manager.processor.tokenizer tokens = model_manager.generate( [image], messages, max_tokens=settings.max_tokens, temperature=settings.temperature, top_k=settings.top_k, top_p=settings.top_p ) response_text = tokenizer.decode(tokens, skip_special_tokens=True) return { "id": generate_id(), "model": "lfm-vision", "response": response_text } # ============================================================================== # Run Server # ============================================================================== if __name__ == "__main__": import uvicorn logger.info(f"Starting server on {settings.host}:{settings.port}") uvicorn.run( "app:app", host=settings.host, port=settings.port, reload=False, log_level=settings.log_level )