Update app.py

app.py

@@ -6,7 +6,7 @@ import sys
 print(f"[BOOT] Python {sys.version}", flush=True)
 
 import base64, os, re, json, subprocess
-from typing import Generator, Optional
+from typing import Generator
 from collections.abc import Iterator
 from pathlib import Path
 from threading import Thread
@@ -29,23 +29,12 @@ if not _installed:
     print("[BOOT] Installing transformers from PyPI...", flush=True)
     subprocess.check_call([sys.executable, "-m", "pip", "install", "transformers>=4.49"])
 
-import urllib3
-urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
-
 try:
     import gradio as gr
     print(f"[BOOT] gradio {gr.__version__}", flush=True)
 except ImportError as e:
     print(f"[BOOT] FATAL: {e}", flush=True); sys.exit(1)
 
-try:
-    import httpx, uvicorn, requests
-    from fastapi import FastAPI, Request
-    from fastapi.responses import HTMLResponse, RedirectResponse, JSONResponse
-    print("[BOOT] All imports OK", flush=True)
-except ImportError as e:
-    print(f"[BOOT] FATAL: {e}", flush=True); sys.exit(1)
-
 import torch
 import spaces
 from transformers import AutoModelForMultimodalLM, AutoProcessor, BatchFeature
@@ -118,9 +107,9 @@ def _load_model(model_name: str):
         import gc; gc.collect()
         print(f"[MODEL] Unloaded previous model", flush=True)
 
-    _processor = AutoProcessor.from_pretrained(model_id, use_fast=False)
+    _processor = AutoProcessor.from_pretrained(model_id, backend="pil")
     _model = AutoModelForMultimodalLM.from_pretrained(
-        model_id, device_map="auto", torch_dtype=torch.bfloat16,
+        model_id, device_map="auto", dtype=torch.bfloat16,
     )
 
     # Build strip tokens list (keep thinking delimiters)
