Update app.py

app.py

@@ -1,43 +1,50 @@
 import os
 import gc
 import asyncio
-import json
 from typing import Dict, List, Optional
-from contextlib import asynccontextmanager
 
 import torch
-from fastapi import FastAPI, HTTPException, Header, Depends, Request
-from fastapi.responses import StreamingResponse, JSONResponse, HTMLResponse
+import gradio as gr
+from contextlib import asynccontextmanager
+
+from fastapi import FastAPI, HTTPException, Header, Depends
 from fastapi.middleware.cors import CORSMiddleware
+from starlette.responses import Response  # noqa: F401
 from pydantic import BaseModel, Field
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
-from threading import Thread
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 
-# ═══════════════════════════════════════════════════
-# MODEL REGISTRY — EDIT THIS
-# ═══════════════════════════════════════════════════
+# --------------------------------------------------------------------
+# EDIT THIS DICTIONARY
+# key = name shown in UI/API
+# value = Hugging Face repo ID or local path
+# --------------------------------------------------------------------
 MODELS: Dict[str, str] = {
     "Qwen2.5-0.5B-Instruct": "Qwen/Qwen2.5-0.5B-Instruct",
     "TinyLlama-1.1B-Chat": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+    # "My Cool Model": "YourUsername/your-model-repo",
+    # "Local Model": "./models/my-local-model",
 }
 
 if not MODELS:
     raise RuntimeError("bro, put at least one model in MODELS.")
 
 DEFAULT_MODEL = next(iter(MODELS))
+
 PRIORITY_API_KEY = os.getenv("PRIORITY_API_KEY", "").strip()
+ALLOW_OPEN_PRIORITY = os.getenv("ALLOW_OPEN_PRIORITY", "0") == "1"
 MAX_NEW_TOKENS_LIMIT = int(os.getenv("MAX_NEW_TOKENS_LIMIT", "1024"))
+MAX_SLIDER = min(8192, max(256, MAX_NEW_TOKENS_LIMIT))
+DEFAULT_MAX_TOKENS = min(512, MAX_SLIDER)
+
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 
-# ═══════════════════════════════════════════════════
-# PYDANTIC SCHEMAS
-# ═══════════════════════════════════════════════════
 class ChatMessage(BaseModel):
     role: str
     content: str
 
+
 class ChatRequest(BaseModel):
     model: str = Field(default=DEFAULT_MODEL)
     messages: List[ChatMessage] = Field(min_length=1)
@@ -46,9 +53,6 @@ class ChatRequest(BaseModel):
     top_p: float = Field(default=0.95, ge=0.0, le=1.0)
 
 
-# ═══════════════════════════════════════════════════
-# MODEL MANAGER
-# ═══════════════════════════════════════════════════
 class ModelManager:
     def __init__(self):
         self.current_model = None
@@ -66,64 +70,152 @@ class ModelManager:
     def load(self, name: str):
         if name == self.current_model and self.model is not None:
             return
+
         if name not in MODELS:
-            raise ValueError(f"Model '{name}' not in MODELS.")
+            raise ValueError(f"Model '{name}' is not in MODELS.")
 
         model_id = MODELS[name]
         self._cleanup()
+
         token = os.getenv("HF_TOKEN") or None
+        trust_remote_code = os.getenv("TRUST_REMOTE_CODE", "0") == "1"
         use_4bit = os.getenv("USE_4BIT", "1") == "1" and torch.cuda.is_available()
 
-        kwargs = {"token": token, "low_cpu_mem_usage": True}
+        model_kwargs = {
+            "token": token,
+            "low_cpu_mem_usage": True,
+            "trust_remote_code": trust_remote_code,
+        }
+
         if use_4bit:
