Update app.py

app.py

@@ -8,8 +8,7 @@ from typing import AsyncIterator, Optional
 import gradio as gr
 import spaces
 import torch
-import uvicorn
-from fastapi import FastAPI, HTTPException
+from fastapi import HTTPException
 from fastapi.responses import JSONResponse, StreamingResponse
 from pydantic import BaseModel, Field
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
@@ -21,11 +20,6 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStream
 MODEL_ID = "Qwen/Qwen3-30B-A3B"
 MODEL_ALIAS = "qwen3-30b-a3b"
 
-# ---------------------------------------------------------------------------
-# Model loading — tokenizer on CPU at startup; model loaded with device_map
-# so ZeroGPU can manage GPU placement per request.
-# ---------------------------------------------------------------------------
-
 print(f"Loading tokenizer for {MODEL_ID} …")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
@@ -38,12 +32,6 @@ model = AutoModelForCausalLM.from_pretrained(
 model.eval()
 print("Model ready.")
 
-# ---------------------------------------------------------------------------
-# FastAPI app
-# ---------------------------------------------------------------------------
-
-app = FastAPI(title="Qwen3-30B-A3B OpenAI-compatible API")
-
 # ---------------------------------------------------------------------------
 # Pydantic schemas
 # ---------------------------------------------------------------------------
@@ -90,17 +78,12 @@ def make_generation_kwargs(request: ChatCompletionRequest) -> dict:
 
 
 # ---------------------------------------------------------------------------
-# GPU generation functions
-# NOTE: ZeroGPU requires at least one @spaces.GPU function to be wired into
-# the Gradio UI (not just defined). We satisfy this by using `gradio_chat`
-# as both the Gradio interface handler AND calling the same underlying logic
-# that the FastAPI routes use.
+# GPU generation functions — these are the ZeroGPU anchors.
 # ---------------------------------------------------------------------------
 
 
 @spaces.GPU
 def gradio_chat(message: str, history: list) -> str:
-    """Gradio-facing chat handler — also acts as the ZeroGPU anchor function."""
     hf_messages = [{"role": "user" if i % 2 == 0 else "assistant", "content": m}
                    for i, m in enumerate([msg for pair in history for msg in pair] + [message])]
     prompt = tokenizer.apply_chat_template(
@@ -122,7 +105,6 @@ def gradio_chat(message: str, history: list) -> str:
 
 @spaces.GPU
 def _generate_response(prompt: str, gen_kwargs: dict) -> str:
-    """Non-streaming generation for FastAPI."""
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     with torch.no_grad():
         output_ids = model.generate(**inputs, **gen_kwargs)
@@ -132,7 +114,6 @@ def _generate_response(prompt: str, gen_kwargs: dict) -> str:
 
 @spaces.GPU
 def _generate_streaming(prompt: str, gen_kwargs: dict, streamer: TextIteratorStreamer):
-    """Streaming generation for FastAPI."""
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     with torch.no_grad():
         model.generate(**inputs, streamer=streamer, **gen_kwargs)
@@ -167,95 +148,99 @@ def stream_chunk(delta_content: str, model_name: str, completion_id: str, finish
 
 
 # ---------------------------------------------------------------------------
-# FastAPI routes
+# Gradio UI
 # ---------------------------------------------------------------------------
 
+with gr.Blocks(title=f"{MODEL_ALIAS} API") as demo:
+    gr.Markdown(f"""
+# {MODEL_ALIAS} — OpenAI-compatible API
 
-@app.get("/v1/models")
-async def list_models():
-    return {
-        "object": "list",
-        "data": [{"id": MODEL_ALIAS, "object": "model", "created": int(time.time()), "owned_by": "qwen"}],
-    }
-
-
-@app.post("/v1/chat/completions")
-async def chat_completions(request: ChatCompletionRequest):
-    try:
-        prompt = build_prompt(request.messages, request.enable_thinking or False)
-        gen_kwargs = make_generation_kwargs(request)
-    except Exception as exc:
-        raise HTTPException(status_code=422, detail=str(exc))
-
-    if request.stream:
-        completion_id = f"chatcmpl-{uuid.uuid4().hex}"
+Point **Paperclip** at `https://<your-space>.hf.space` with model `{MODEL_ALIAS}`.
 
-        async def token_generator() -> AsyncIterator[str]:
-            role_chunk = {
-                "id": completion_id, "object": "chat.completion.chunk",
-                "created": int(time.time()), "model": request.model,
-                "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}],
-            }
-            yield f"data: {json.dumps(role_chunk)}\n\n"
+| Method | Path | Description |
+|--------|------|-------------|
+| GET | `/v1/models` | List models |
+| POST | `/v1/chat/completions` | Chat (streaming & non-streaming) |
+| GET | `/health` | Health check |
 
-            streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-            thread = Thread(target=_generate_streaming, args=(prompt, gen_kwargs, streamer), daemon=True)
-            thread.start()
+You can also chat directly below.
+""")
+    gr.ChatInterface(fn=gradio_chat)
 
-            try:
-                for token_text in streamer:
-                    if token_text:
-                        yield stream_chunk(token_text, request.model, completion_id)
-                    await asyncio.sleep(0)
-            finally:
-                thread.join()
+# ---------------------------------------------------------------------------
+# Register FastAPI-style routes onto Gradio's own FastAPI app via app_kwargs,
+# BEFORE launch() — so ZeroGPU's launch hook still does its registration scan.
+# ---------------------------------------------------------------------------
 
-            yield stream_chunk("", request.model, completion_id, finish_reason="stop")
-            yield "data: [DONE]\n\n"
 
-        return StreamingResponse(token_generator(), media_type="text/event-stream",
-                                 headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})
+def add_custom_routes(fastapi_app):
+    @fastapi_app.get("/v1/models")
+    async def list_models():
+        return {
+            "object": "list",
+            "data": [{"id": MODEL_ALIAS, "object": "model", "created": int(time.time()), "owned_by": "qwen"}],
+        }
 
-    try:
-        content = _generate_response(prompt, gen_kwargs)
-    except Exception as exc:
-        raise HTTPException(status_code=500, detail=f"Generation failed: {exc}")
+    @fastapi_app.post("/v1/chat/completions")
+    async def chat_completions(request: ChatCompletionRequest):
+        try:
+            prompt = build_prompt(request.messages, request.enable_thinking or False)
+            gen_kwargs = make_generation_kwargs(request)
+        except Exception as exc:
+            raise HTTPException(status_code=422, detail=str(exc))
 
-    return JSONResponse(chat_completion_object(content, request.model))
+        if request.stream:
+            completion_id = f"chatcmpl-{uuid.uuid4().hex}"
 
+            async def token_generator() -> AsyncIterator[str]:
+                role_chunk = {
+                    "id": completion_id, "object": "chat.completion.chunk",
+                    "created": int(time.time()), "model": request.model,
+                    "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}],
+                }
+                yield f"data: {json.dumps(role_chunk)}\n\n"
 
-@app.get("/health")
-async def health():
-    return {"status": "ok", "model": MODEL_ID}
+                streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+                thread = Thread(target=_generate_streaming, args=(prompt, gen_kwargs, streamer), daemon=True)
+                thread.start()
 
+                try:
+                    for token_text in streamer:
+                        if token_text:
+                            yield stream_chunk(token_text, request.model, completion_id)
+                        await asyncio.sleep(0)
+                finally:
+                    thread.join()
 
-# ---------------------------------------------------------------------------
-# Gradio UI — must have a real wired @spaces.GPU function for ZeroGPU
-# ---------------------------------------------------------------------------
+                yield stream_chunk("", request.model, completion_id, finish_reason="stop")
+                yield "data: [DONE]\n\n"
 
-with gr.Blocks(title=f"{MODEL_ALIAS} API") as demo:
-    gr.Markdown(f"""
-# {MODEL_ALIAS} — OpenAI-compatible API
+            return StreamingResponse(token_generator(), media_type="text/event-stream",
+                                     headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})
 
-Point **Paperclip** at `https://<your-space>.hf.space` with model `{MODEL_ALIAS}`.
+        try:
+            content = _generate_response(prompt, gen_kwargs)
+        except Exception as exc:
+            raise HTTPException(status_code=500, detail=f"Generation failed: {exc}")
 
-| Method | Path | Description |
-|--------|------|-------------|
-| GET | `/v1/models` | List models |
-| POST | `/v1/chat/completions` | Chat (streaming & non-streaming) |
-| GET | `/health` | Health check |
+        return JSONResponse(chat_completion_object(content, request.model))
 
-Chat UI is served at the root path below; API routes remain available alongside it.
-""")
-    # This ChatInterface wires `gradio_chat` into Gradio's event system,
-    # which is what ZeroGPU's startup scanner requires.
-    gr.ChatInterface(fn=gradio_chat)
+    @fastapi_app.get("/health")
+    async def health():
+        return {"status": "ok", "model": MODEL_ID}
 
-app = gr.mount_gradio_app(app, demo, path="/")
 
 # ---------------------------------------------------------------------------
-# Entry-point
+# Entry-point — demo.launch() is what triggers spaces.one_launch's patched
+# hook, which is what registers @spaces.GPU functions with the HF platform.
 # ---------------------------------------------------------------------------
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    demo.queue()
+    add_custom_routes(demo.app)
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        app_kwargs={"docs_url": None},
+    )
+