Spaces:
Running
Running
| 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 ──────────────────────────────────────────────────────────────────── | |
| 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": "<base64_string_here>", | |
| "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." | |
| } | |
| }) | |
| def health(): | |
| return jsonify({"status": "ok", "model_status": model_status}) | |
| 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) |