-            dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
-            kwargs.update(load_in_4bit=True, bnb_4bit_compute_dtype=dtype, bnb_4bit_quant_type="nf4")
+            compute_dtype = (
+                torch.bfloat16
+                if torch.cuda.is_available() and torch.cuda.is_bf16_supported()
+                else torch.float16
+            )
+            model_kwargs.update(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype=compute_dtype,
+                bnb_4bit_quant_type="nf4",
+            )
         else:
-            kwargs["torch_dtype"] = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16
+            model_kwargs["torch_dtype"] = (
+                torch.bfloat16
+                if torch.cuda.is_available() and torch.cuda.is_bf16_supported()
+                else torch.float16
+            )
+
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_id,
+            token=token,
+            trust_remote_code=trust_remote_code,
+        )
 
-        self.tokenizer = AutoTokenizer.from_pretrained(model_id, token=token)
         if self.tokenizer.pad_token is None:
             self.tokenizer.pad_token = self.tokenizer.eos_token
+
         if getattr(self.tokenizer, "pad_token_id", None) is None:
             self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+
+        if getattr(self.tokenizer, "eos_token_id", None) is None:
+            self.tokenizer.eos_token_id = self.tokenizer.pad_token_id or 0
+
         self.tokenizer.padding_side = "left"
 
-        self.model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", **kwargs)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            device_map="auto",
+            **model_kwargs,
+        )
         self.model.eval()
         self.current_model = name
 
-    def build_prompt(self, messages: List[ChatMessage]) -> str:
-        raw = [m.model_dump() for m in messages]
+    def _build_prompt(self, messages: List[ChatMessage]) -> str:
+        raw_messages = [m.model_dump() for m in messages]
+
         try:
-            return self.tokenizer.apply_chat_template(raw, tokenize=False, add_generation_prompt=True)
+            return self.tokenizer.apply_chat_template(
+                raw_messages,
+                tokenize=False,
+                add_generation_prompt=True,
+            )
         except Exception:
             lines = []
             for m in messages:
-                r = m.role.strip().lower()
-                if r == "system": lines.append(f"System: {m.content}")
-                elif r == "user": lines.append(f"User: {m.content}")
-                elif r == "assistant": lines.append(f"Assistant: {m.content}")
-                else: lines.append(f"{m.role}: {m.content}")
+                role = m.role.strip().lower()
+                if role == "system":
+                    lines.append(f"System: {m.content}")
+                elif role == "user":
+                    lines.append(f"User: {m.content}")
+                elif role == "assistant":
+                    lines.append(f"Assistant: {m.content}")
+                else:
+                    lines.append(f"{m.role}: {m.content}")
             lines.append("Assistant:")
             return "\n".join(lines)
 
-    def get_max_input_length(self) -> int:
-        ml = getattr(self.tokenizer, "model_max_length", None)
+    def generate(self, request: ChatRequest) -> str:
         try:
-            ml = int(ml) if ml is not None else 4096
-        except Exception:
-            ml = 4096
-        if ml <= 0 or ml > 1_000_000:
-            ml = 4096
-        return min(4096, ml)
+            self.load(request.model)
 
+            prompt = self._build_prompt(request.messages)
 
-model_manager = ModelManager()
+            model_max = getattr(self.tokenizer, "model_max_length", None)
+            try:
+                model_max = int(model_max) if model_max is not None else 4096
+            except Exception:
+                model_max = 4096
+
+            if model_max <= 0 or model_max > 1_000_000:
+                model_max = 4096
+
+            max_input_length = min(4096, model_max)
+
+            inputs = self.tokenizer(
+                prompt,
+                return_tensors="pt",
+                truncation=True,
+                max_length=max_input_length,
+            )
+            inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
+
+            input_len = inputs["input_ids"].shape[1]
+
+            max_new_tokens = max(
+                1,
+                min(
+                    request.max_new_tokens,
+                    MAX_NEW_TOKENS_LIMIT,
+                    max_input_length - input_len,
+                ),
+            )
+
+            gen_kwargs = {
+                "max_new_tokens": max_new_tokens,
+                "pad_token_id": self.tokenizer.pad_token_id,
+                "eos_token_id": self.tokenizer.eos_token_id,
+            }
+
+            if request.temperature > 0:
+                gen_kwargs.update(
+                    do_sample=True,
+                    temperature=request.temperature,
+                    top_p=request.top_p if request.top_p > 0 else 1.0,
+                )
+            else:
+                gen_kwargs.update(do_sample=False)
+
+            with torch.inference_mode():
+                output_ids = self.model.generate(**inputs, **gen_kwargs)
+
+            new_tokens = output_ids[0, inputs["input_ids"].shape[1]:]
+            return self.tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
+
+        except Exception:
+            self._cleanup()
+            raise
 
 
-# ═══════════════════════════════════════════════════
-# PRIORITY QUEUE BROKER
-# ═══════════════════════════════════════════════════
 class PriorityBroker:
     def __init__(self):
         self.queue = asyncio.PriorityQueue()