@@ -352,202 +341,55 @@ def generate_reply(
 
 
 # ══════════════════════════════════════════════════════════════════════════════
-# 6.  GRADIO BLOCKS (hidden — serves API for frontend)
+# 6.  GRADIO CHAT INTERFACE — ZeroGPU compatible (must use demo.launch())
 # ══════════════════════════════════════════════════════════════════════════════
-with gr.Blocks(title="Gemma 4 Playground") as gradio_demo:
-    thinking_toggle = gr.Radio(
-        choices=["⚡ Fast Mode  (direct answer)",
-                 "🧠 Thinking Mode  (chain-of-thought reasoning)"],
-        value="⚡ Fast Mode  (direct answer)",
-        visible=False,
-    )
-    image_input    = gr.Textbox(value="", visible=False)
-    system_prompt  = gr.Textbox(value=PRESETS["general"], visible=False)
-    max_new_tokens = gr.Slider(minimum=64, maximum=8192, value=4096, visible=False)
-    temperature    = gr.Slider(minimum=0.0, maximum=1.5, value=0.6, visible=False)
-    top_p          = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, visible=False)
-    model_selector = gr.Dropdown(
-        choices=list(MODELS.keys()),
-        value=DEFAULT_MODEL,
-        visible=False,
-    )
-
-    gr.ChatInterface(
-        fn=generate_reply,
-        api_name="chat",
-        additional_inputs=[
-            thinking_toggle, image_input,
-            system_prompt, max_new_tokens, temperature, top_p,
-            model_selector,
+demo = gr.ChatInterface(
+    fn=generate_reply,
+    chatbot=gr.Chatbot(
+        scale=1,
+        latex_delimiters=[
+            {"left": "$$", "right": "$$", "display": True},
+            {"left": "$", "right": "$", "display": False},
+            {"left": "\\(", "right": "\\)", "display": False},
+            {"left": "\\[", "right": "\\]", "display": True},
         ],
-    )
-
-
-# ══════════════════════════════════════════════════════════════════════════════
-# 7.  FASTAPI — index.html + HF OAuth + API endpoints
-# ══════════════════════════════════════════════════════════════════════════════
-import pathlib, secrets
-
-fapp    = FastAPI()
-SESSIONS: dict[str, dict] = {}
-HTML    = pathlib.Path(__file__).parent / "index.html"
-
-CLIENT_ID     = os.getenv("OAUTH_CLIENT_ID", "")
-CLIENT_SECRET = os.getenv("OAUTH_CLIENT_SECRET", "")
-SPACE_HOST    = os.getenv("SPACE_HOST", "localhost:7860")
-REDIRECT_URI  = f"https://{SPACE_HOST}/login/callback"
-
-print(f"[OAuth] CLIENT_ID set: {bool(CLIENT_ID)}")
-print(f"[OAuth] SPACE_HOST: {SPACE_HOST}")
-HF_AUTH_URL   = "https://huggingface.co/oauth/authorize"
-HF_TOKEN_URL  = "https://huggingface.co/oauth/token"
-HF_USER_URL   = "https://huggingface.co/oauth/userinfo"
-SCOPES        = os.getenv("OAUTH_SCOPES", "openid profile")
-
-from urllib.parse import urlencode
-
-def _sid(req: Request) -> Optional[str]:
-    return req.cookies.get("mc_session")
-
-def _user(req: Request) -> Optional[dict]:
-    sid = _sid(req)
-    return SESSIONS.get(sid) if sid else None
-
-@fapp.get("/")
-async def root(request: Request):
-    html = HTML.read_text(encoding="utf-8") if HTML.exists() else "<h2>index.html missing</h2>"
-    return HTMLResponse(html)
-
-@fapp.get("/oauth/user")
-async def oauth_user(request: Request):
-    u = _user(request)
-    return JSONResponse(u) if u else JSONResponse({"logged_in": False}, status_code=401)
-
-@fapp.get("/oauth/login")
-async def oauth_login(request: Request):
-    if not CLIENT_ID:
-        return RedirectResponse("/?oauth_error=not_configured")
-    state = secrets.token_urlsafe(16)
-    params = {"response_type":"code","client_id":CLIENT_ID,"redirect_uri":REDIRECT_URI,"scope":SCOPES,"state":state}
-    return RedirectResponse(f"{HF_AUTH_URL}?{urlencode(params)}", status_code=302)
-
-@fapp.get("/login/callback")
-async def oauth_callback(code: str = "", error: str = "", state: str = ""):
-    if error or not code:
-        return RedirectResponse("/?auth_error=1")
-    basic = base64.b64encode(f"{CLIENT_ID}:{CLIENT_SECRET}".encode()).decode()
-    async with httpx.AsyncClient() as client:
-        tok = await client.post(HF_TOKEN_URL, data={"grant_type":"authorization_code","code":code,"redirect_uri":REDIRECT_URI},
-                                headers={"Accept":"application/json","Authorization":f"Basic {basic}"})
-        if tok.status_code != 200:
-            return RedirectResponse("/?auth_error=1")
-        access_token = tok.json().get("access_token", "")
-        if not access_token:
-            return RedirectResponse("/?auth_error=1")
-        uinfo = await client.get(HF_USER_URL, headers={"Authorization":f"Bearer {access_token}"})
-        if uinfo.status_code != 200:
-            return RedirectResponse("/?auth_error=1")
-        user = uinfo.json()
-
-    sid = secrets.token_urlsafe(32)
-    SESSIONS[sid] = {
-        "logged_in": True,
-        "username": user.get("preferred_username", user.get("name", "User")),
-        "name": user.get("name", ""),
-        "avatar": user.get("picture", ""),
-        "profile": f"https://huggingface.co/{user.get('preferred_username', '')}",
-    }
-    resp = RedirectResponse("/")
-    resp.set_cookie("mc_session", sid, httponly=True, samesite="lax", secure=True, max_age=60*60*24*7)
-    return resp
-
-@fapp.get("/oauth/logout")
-async def oauth_logout(request: Request):
-    sid = _sid(request)
-    if sid and sid in SESSIONS: del SESSIONS[sid]
-    resp = RedirectResponse("/")
-    resp.delete_cookie("mc_session")
-    return resp
-
-# ── Model Info API (for frontend model selector) ────────────────────────
-@fapp.get("/api/models")
-async def api_models():
-    return JSONResponse({
-        "models": {k: v for k, v in MODELS.items()},
-        "current": _loaded_model_name,
-        "default": DEFAULT_MODEL,
-    })
-
-@fapp.get("/health")
-async def health():
-    return {
-        "status": "ok",
-        "model_loaded": _loaded_model_name,
-        "model_arch": MODELS.get(_loaded_model_name, {}).get("arch", "unknown"),
-        "gpu": torch.cuda.is_available(),
-    }
-
-# ── Web Search API (Brave) ──────────────────────────────────────────────
-BRAVE_API_KEY = os.getenv("BRAVE_API_KEY", "")
-
-@fapp.post("/api/search")
-async def api_search(request: Request):
-    body = await request.json()
-    query = body.get("query", "").strip()
-    if not query:
-        return JSONResponse({"error": "empty query"}, status_code=400)
-    key = BRAVE_API_KEY
-    if not key:
-        return JSONResponse({"error": "BRAVE_API_KEY not set"}, status_code=500)
-    try:
-        r = requests.get(
-            "https://api.search.brave.com/res/v1/web/search",
-            headers={"X-Subscription-Token": key, "Accept": "application/json"},
-            params={"q": query, "count": 5}, timeout=10,
-        )
-        r.raise_for_status()
-        results = r.json().get("web", {}).get("results", [])
-        items = []
-        for item in results[:5]:
-            items.append({
-                "title": item.get("title", ""),
-                "desc": item.get("description", ""),
-                "url": item.get("url", ""),
-            })
-        return JSONResponse({"results": items})
-    except Exception as e:
-        return JSONResponse({"error": str(e)}, status_code=500)
-
-# ── PDF Text Extraction ─────────────────────────────────────────────────
-@fapp.post("/api/extract-pdf")
-async def api_extract_pdf(request: Request):
-    try:
-        body = await request.json()
-        b64 = body.get("data", "")
-        if "," in b64:
-            b64 = b64.split(",", 1)[1]
-        import io
-        pdf_bytes = base64.b64decode(b64)
-        text = ""
-        try:
-            import fitz  # PyMuPDF
-            doc = fitz.open(stream=pdf_bytes, filetype="pdf")
-            for page in doc:
-                text += page.get_text() + "\n"
-        except ImportError:
-            content = pdf_bytes.decode("utf-8", errors="ignore")
-            text = re.sub(r'[^\x20-\x7E\n\r\uAC00-\uD7A3\u3040-\u309F\u30A0-\u30FF]', '', content)
-        text = text.strip()[:8000]
-        return JSONResponse({"text": text, "chars": len(text)})
-    except Exception as e:
-        return JSONResponse({"error": str(e)}, status_code=500)
-
+    ),
+    textbox=gr.Textbox(placeholder="Message Gemma 4…", lines=1, scale=7),
+    additional_inputs=[
+        gr.Radio(
+            choices=["⚡ Fast Mode  (direct answer)",
+                     "🧠 Thinking Mode  (chain-of-thought reasoning)"],
+            value="⚡ Fast Mode  (direct answer)",
+            label="Mode",
+        ),
+        gr.Textbox(value="", label="Image (base64 or URL)", visible=False),
+        gr.Textbox(value=PRESETS["general"], label="System Prompt", lines=2),
+        gr.Slider(minimum=64, maximum=8192, value=4096, step=64, label="Max Tokens"),
+        gr.Slider(minimum=0.0, maximum=1.5, value=0.6, step=0.05, label="Temperature"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-P"),
+        gr.Dropdown(
+            choices=list(MODELS.keys()),
+            value=DEFAULT_MODEL,
+            label="Model",
+            info="26B-A4B: MoE 3.8B active (fast) | 31B: Dense (best quality)",
+        ),
+    ],
+    additional_inputs_accordion=gr.Accordion("⚙️ Settings", open=False),
+    title="💎 Gemma 4 Playground",
+    description="Google DeepMind Gemma 4 — Dense 31B or MoE 26B-A4B · Vision · Thinking · Apache 2.0",
+    examples=[
+        ["Explain how Gemma 4 achieves frontier-level performance with Mixture-of-Experts architecture."],
+        ["Write a Python async web scraper with retry logic and rate limiting."],
+        ["한국의 K-pop이 세계적으로 성공한 이유를 문화적, 경제적 관점에서 분석해주세요."],
+        ["Solve: What is the sum of all prime numbers less than 100?"],
+    ],
+    run_examples_on_click=False,
+    cache_examples=False,
+)
 
 # ══════════════════════════════════════════════════════════════════════════════
-# 8.  MOUNT & LAUNCH
+# 7.  LAUNCH  — must use demo.launch() for ZeroGPU @spaces.GPU registration
 # ══════════════════════════════════════════════════════════════════════════════
-app = gr.mount_gradio_app(fapp, gradio_demo, path="/gradio")
-
 if __name__ == "__main__":
     print(f"[BOOT] Gemma 4 Playground · Default: {DEFAULT_MODEL}", flush=True)
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860)