Update server.py

server.py

@@ -3,6 +3,8 @@ import io
 import base64
 import ctypes
 import threading
+import json
+import time
 from flask import Flask, request, jsonify, Response
 from flask_cors import CORS
 
@@ -13,7 +15,6 @@ _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:
@@ -68,153 +69,78 @@ def load_model():
         model_status = "error"
 
 
-# ─── Image analysis via Pillow (no model) ──────────────────────────────────────
+# ─── 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()
-
-    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)
-
+    return dict(w=w, h=h, fmt=fmt)
 
 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)
+        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."
 
-# ─── 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
+# ─── Real model inference (litert_lm Generator) ────────────────────────────────
 
-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:
+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:
-        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)
+        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]:
+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
+        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")
-        return ask, (f.read() if f else None)
+        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:
@@ -222,128 +148,89 @@ def extract_request() -> tuple[str, bytes | None]:
             image_bytes = base64.b64decode(raw)
         except Exception:
             pass
-    return ask, image_bytes
+            
+    return ask, image_bytes, stream
 
 
 # ─── Routes ────────────────────────────────────────────────────────────────────
 
-@app.route("/favicon.ico")
-def favicon():
-    return "", 204
-
-
 @app.route("/")
 def index():
     return jsonify({
-        "service":      "Gemma 4 API",
+        "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",
-        },
+        "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")
-@app.route("/api/health")
 def health():
-    info = {"status": "ok", "model_status": model_status}
-    if MODEL_PATH:
-        info["model_path"] = MODEL_PATH
-    return jsonify(info)
+    return jsonify({"status": "ok", "model_status": model_status})
 
 
 @app.route("/gemma", methods=["GET", "POST"])
 def gemma():
-    ask, image_bytes = extract_request()
+    ask, image_bytes, stream = 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.",
-    })
+    # 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):
+                    # Standard SSE format for frontend parsing
+                    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:
+        # Wait for all chunks and join them into one string
+        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 ─────────────────────────────────────────────────────────────────────
@@ -352,5 +239,4 @@ 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)
+    app.run(host="0.0.0.0", port=port, debug=False)