@@ -135,326 +227,287 @@ class PriorityBroker:
 
     async def stop(self):
         self.counter += 1
-        await self.queue.put((999999, self.counter, None, None, None))
+        await self.queue.put((999999, self.counter, None, None))
         if self.task:
             await self.task
 
-    async def enqueue(self, priority: int, request: ChatRequest, stream: bool = False):
+    async def enqueue(self, priority: int, request: ChatRequest) -> str:
         loop = asyncio.get_running_loop()
         future = loop.create_future()
         self.counter += 1
-        await self.queue.put((priority, self.counter, future, request, stream))
+        await self.queue.put((priority, self.counter, future, request))
         return await future
 
     async def _worker(self):
         while True:
             item = await self.queue.get()
-            priority, counter, future, request, stream = item
+            priority, counter, future, request = item
 
             if future is None:
                 self.queue.task_done()
                 break
+
             if future.cancelled():
                 self.queue.task_done()
                 continue
 
             try:
-                if stream:
-                    # For streaming, we set result to a generator function
-                    gen_fn = lambda: self._generate_stream(request)
-                    future.set_result(gen_fn)
-                else:
-                    result = await asyncio.to_thread(self._generate_full, request)
-                    future.set_result(result)
+                result = await asyncio.to_thread(model_manager.generate, request)
+                future.set_result(result)
             except Exception as exc:
                 future.set_exception(exc)
             finally:
                 self.queue.task_done()
 
-    def _generate_full(self, request: ChatRequest) -> str:
-        try:
-            model_manager.load(request.model)
-            prompt = model_manager.build_prompt(request.messages)
-            max_input = model_manager.get_max_input_length()
-            inputs = model_manager.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=max_input)
-            inputs = {k: v.to(model_manager.model.device) for k, v in inputs.items()}
-
-            input_len = inputs["input_ids"].shape[1]
-            max_new = max(1, min(request.max_new_tokens, MAX_NEW_TOKENS_LIMIT, max_input - input_len))
-
-            gen_kwargs = {"max_new_tokens": max_new, "pad_token_id": model_manager.tokenizer.pad_token_id, "eos_token_id": model_manager.tokenizer.eos_token_id}
-            if request.temperature > 0:
-                gen_kwargs.update(do_sample=True, temperature=request.temperature, top_p=request.top_p if request.top_p > 0 else 1.0)
-            else:
-                gen_kwargs.update(do_sample=False)
-
-            with torch.inference_mode():
-                out = model_manager.model.generate(**inputs, **gen_kwargs)
-            new = out[0, inputs["input_ids"].shape[1]:]
-            return model_manager.tokenizer.decode(new, skip_special_tokens=True).strip()
-        except Exception:
-            model_manager._cleanup()
-            raise
-
-    def _generate_stream(self, request: ChatRequest):
-        """Returns a generator that yields tokens. Called in a thread."""
-        try:
-            model_manager.load(request.model)
-            prompt = model_manager.build_prompt(request.messages)
-            max_input = model_manager.get_max_input_length()
-            inputs = model_manager.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=max_input)
-            inputs = {k: v.to(model_manager.model.device) for k, v in inputs.items()}
-
-            input_len = inputs["input_ids"].shape[1]
-            max_new = max(1, min(request.max_new_tokens, MAX_NEW_TOKENS_LIMIT, max_input - input_len))
-
-            streamer = TextIteratorStreamer(model_manager.tokenizer, skip_prompt=True, skip_special_tokens=True)
-
-            gen_kwargs = {"max_new_tokens": max_new, "pad_token_id": model_manager.tokenizer.pad_token_id, "eos_token_id": model_manager.tokenizer.eos_token_id, "streamer": streamer}
-            if request.temperature > 0:
-                gen_kwargs.update(do_sample=True, temperature=request.temperature, top_p=request.top_p if request.top_p > 0 else 1.0)
-            else:
-                gen_kwargs.update(do_sample=False)
-
-            def run_gen():
-                with torch.inference_mode():
-                    model_manager.model.generate(**inputs, **gen_kwargs)
-
-            t = Thread(target=run_gen)
-            t.start()
-
-            for token_text in streamer:
-                yield token_text
-
-            t.join(timeout=5)
-        except Exception:
-            model_manager._cleanup()
-            raise
-
 
