server.py

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": "<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."
        }
    })


@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)