Upload server.py with huggingface_hub

server.py

@@ -0,0 +1,356 @@
+import os
+import io
+import base64
+import ctypes
+import threading
+from flask import Flask, request, jsonify, Response
+from flask_cors import CORS
+
+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.
+# Pre-load a stub so the dynamic linker is satisfied without a real GPU.
+_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"
+
+
+# ─── Image analysis via Pillow (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()
+
+    thumb  = img.convert("RGB").resize((64, 64))
+    pixels = list(thumb.getdata())
+    n      = len(pixels)
+    r = sum(p[0] for p in pixels) // n
+    g = sum(p[1] for p in pixels) // n
+    b = sum(p[2] for p in pixels) // n
+
+    lum = (r * 299 + g * 587 + b * 114) // 1000
+    tone = "bright" if lum > 200 else "medium" if lum > 130 else "dark" if lum > 60 else "very dark"
+
+    diff = max(r, g, b) - min(r, g, b)
+    if diff < 25:
+        hue = "neutral / grayscale"
+    elif max(r, g, b) == r and r - g > 30:
+        hue = "red / warm"
+    elif max(r, g, b) == g and g - b > 20:
+        hue = "green"
+    elif max(r, g, b) == b:
+        hue = "blue / cool"
+    elif r > 180 and g > 150 and b < 100:
+        hue = "yellow / orange"
+    elif r > 150 and b > 150 and g < 100:
+        hue = "purple / violet"
+    else:
+        hue = "mixed"
+
+    ratio       = w / h if h else 1
+    orientation = "landscape" if ratio > 1.4 else "portrait" if ratio < 0.72 else "square"
+
+    return dict(w=w, h=h, fmt=fmt, r=r, g=g, b=b, tone=tone, hue=hue, orientation=orientation)
+
+
+def _describe_image(ask: str, image_bytes: bytes) -> str:
+    try:
+        i   = _analyze_image(image_bytes)
+        q   = ask.lower()
+        out = [
+            f"This is a {i['orientation']} {i['fmt']} image ({i['w']}×{i['h']} px).",
+            f"Overall tone: {i['tone']}. Dominant color: {i['hue']}. Average RGB: ({i['r']}, {i['g']}, {i['b']}).",
+        ]
+        if any(w in q for w in ["color", "colour", "kulay"]):
+            out.append(f"Main color: {i['hue']} — RGB({i['r']}, {i['g']}, {i['b']}).")
+        elif any(w in q for w in ["size", "dimension", "laki", "sukat"]):
+            out.append(f"Dimensions: {i['w']}×{i['h']} pixels.")
+        elif any(w in q for w in ["bright", "dark", "liwanag"]):
+            out.append(f"Brightness: {i['tone']}.")
+
+        out.append(
+            f"For full scene/object understanding, load the real model: "
+            f"GET /gemma/download  or  set GEMMA_MODEL_PATH."
+        )
+        return "\n".join(out)
+    except Exception as e:
+        return f"Could not analyze image: {e}"
+
+
+# ─── Fallback text replies (no model) ─────────────────────────────────────────
+
+def _hint() -> str:
+    if model_status == "no_litert_lm":
+        return "litert_lm is not installed. Run: pip install litert-lm"
+    if model_status in ("no_model_path", "model_file_missing"):
+        return (f"Model not found at '{MODEL_PATH}'. "
+                f"Run: GET /gemma/download  to fetch it automatically.")
+    return "Connect a Gemma 4 model to enable full responses."
+
+_REPLIES = [
+    lambda h: f"Hello! I'm the Gemma 4 API. {h}",
+    lambda h: f"Gemma 4 is Google's open multimodal model — text + image input, runs on-device via LiteRT. {h}",
+    lambda h: f"Send images as base64 JSON (field 'image') or multipart/form-data file upload. {h}",
+    lambda h: f"API: GET /gemma?ask=... or POST /gemma {{\"ask\":\"...\",\"image\":\"<base64>\"}}. {h}",
+]
+_idx = 0
+
+def _text_reply(ask: str) -> str:
+    global _idx
+    h, q = _hint(), ask.lower()
+    if any(w in q for w in ["hello", "hi", "hey", "kumusta"]):
+        return _REPLIES[0](h)
+    if any(w in q for w in ["gemma", "model", "what are you"]):
+        return _REPLIES[1](h)
+    if any(w in q for w in ["image", "photo", "picture", "larawan"]):
+        return _REPLIES[2](h)
+    if any(w in q for w in ["how", "api", "endpoint", "use", "query"]):
+        return _REPLIES[3](h)
+    r = _REPLIES[_idx % len(_REPLIES)](h)
+    _idx += 1
+    return r
+
+
+# ─── Real model inference (litert_lm) ─────────────────────────────────────────
+
+def _run_real_model(ask: str, image_bytes: bytes | None) -> str:
+    import litert_lm
+    with engine_lock:
+        try:
+            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
+
+                out = []
+                for chunk in conv.send_message_async(msg):
+                    for part in chunk.get("content", []):
+                        if part.get("type") == "text":
+                            out.append(part.get("text", ""))
+                return "".join(out) or "(empty response)"
+        except Exception as e:
+            return f"Model error: {e}"
+
+
+def run_model(ask: str, image_bytes: bytes | None) -> str:
+    if engine is not None and model_status == "ready":
+        return _run_real_model(ask, image_bytes)
+    if image_bytes:
+        return _describe_image(ask, image_bytes)
+    return _text_reply(ask)
+
+
+# ─── Request extraction ────────────────────────────────────────────────────────
+
+def extract_request() -> tuple[str, bytes | None]:
+    if request.method == "GET":
+        return request.args.get("ask", "").strip(), None
+
+    ct = request.content_type or ""
+    if "multipart/form-data" in ct:
+        ask = request.form.get("ask", "").strip()
+        f   = request.files.get("image")
+        return ask, (f.read() if f else None)
+
+    data        = request.get_json(silent=True) or {}
+    ask         = data.get("ask", "").strip()
+    image_bytes = None
+    raw         = data.get("image", "")
+    if raw:
+        try:
+            if "," in raw:
+                raw = raw.split(",", 1)[1]
+            image_bytes = base64.b64decode(raw)
+        except Exception:
+            pass
+    return ask, image_bytes
+
+
+# ─── Routes ────────────────────────────────────────────────────────────────────
+
+@app.route("/favicon.ico")
+def favicon():
+    return "", 204
+
+
+@app.route("/")
+def index():
+    return jsonify({
+        "service":      "Gemma 4 API",
+        "model_status": model_status,
+        "model_path":   MODEL_PATH or None,
+        "endpoints": {
+            "GET  /gemma?ask=hello":             "Text query",
+            "POST /gemma  {ask, image?}":        "JSON body — text + optional base64 image",
+            "POST /gemma  multipart/form-data":  "File upload — text + optional image file",
+            "GET  /gemma/download":              "Download Gemma 4 model from HuggingFace",
+            "GET  /health":                      "Health check",
+        },
+    })
+
+
+@app.route("/health")
+@app.route("/api/health")
+def health():
+    info = {"status": "ok", "model_status": model_status}
+    if MODEL_PATH:
+        info["model_path"] = MODEL_PATH
+    return jsonify(info)
+
+
+@app.route("/gemma", methods=["GET", "POST"])
+def gemma():
+    ask, image_bytes = extract_request()
+    if not ask:
+        return jsonify({"error": "Missing 'ask' parameter"}), 400
+    response = run_model(ask, image_bytes)
+    return jsonify({
+        "ask":          ask,
+        "response":     response,
+        "has_image":    image_bytes is not None,
+        "model_status": model_status,
+    })
+
+
+# ─── Download state ────────────────────────────────────────────────────────────
+
+_dl: dict = {"status": "idle", "path": None, "error": None, "bytes_done": 0}
+_dl_lock = threading.Lock()
+
+
+def _do_download(save_to: str):
+    global _dl
+    with _dl_lock:
+        _dl.update(status="downloading", path=save_to, error=None, bytes_done=0)
+    try:
+        from huggingface_hub import hf_hub_download
+        save_dir = os.path.dirname(os.path.abspath(save_to))
+        os.makedirs(save_dir, exist_ok=True)
+        result_path = hf_hub_download(
+            repo_id   = HF_REPO,
+            filename  = HF_FILE,
+            local_dir = save_dir,
+        )
+        # hf_hub_download saves to local_dir/<filename>
+        canonical = os.path.join(save_dir, HF_FILE)
+        if result_path != save_to and os.path.exists(result_path):
+            os.replace(result_path, save_to)
+        with _dl_lock:
+            _dl.update(status="done", path=save_to, bytes_done=os.path.getsize(save_to))
+        print(f"[INFO] Model downloaded → {save_to}", flush=True)
+    except Exception as e:
+        with _dl_lock:
+            _dl.update(status="error", error=str(e))
+        print(f"[ERROR] Download failed: {e}", flush=True)
+
+
+@app.route("/gemma/download", methods=["GET", "POST"])
+def download_model():
+    """Start (or check) background download of the Gemma 4 model from HuggingFace."""
+    save_to = request.args.get("path", "").strip() or _DEFAULT_PATH
+
+    # Status check only (no action)
+    if request.args.get("status"):
+        with _dl_lock:
+            info = dict(_dl)
+        info["model_path"] = save_to
+        return jsonify(info)
+
+    # Already on disk
+    if os.path.exists(save_to):
+        size_mb = os.path.getsize(save_to) // (1024 * 1024)
+        return jsonify({
+            "status":     "already_exists",
+            "model_path": save_to,
+            "size_mb":    size_mb,
+            "next_step":  f"Set env var GEMMA_MODEL_PATH={save_to} and restart the server.",
+        })
+
+    with _dl_lock:
+        current = _dl["status"]
+
+    if current == "downloading":
+        return jsonify({
+            "status":  "downloading",
+            "message": "Download already in progress. Poll GET /gemma/download?status=1",
+        })
+
+    # Kick off background download
+    threading.Thread(target=_do_download, args=(save_to,), daemon=True).start()
+    return jsonify({
+        "status":     "started",
+        "model":      f"{HF_REPO}/{HF_FILE}",
+        "saving_to":  save_to,
+        "size":       "~2.5 GB — will take a few minutes",
+        "poll":       "GET /gemma/download?status=1",
+        "next_step":  f"When done, set GEMMA_MODEL_PATH={save_to} and restart the server.",
+    })
+
+
+# ─── 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)
+    print(f"[INFO] Model: {MODEL_PATH or '(not set — GET /gemma/download to fetch)'}", flush=True)
+    app.run(host="0.0.0.0", port=port, debug=False)