+model_manager = ModelManager()
 broker = PriorityBroker()
 
 
-# ═══════════════════════════════════════════════════
-# FASTAPI APP
-# ══════════════════���════════════════════════════════
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     await broker.start()
     yield
     await broker.stop()
 
+
 app = FastAPI(title="Priority Model Chat", lifespan=lifespan)
-app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"])
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
 
 def require_priority_key(
     x_api_key: Optional[str] = Header(default=None, alias="x-api-key"),
     authorization: Optional[str] = Header(default=None),
 ):
+    if ALLOW_OPEN_PRIORITY:
+        return
+
     provided = None
+
     if x_api_key:
         provided = x_api_key.strip()
     elif authorization and authorization.lower().startswith("bearer "):
         provided = authorization[7:].strip()
+
     if not PRIORITY_API_KEY:
-        raise HTTPException(503, "Set PRIORITY_API_KEY in Space secrets first.")
+        raise HTTPException(
+            status_code=503,
+            detail=(
+                "Set PRIORITY_API_KEY in Space secrets, "
+                "or set ALLOW_OPEN_PRIORITY=1 for testing."
+            ),
+        )
+
     if provided != PRIORITY_API_KEY:
-        raise HTTPException(401, "Invalid or missing priority API key.")
+        raise HTTPException(
+            status_code=401,
+            detail="Invalid or missing priority API key.",
+        )
 
 
 @app.get("/health")
 async def health():
-    return {"ok": True, "device": DEVICE, "models": list(MODELS.keys()), "queue_size": broker.queue.qsize()}
+    return {
+        "ok": True,
+        "device": DEVICE,
+        "models": list(MODELS.keys()),
+        "queue_size": broker.queue.qsize(),
+        "priority_open": ALLOW_OPEN_PRIORITY,
+    }
+
 
 @app.get("/api/models")
 async def models():
     return {"models": MODELS}
 
+
 @app.post("/api/chat")
 async def chat(req: ChatRequest):
-    if not req.model: req.model = DEFAULT_MODEL
+    if not req.model:
+        req.model = DEFAULT_MODEL
+
     if req.model not in MODELS:
-        raise HTTPException(404, f"Model '{req.model}' not found. Available: {list(MODELS.keys())}")
+        raise HTTPException(
+            status_code=404,
+            detail=f"Model '{req.model}' not found. Available: {list(MODELS.keys())}",
+        )
+
     try:
-        answer = await broker.enqueue(priority=1, request=req, stream=False)
-        return {"model": req.model, "response": answer}
+        answer = await broker.enqueue(priority=1, request=req)
+        return {
+            "model": req.model,
+            "response": answer,
+        }
     except Exception as exc:
