import os import io import base64 import ctypes import threading import json import time import uuid from flask import Flask, request, jsonify, Response from flask_cors import CORS # --- Model Configuration --- HF_REPO = "paulsp94/Qwen3.5-2B-LiteRT-LM" HF_FILE = "model.litertlm" _SERVER_DIR = os.path.dirname(os.path.abspath(__file__)) _DEFAULT_PATH = os.path.join(_SERVER_DIR, "models", "qwen", HF_FILE) # litert_lm links against libvulkan.so.1 even on CPU-only runs. _vk_stub = os.path.join(_SERVER_DIR, "libvulkan.so.1") if os.path.exists(_vk_stub): try: ctypes.CDLL(_vk_stub, mode=ctypes.RTLD_GLOBAL) except OSError as e: print(f"[WARN] Could not preload vulkan stub: {e}", flush=True) # Suppress verbose C++ logs from litert_lm os.environ.setdefault("GLOG_minloglevel", "3") MODEL_PATH = os.environ.get("GEMMA_MODEL_PATH", _DEFAULT_PATH).strip() MODEL_ID = "qwen3.5-2b" model_status = "loading" engine = None _lm = None engine_lock = threading.Lock() app = Flask(__name__) CORS(app) # ─── Model loading ───────────────────────────────────────────────────────────── def load_model(): global engine, model_status, _lm if not MODEL_PATH: print("[INFO] GEMMA_MODEL_PATH not set — no model loaded", flush=True) model_status = "no_model_path" return try: import litert_lm as lm _lm = lm except ImportError: print("[INFO] litert_lm not installed — running in mock mode", flush=True) model_status = "no_litert_lm" return # Try to silence logs if the API exists try: _lm.set_min_log_severity(_lm.LogSeverity.SILENT) except Exception: pass if not os.path.exists(MODEL_PATH): print(f"[WARN] Model file not found: {MODEL_PATH}", flush=True) model_status = "model_file_missing" return try: # The litert_lm Engine API. Build args defensively depending # on what the installed version exposes. try: cpu_backend = _lm.interfaces.CPU() engine = _lm.Engine( MODEL_PATH, backend=cpu_backend, vision_backend=cpu_backend, ) except (AttributeError, TypeError): # Fallback: simpler constructor signature engine = _lm.Engine(MODEL_PATH) model_status = "ready" print(f"[INFO] Model ready → {MODEL_PATH}", flush=True) except Exception as e: print(f"[ERROR] Failed to load model: {e}", flush=True) model_status = "error" # ─── OpenAI Request Parsing ──────────────────────────────────────────────────── def parse_openai_messages(messages: list): """Parses OpenAI formatted messages into a flat text prompt and optional image.""" prompt_text = "" image_bytes = None for msg in messages: role = msg.get("role", "user") content = msg.get("content", "") if isinstance(content, str): prompt_text += f"{role}: {content}\n" elif isinstance(content, list): prompt_text += f"{role}:\n" for part in content: ptype = part.get("type") if ptype == "text": prompt_text += part.get("text", "") + "\n" elif ptype == "image_url": url = part.get("image_url", {}).get("url", "") if url.startswith("data:image"): try: b64_data = url.split(",", 1)[1] image_bytes = base64.b64decode(b64_data) except Exception as e: print(f"[WARN] Failed to decode base64 image: {e}", flush=True) prompt_text += "assistant: " return prompt_text.strip(), image_bytes # ─── Inference Engine ────────────────────────────────────────────────────────── def _run_real_model_generator(ask: str, image_bytes): """Yields text chunks as they are generated by the model.""" # Qwen 3.5 2B is text-only; image_bytes are ignored intentionally. with engine_lock: conv = None try: conv = engine.create_conversation() # Support both context-manager and plain object styles if hasattr(conv, "__enter__"): conv_obj = conv.__enter__() else: conv_obj = conv stream = conv_obj.send_message_async(ask) for chunk in stream: # Chunk may be a plain string or a structured dict if isinstance(chunk, str): if chunk: yield chunk elif isinstance(chunk, dict): for part in chunk.get("content", []): if part.get("type") == "text": text = part.get("text", "") if text: yield text else: # Try common attribute names text = getattr(chunk, "text", None) if text: yield text finally: if conv is not None and hasattr(conv, "__exit__"): try: conv.__exit__(None, None, None) except Exception: pass def _run_mock_generator(ask: str, has_image: bool): """Fallback generator when the model is missing/loading.""" msg = (f"[MOCK] Model status: {model_status}. " f"Vision included: {has_image}. " f"Connect litert_lm + model file for real output.") for word in msg.split(): yield word + " " time.sleep(0.02) # ─── Routes ──────────────────────────────────────────────────────────────────── @app.route("/health", methods=["GET"]) def health(): return jsonify({"status": model_status}), 200 @app.route("/v1/models", methods=["GET"]) def list_models(): return jsonify({ "object": "list", "data": [{ "id": MODEL_ID, "object": "model", "created": int(time.time()), "owned_by": "paulsp94" }] }) @app.route("/v1/chat/completions", methods=["POST"]) def chat_completions(): data = request.get_json(silent=True) or {} messages = data.get("messages", []) stream = bool(data.get("stream", False)) if not messages: return jsonify({ "error": {"message": "Missing 'messages' array", "type": "invalid_request_error"} }), 400 ask, image_bytes = parse_openai_messages(messages) use_mock = engine is None or model_status != "ready" req_model = data.get("model", MODEL_ID) cmpl_id = f"chatcmpl-{uuid.uuid4().hex}" created_time = int(time.time()) if stream: def stream_response(): init_chunk = { "id": cmpl_id, "object": "chat.completion.chunk", "created": created_time, "model": req_model, "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}] } yield f"data: {json.dumps(init_chunk)}\n\n" try: gen = (_run_mock_generator(ask, bool(image_bytes)) if use_mock else _run_real_model_generator(ask, image_bytes)) for text_chunk in gen: chunk = { "id": cmpl_id, "object": "chat.completion.chunk", "created": created_time, "model": req_model, "choices": [{"index": 0, "delta": {"content": text_chunk}, "finish_reason": None}] } yield f"data: {json.dumps(chunk)}\n\n" except Exception as e: err_chunk = { "id": cmpl_id, "object": "chat.completion.chunk", "created": created_time, "model": req_model, "choices": [{"index": 0, "delta": {"content": f"[ERROR] {e}"}, "finish_reason": "stop"}] } yield f"data: {json.dumps(err_chunk)}\n\n" final_chunk = { "id": cmpl_id, "object": "chat.completion.chunk", "created": created_time, "model": req_model, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}] } yield f"data: {json.dumps(final_chunk)}\n\n" yield "data: [DONE]\n\n" return Response(stream_response(), mimetype="text/event-stream") else: try: gen = (_run_mock_generator(ask, bool(image_bytes)) if use_mock else _run_real_model_generator(ask, image_bytes)) full_text = "".join(gen) response = { "id": cmpl_id, "object": "chat.completion", "created": created_time, "model": req_model, "choices": [{ "index": 0, "message": {"role": "assistant", "content": full_text}, "finish_reason": "stop" }], "usage": { "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0 } } return jsonify(response) except Exception as e: return jsonify({ "error": {"message": f"Model error: {e}", "type": "server_error"} }), 500 # ─── Entry ───────────────────────────────────────────────────────────────────── if __name__ == "__main__": port = int(os.environ.get("PORT", 7860)) threading.Thread(target=load_model, daemon=True).start() print(f"[INFO] Qwen 3.5 2B OpenAI-Compatible API listening on :{port}", flush=True) app.run(host="0.0.0.0", port=port, debug=False, threaded=True)