import os import io import base64 import ctypes import threading import json import time from flask import Flask, request, jsonify, Response from flask_cors import CORS # --- KEY CHANGE: Balik sa 2B Model --- HF_REPO = "litert-community/gemma-4-E2B-it-litert-lm" HF_FILE = "gemma-4-E2B-it.litertlm" _SERVER_DIR = os.path.dirname(os.path.abspath(__file__)) _DEFAULT_PATH = os.path.join(_SERVER_DIR, "models", "gemma", 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: pass # 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_status = "loading" engine = None _engine_ctx = None engine_lock = threading.Lock() app = Flask(__name__) CORS(app) # ─── Model loading ───────────────────────────────────────────────────────────── def load_model(): global engine, model_status, _engine_ctx 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.set_min_log_severity(_lm.LogSeverity.SILENT) except ImportError: print("[INFO] litert_lm not installed — no model loaded", flush=True) model_status = "no_litert_lm" return 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: _engine_ctx = _lm.Engine( MODEL_PATH, backend=_lm.interfaces.CPU(), vision_backend=_lm.interfaces.CPU(), ) engine = _engine_ctx.__enter__() 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" # ─── Fallback functions (No Model) ────────────────────────────────────────────── def _analyze_image(image_bytes: bytes) -> dict: from PIL import Image img = Image.open(io.BytesIO(image_bytes)) w, h = img.size fmt = (img.format or "image").lower() return dict(w=w, h=h, fmt=fmt) def _describe_image(ask: str, image_bytes: bytes) -> str: try: i = _analyze_image(image_bytes) return f"[MOCK] This is a {i['fmt']} image ({i['w']}×{i['h']} px). Connect the real model for full vision analysis." except Exception as e: return f"Could not analyze image: {e}" def _text_reply(ask: str) -> str: return f"[MOCK] Hello! I received: '{ask}'. Connect the Gemma model to see real answers." # ─── Real model inference (litert_lm Generator) ──────────────────────────────── def _run_real_model_generator(ask: str, image_bytes: bytes | None): """Yields text chunks as they are generated by the model.""" import litert_lm # engine_lock ensures only 1 request processes at a time to prevent RAM crashes with engine_lock: with engine.create_conversation() as conv: if image_bytes: msg = litert_lm.Contents.of( litert_lm.Content.ImageBytes(image_bytes), litert_lm.Content.Text(ask), ) else: msg = ask for chunk in conv.send_message_async(msg): for part in chunk.get("content", []): if part.get("type") == "text": text = part.get("text", "") if text: yield text # ─── Request extraction ──────────────────────────────────────────────────────── def extract_request() -> tuple[str, bytes | None, bool]: # Check query params first (e.g. ?stream=true) stream = str(request.args.get("stream", "")).lower() in ["true", "1", "yes"] if request.method == "GET": return request.args.get("ask", "").strip(), None, stream ct = request.content_type or "" # Handle Form Data (File Uploads) if "multipart/form-data" in ct: ask = request.form.get("ask", "").strip() f = request.files.get("image") if "stream" in request.form: stream = str(request.form.get("stream")).lower() in ["true", "1", "yes"] return ask, (f.read() if f else None), stream # Handle JSON data = request.get_json(silent=True) or {} ask = data.get("ask", "").strip() image_bytes = None raw = data.get("image", "") if "stream" in data: stream = str(data.get("stream")).lower() in ["true", "1", "yes"] if raw: try: if "," in raw: raw = raw.split(",", 1)[1] image_bytes = base64.b64decode(raw) except Exception: pass return ask, image_bytes, stream # ─── Routes ──────────────────────────────────────────────────────────────────── @app.route("/") def index(): return jsonify({ "service": "Gemma 4 API (2B Model)", "model_status": model_status, "guide": { "1. Text Only Chat": { "GET_Example": "/gemma?ask=Hello&stream=false", "POST_JSON": {"ask": "What is AI?", "stream": True} }, "2. Image with Text": { "POST_JSON": { "ask": "Describe this image", "image": "", "stream": False }, "POST_FormData": { "ask": "What color is this?", "image": "@file.jpg (File Upload)", "stream": "true" } }, "Streaming Info": "Set 'stream': true to receive Server-Sent Events (SSE) by token. Set false to wait for 1 full JSON response." } }) @app.route("/health") def health(): return jsonify({"status": "ok", "model_status": model_status}) @app.route("/gemma", methods=["GET", "POST"]) def gemma(): ask, image_bytes, stream = extract_request() if not ask: return jsonify({"error": "Missing 'ask' parameter"}), 400 # Fallback if model isn't ready if engine is None or model_status != "ready": fallback_msg = _describe_image(ask, image_bytes) if image_bytes else _text_reply(ask) if stream: def mock_stream(): for word in fallback_msg.split(): yield f"data: {json.dumps({'text': word + ' '})}\n\n" time.sleep(0.05) yield "data: [DONE]\n\n" return Response(mock_stream(), mimetype="text/event-stream") else: return jsonify({"ask": ask, "response": fallback_msg, "has_image": bool(image_bytes)}) # Real Model Logic if stream: def generate_stream(): try: for text_chunk in _run_real_model_generator(ask, image_bytes): yield f"data: {json.dumps({'text': text_chunk})}\n\n" yield "data: [DONE]\n\n" except Exception as e: yield f"data: {json.dumps({'error': str(e)})}\n\n" return Response(generate_stream(), mimetype="text/event-stream") else: try: full_text = "".join(list(_run_real_model_generator(ask, image_bytes))) return jsonify({ "ask": ask, "response": full_text, "has_image": image_bytes is not None, "model_status": model_status }) except Exception as e: return jsonify({"error": f"Model error: {e}"}), 500 # ─── Entry ───────────────────────────────────────────────────────────────────── if __name__ == "__main__": port = int(os.environ.get("PORT", 5173)) threading.Thread(target=load_model, daemon=True).start() print(f"[INFO] Gemma API on :{port}", flush=True) app.run(host="0.0.0.0", port=port, debug=False)