-        raise HTTPException(500, f"{type(exc).__name__}: {exc}")
+        raise HTTPException(status_code=500, detail=f"{type(exc).__name__}: {exc}")
+
 
 @app.post("/api/priority")
-async def priority_chat(req: ChatRequest, _: None = Depends(require_priority_key)):
-    if not req.model: req.model = DEFAULT_MODEL
+async def priority_chat(
+    req: ChatRequest,
+    _: None = Depends(require_priority_key),
+):
+    if not req.model:
+        req.model = DEFAULT_MODEL
+
     if req.model not in MODELS:
-        raise HTTPException(404, f"Model '{req.model}' not found. Available: {list(MODELS.keys())}")
+        raise HTTPException(
+            status_code=404,
+            detail=f"Model '{req.model}' not found. Available: {list(MODELS.keys())}",
+        )
+
     try:
-        answer = await broker.enqueue(priority=0, request=req, stream=False)
-        return {"model": req.model, "response": answer, "priority": True}
+        answer = await broker.enqueue(priority=0, request=req)
+        return {
+            "model": req.model,
+            "response": answer,
+            "priority": True,
+        }
     except Exception as exc:
-        raise HTTPException(500, f"{type(exc).__name__}: {exc}")
+        raise HTTPException(status_code=500, detail=f"{type(exc).__name__}: {exc}")
 
 
-# ═══════════════════════════════════════════════════
-# STREAMING ENDPOINTS (SSE)
-# ═══════════════════════════════════════════════════
-async def sse_generator(gen_fn):
-    """Wraps a sync token generator into SSE format."""
-    try:
-        # gen_fn is a lambda that returns a generator when called
-        token_iter = await asyncio.to_thread(gen_fn)
-        for token_text in token_iter:
-            data = json.dumps({"token": token_text})
-            yield f"data: {data}\n\n"
-        yield "data: [DONE]\n\n"
-    except Exception as exc:
-        err = json.dumps({"error": f"{type(exc).__name__}: {exc}"})
-        yield f"data: {err}\n\n"
-        yield "data: [DONE]\n\n"
+CSS = """
+.gradio-container {
+    max-width: 950px;
+    margin: auto;
+}
+footer {
+    display: none !important;
+}
+"""
 
 
-@app.post("/api/chat/stream")
-async def chat_stream(req: ChatRequest):
-    if not req.model: req.model = DEFAULT_MODEL
-    if req.model not in MODELS:
-        raise HTTPException(404, f"Model '{req.model}' not found.")
-    try:
-        gen_fn = await broker.enqueue(priority=1, request=req, stream=True)
-        return StreamingResponse(sse_generator(gen_fn), media_type="text/event-stream",
-                                 headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})
-    except Exception as exc:
-        raise HTTPException(500, f"{type(exc).__name__}: {exc}")
+async def ui_respond(
+    model_name,
+    user_message,
+    history,
+    system,
+    temperature,
+    top_p,
+    max_new_tokens,
+):
+    if not user_message or not str(user_message).strip():
+        yield history, ""
+        return
 
+    if model_name not in MODELS:
+        model_name = DEFAULT_MODEL
 
-@app.post("/api/priority/stream")
-async def priority_stream(req: ChatRequest, _: None = Depends(require_priority_key)):
-    if not req.model: req.model = DEFAULT_MODEL
-    if req.model not in MODELS:
-        raise HTTPException(404, f"Model '{req.model}' not found.")
-    try:
-        gen_fn = await broker.enqueue(priority=0, request=req, stream=True)
-        return StreamingResponse(sse_generator(gen_fn), media_type="text/event-stream",
-                                 headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})
-    except Exception as exc:
-        raise HTTPException(500, f"{type(exc).__name__}: {exc}")
-
-
-# ═══════════════════════════════════════════════════
-# BUILT-IN CHAT UI (served at /)
-# ═══════════════════════════════════════════════════
-CHAT_UI_HTML = """<!DOCTYPE html>
-<html lang="en">
-<head>
-<meta charset="UTF-8"><meta name="viewport" content="width=device-width,initial-scale=1">
-<title>⚡ Priority Model Chat</title>
-<style>
-*{margin:0;padding:0;box-sizing:border-box}
-body{background:#0a0a1e;color:#e0e0e0;font-family:'Segoe UI',system-ui,sans-serif;height:100vh;display:flex;flex-direction:column}
-.header{text-align:center;padding:1rem;border-bottom:1px solid rgba(0,229,255,0.15)}
-.header h1{font-size:1.5rem;background:linear-gradient(90deg,#00e5ff,#7c4dff);-webkit-background-clip:text;-webkit-text-fill-color:transparent}
-.header p{color:#888;font-size:0.8rem;margin-top:0.2rem}
-.chat-area{flex:1;overflow-y:auto;padding:1rem;display:flex;flex-direction:column;gap:0.6rem;scroll-behavior:smooth}
-.msg{max-width:80%;padding:0.7rem 1rem;border-radius:12px;line-height:1.5;word-wrap:break-word;white-space:pre-wrap;font-size:0.95rem}
-.msg.user{align-self:flex-end;background:rgba(0,229,255,0.12);border:1px solid rgba(0,229,255,0.25);color:#d0f0ff}
-.msg.assistant{align-self:flex-start;background:rgba(255,255,255,0.06);border:1px solid rgba(255,255,255,0.08);color:#ddd}
-.msg.sys{align-self:center;font-size:0.75rem;color:#666;font-style:italic}
-.msg.err{align-self:center;background:rgba(255,50,50,0.12);border:1px solid rgba(255,50,50,0.25);color:#ff6b6b}
-.cursor{display:inline-block;width:2px;height:1em;background:#00e5ff;animation:blink 0.8s infinite;vertical-align:text-bottom;margin-left:2px}
-@keyframes blink{0%,100%{opacity:1}50%{opacity:0}}
-.controls{padding:0.5rem 1rem;display:flex;gap:0.8rem;flex-wrap:wrap;align-items:center;border-top:1px solid rgba(255,255,255,0.06);background:rgba(10,10,30,0.5)}
-.controls label{color:#aaa;font-size:0.78rem;display:flex;align-items:center;gap:0.3rem}
-.controls select,.controls input[type=number]{background:#12122a;border:1px solid rgba(255,255,255,0.12);border-radius:6px;padding:0.25rem 0.4rem;color:#fff;font-size:0.85rem}
-.input-row{display:flex;gap:0.5rem;padding:0.75rem 1rem;border-top:1px solid rgba(0,229,255,0.1);background:#0d0d24}
-.input-row textarea{flex:1;background:#12122a;border:1px solid rgba(0,229,255,0.2);border-radius:8px;padding:0.6rem;color:#fff;font-size:0.95rem;resize:none;outline:none;min-height:44px;max-height:120px;font-family:inherit}
-.input-row textarea:focus{border-color:#00e5ff;box-shadow:0 0 8px rgba(0,229,255,0.15)}
-.input-row button{background:linear-gradient(135deg,#00e5ff,#7c4dff);border:none;border-radius:8px;padding:0 1.5rem;color:#000;font-weight:700;cursor:pointer;font-size:0.9rem;transition:all 0.2s}
-.input-row button:hover{transform:translateY(-1px);box-shadow:0 4px 12px rgba(0,229,255,0.25)}
-.input-row button:disabled{opacity:0.4;cursor:not-allowed;transform:none}
-.clear-btn{background:transparent;border:1px solid rgba(255,100,100,0.25);color:#ff6b6b;border-radius:6px;padding:0.25rem 0.7rem;cursor:pointer;font-size:0.78rem;margin-left:auto;transition:all 0.2s}
-.clear-btn:hover{background:rgba(255,100,100,0.08)}
-</style>
-</head>
-<body>
-<div class="header"><h1>⚡ Priority Model Chat</h1><p>streaming • priority queue • no gradio</p></div>
-<div class="chat-area" id="chat"></div>
-<div class="controls">
-  <label>Model:<select id="mdl"></select></label>
-  <label>Temp:<input type="number" id="tmp" value="0.7" min="0" max="2" step="0.05" style="width:55px"></label>
-  <label>Tokens:<input type="number" id="tok" value="512" min="16" max="2048" step="16" style="width:65px"></label>
-  <button class="clear-btn" onclick="clearChat()">🧹 Clear</button>
-</div>
-<div class="input-row">
-  <textarea id="inp" placeholder="Type a message… (Enter to send, Shift+Enter for newline)" rows="1"></textarea>
-  <button id="btn" onclick="send()">Send ⚡</button>
-</div>
-<script>
-const chat=document.getElementById('chat'),inp=document.getElementById('inp'),btn=document.getElementById('btn');
-let hist=[];
-fetch('/api/models').then(r=>r.json()).then(d=>{
-  const s=document.getElementById('mdl');
-  Object.keys(d.models).forEach(k=>{const o=document.createElement('option');o.value=k;o.textContent=k;s.appendChild(o)});
-}).catch(()=>{});
-
-function addMsg(role,text,cls){const d=document.createElement('div');d.className='msg '+role+(cls?' '+cls:'');d.textContent=text;chat.appendChild(d);chat.scrollTop=chat.scrollHeight;return d}
-function clearChat(){hist=[];chat.innerHTML='<div class="msg sys">chat cleared.</div>'}
-
-inp.addEventListener('keydown',e=>{if(e.key==='Enter'&&!e.shiftKey){e.preventDefault();send()}});
-inp.addEventListener('input',()=>{inp.style.height='auto';inp.style.height=Math.min(inp.scrollHeight,120)+'px'});
-
-async function send(){
-  const text=inp.value.trim();if(!text)return;
-  addMsg('user',text);
-  hist.push({role:'user',content:text});
-  inp.value='';inp.style.height='auto';btn.disabled=true;
-
-  const assistantDiv=addMsg('assistant','');
-  const cursor=document.createElement('span');cursor.className='cursor';assistantDiv.appendChild(cursor);
-
-  try{
-    const resp=await fetch('/api/priority/stream',{
-      method:'POST',
-      headers:{'Content-Type':'application/json','x-api-key':document.getElementById('apiKey')?.value||''},
-      body:JSON.stringify({messages:hist,model:document.getElementById('mdl').value,
-        temperature:parseFloat(document.getElementById('tmp').value)||0.7,
-        max_new_tokens:parseInt(document.getElementById('tok').value)||512})
-    });
-
-    if(!resp.ok){
-      const err=await resp.text();
-      cursor.remove();assistantDiv.textContent='💀 '+err;assistantDiv.classList.add('err');
-      btn.disabled=false;return;
-    }
+    clean_history = []
+    for item in history or []:
+        if not isinstance(item, dict):
+            continue
 
-    const reader=resp.body.getReader();
-    const decoder=new TextDecoder();
-    let fullText='';let buffer='';
-
-    while(true){
-      const{done,value}=await reader.read();
-      if(done)break;
-      buffer+=decoder.decode(value,{stream:true});
-      const lines=buffer.split('\\n');
-      buffer=lines.pop()||'';
-      for(const line of lines){
-        if(!line.startsWith('data: '))continue;
-        const payload=line.slice(6).trim();
-        if(payload==='[DONE]'){cursor.remove();hist.push({role:'assistant',content:fullText});btn.disabled=false;return}
-        try{
-          const parsed=JSON.parse(payload);
-          if(parsed.error){cursor.remove();assistantDiv.textContent='💀 '+parsed.error;assistantDiv.classList.add('err');btn.disabled=false;return}
-          if(parsed.token!==undefined){fullText+=parsed.token;assistantDiv.textContent=fullText;assistantDiv.appendChild(cursor);chat.scrollTop=chat.scrollHeight}
-        }catch(e){}
-      }
-    }
-    cursor.remove();
-    if(fullText)hist.push({role:'assistant',content:fullText});
-  }catch(e){
-    cursor.remove();assistantDiv.textContent='💀 Network error: '+e.message;assistantDiv.classList.add('err');
-  }finally{btn.disabled=false;inp.focus()}
-}
-</script>
-</body>
-</html>"""
+        role = item.get("role")
+        content = str(item.get("content", ""))
+
+        if role not in {"user", "assistant", "system"}:
+            continue
+
+        if content == "⏳ cooking...":
+            continue
+
+        clean_history.append({"role": role, "content": content})
+
+    api_messages = []
+
+    if system and str(system).strip():
+        api_messages.append({"role": "system", "content": str(system).strip()})
 
-@app.get("/", response_class=HTMLResponse)
-async def serve_ui():
-    return CHAT_UI_HTML
+    api_messages.extend(clean_history)
+    api_messages.append({"role": "user", "content": str(user_message).strip()})
+
+    request = ChatRequest(
+        model=model_name,
+        messages=[ChatMessage(**m) for m in api_messages],
+        max_new_tokens=int(max_new_tokens or DEFAULT_MAX_TOKENS),
+        temperature=float(temperature or 0.7),
+        top_p=float(top_p or 0.95),
+    )
+
+    display_history = clean_history + [
+        {"role": "user", "content": str(user_message).strip()},
+        {"role": "assistant", "content": "⏳ cooking..."},
+    ]
+
+    yield display_history, ""
+
+    try:
+        answer = await broker.enqueue(priority=1, request=request)
+        display_history[-1]["content"] = answer
+    except Exception as exc:
+        display_history[-1]["content"] = f"💀 {type(exc).__name__}: {exc}"
+
+    yield display_history, ""
+
+
+with gr.Blocks(css=CSS) as demo:
+    gr.Markdown("## ⚡ Priority Model Chat")
+    gr.Markdown("chat UI uses the normal queue. priority API is separate.")
+
+    chatbot = gr.Chatbot(
+        type="messages",
+        height=520,
+    )
+
+    with gr.Row():
+        msg = gr.Textbox(
+            label="Message",
+            placeholder="type something...",
+            lines=2,
+            scale=5,
+        )
+        send = gr.Button("Send", variant="primary", scale=1)
+
+    with gr.Accordion("⚙️ model + generation settings", open=False):
+        model = gr.Dropdown(
+            choices=list(MODELS.keys()),
+            value=DEFAULT_MODEL,
+            label="Model",
+        )
+        system = gr.Textbox(
+            label="System prompt",
+            value="You are a helpful assistant.",
+            lines=2,
+        )
+        temperature = gr.Slider(
+            0.0,
+            2.0,
+            value=0.7,
+            step=0.05,
+            label="Temperature",
+        )
+        top_p = gr.Slider(
+            0.0,
+            1.0,
+            value=0.95,
+            step=0.01,
+            label="Top-p",
+        )
+        max_new_tokens = gr.Slider(
+            16,
+            MAX_SLIDER,
+            value=DEFAULT_MAX_TOKENS,
+            step=16,
+            label="Max new tokens",
+        )
+
+    clear = gr.Button("🧹 Clear chat")
+    clear.click(lambda: ([], ""), inputs=[], outputs=[chatbot, msg])
+
+    inputs = [model, msg, chatbot, system, temperature, top_p, max_new_tokens]
+    outputs = [chatbot, msg]
+
+    msg.submit(ui_respond, inputs, outputs)
+    send.click(ui_respond, inputs, outputs)
+
+
+demo.queue()
+app = gr.mount_gradio_app(app, demo, path="/")