"""THOX ZeroGPU model Space — one model behind an OpenAI-compatible /v1 endpoint. This single file is deployed verbatim to every THOX ZeroGPU model Space. The only per-Space difference is `thox_space_config.py`, which names the model to serve. Why FastAPI *and* Gradio: Open WebUI consumes this Space as an ordinary "OpenAI API" connection, which means it needs `/v1/models` and `/v1/chat/completions`. Gradio alone does not speak that protocol, so the FastAPI app owns the `/v1` routes and the Gradio Blocks UI is mounted at `/` for humans. ZeroGPU only schedules Gradio SDK Spaces, so the Gradio app must genuinely be present -- it is not decoration. """ from __future__ import annotations import json import re import os import time import uuid from threading import Thread from typing import Any, Iterator # Disable Gradio 5 server-side rendering BEFORE gradio is imported. # On a Spaces Gradio SDK Space, SSR spawns a Node.js front-end server that claims the # port this app needs, producing "[Errno 98] address already in use" followed by # "Stopping Node.js server...". We serve the Blocks app from our own uvicorn process, so # the SSR layer has nothing to add here and everything to collide with. os.environ.setdefault("GRADIO_SSR_MODE", "false") import gradio as gr # noqa: E402 import spaces import torch import uvicorn from fastapi import FastAPI, Request from fastapi.responses import JSONResponse, StreamingResponse from pydantic import BaseModel, Field from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer import thox_auth # noqa: E402 from thox_space_config import ( ADAPTER_ID, CAVEAT, CHAT_TEMPLATE_KWARGS, DESCRIPTION, MAX_CONTEXT_TOKENS, MODEL_ID, SERVED_MODEL_NAME, SYSTEM_PROMPT, ) # -------------------------------------------------------------------------------------- # Model loading (module scope, eager -- see ZeroGPU rules) # -------------------------------------------------------------------------------------- # ZeroGPU monkey-patches torch so `cuda` resolves at module scope and the weights are # registered with the scheduler, then streamed into VRAM when a @spaces.GPU call lands. # Loading lazily on first request would just move this cost onto the first user. DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") DTYPE = torch.bfloat16 if DEVICE.type == "cuda" else torch.float32 print(f"[thox] loading base model: {MODEL_ID}", flush=True) TOKENIZER = AutoTokenizer.from_pretrained(MODEL_ID) if TOKENIZER.pad_token_id is None: TOKENIZER.pad_token = TOKENIZER.eos_token # ASSEMBLE ON CPU, MOVE ONCE AT THE END. # # ZeroGPU intercepts .to("cuda") at module scope and *registers* tensors with its # scheduler rather than allocating them -- there is no GPU attached outside a # @spaces.GPU call. PEFT does not go through that path: load_peft_weights calls # safe_load_file(..., device=), which is a REAL CUDA allocation, so # attaching an adapter to an already-"cuda" model dies at import with # "RuntimeError: No CUDA GPUs are available". # # Building base + adapter + merge entirely on CPU and moving the finished module once # keeps every CUDA touch inside the single call ZeroGPU is prepared for. MODEL = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=DTYPE) if ADAPTER_ID: # THOX ships several of its models as LoRA adapters rather than merged weights. from peft import PeftModel print(f"[thox] applying LoRA adapter on CPU: {ADAPTER_ID}", flush=True) # PEFT picks its own load device: load_peft_weights() -> infer_device(), which # branches on torch.cuda.is_available(). `import spaces` monkey-patches that to # return True unconditionally, so PEFT selects "cuda" and safe_load_file() attempts # a real allocation in a process with no GPU attached -- "RuntimeError: No CUDA GPUs # are available", at import, before anything can be served. Putting the model on CPU # does not help, because PEFT never consults the model. # # So the flag PEFT actually reads is told the truth for the duration of the load. # Restored immediately after: the rest of the app (and ZeroGPU itself) depends on the # patched value being True. _patched_is_available = torch.cuda.is_available torch.cuda.is_available = lambda: False try: MODEL = PeftModel.from_pretrained(MODEL, ADAPTER_ID) MODEL = MODEL.merge_and_unload() # merge so generation runs at full speed finally: torch.cuda.is_available = _patched_is_available MODEL = MODEL.to(DEVICE) MODEL.eval() print(f"[thox] ready: serving '{SERVED_MODEL_NAME}' on {DEVICE}", flush=True) # -------------------------------------------------------------------------------------- # Generation # -------------------------------------------------------------------------------------- def _encode(messages: list[dict], tools: list | None = None): """Apply the chat template, returning the dict form. return_dict=True is passed explicitly: transformers v5 flipped the default, so the older `...(return_tensors="pt").to(device)` idiom yields a BatchEncoding whose `.shape` lookup falls through __getattr__ into KeyError('shape') -- which crosses the ZeroGPU process boundary as a bare AttributeError with the cause stripped off. Asking for the dict form is correct on both 4.x and 5.x. CHAT_TEMPLATE_KWARGS carries per-model template switches (e.g. Qwen3's enable_thinking). Unknown keys are simply template variables, so a model whose template ignores them is unaffected. """ kwargs = dict(CHAT_TEMPLATE_KWARGS) # Passing tools= lets the model's OWN chat template render the tool schema in the # exact form it was trained on. Qwen2.5-Coder and Hermes-3 both support this; a # template that does not simply ignores the argument. Hand-writing tool descriptions # into the prompt instead would fight whatever format the model actually learned. if tools: kwargs["tools"] = tools return TOKENIZER.apply_chat_template( messages, add_generation_prompt=True, return_tensors="pt", return_dict=True, **kwargs, ) def _estimate_duration(messages: list[dict], max_new_tokens: int, temperature: float, top_p: float, tools: list | None = None) -> int: """Request only the GPU seconds this call plausibly needs. NOTE: @spaces.GPU(duration=...) invokes this with the SAME arguments as the wrapped function, so this signature MUST track _stream_tokens exactly. Adding a parameter to one without the other fails at request time inside the ZeroGPU wrapper ("takes 4 positional arguments but 5 were given"), not at import. ZeroGPU pre-checks the *requested* duration against the caller's remaining quota, not the actual runtime -- so leaving everything at the 60s default makes short requests fail with "quota exceeded" once a user drops below 60s. Shorter requests also rank higher in the node-level queue. """ return int(min(120, max(15, (max_new_tokens / 12) + 10))) @spaces.GPU(duration=_estimate_duration) def _stream_tokens(messages: list[dict], max_new_tokens: int, temperature: float, top_p: float, tools: list | None = None) -> Iterator[tuple[str, str]]: """Yield ("delta", text) pairs, then exactly one ("finish", reason) pair. Runs in a separate ZeroGPU worker process. Everything crossing this boundary is pickled, so only plain strings are yielded -- never CUDA tensors, which would trigger a blocked torch.cuda._lazy_init() in the parent process. The trailing ("finish", reason) pair exists because the stop reason is only knowable inside the worker: the caller sees text, not token counts, and cannot tell a model that chose to stop from one that ran into the max_new_tokens ceiling. Clients rely on finish_reason to detect truncation, so reporting "stop" for a truncated response makes a cut-off answer look complete. """ # return_dict=True explicitly, and index the result. # transformers v5 flipped apply_chat_template's default to return_dict=True, so the # older `...(return_tensors="pt").to(device)` idiom yields a BatchEncoding whose # `.shape` lookup falls through __getattr__ to KeyError('shape') -- which crosses the # ZeroGPU process boundary as a bare AttributeError with the cause stripped off. # Asking for the dict form is correct on both 4.x and 5.x. enc = _encode(messages, tools) input_ids = enc["input_ids"] attention_mask = enc.get("attention_mask") # Guard the context window so an over-long conversation degrades predictably rather # than producing silent garbage. if input_ids.shape[-1] > MAX_CONTEXT_TOKENS: input_ids = input_ids[:, -MAX_CONTEXT_TOKENS:] if attention_mask is not None: attention_mask = attention_mask[:, -MAX_CONTEXT_TOKENS:] input_ids = input_ids.to(MODEL.device) if attention_mask is not None: attention_mask = attention_mask.to(MODEL.device) streamer = TextIteratorStreamer(TOKENIZER, skip_prompt=True, skip_special_tokens=True) kwargs: dict[str, Any] = dict( input_ids=input_ids, attention_mask=attention_mask, streamer=streamer, max_new_tokens=max_new_tokens, do_sample=temperature > 0, pad_token_id=TOKENIZER.pad_token_id, ) if temperature > 0: kwargs["temperature"] = temperature kwargs["top_p"] = top_p # Capture generate()'s output so the true stop reason is available after the stream # drains. Comparing generated length against the ceiling is exact, where inferring it # from the decoded text would not be. result: dict[str, Any] = {} def _run() -> None: try: result["ids"] = MODEL.generate(**kwargs) except Exception as exc: # noqa: BLE001 result["error"] = exc thread = Thread(target=_run) thread.start() for delta in streamer: if delta: yield "delta", delta thread.join() if "error" in result: raise result["error"] generated = int(result["ids"].shape[-1] - input_ids.shape[-1]) if "ids" in result else 0 yield "finish", ("length" if generated >= max_new_tokens else "stop") def _normalise_messages(raw: list[dict]) -> list[dict]: """Coerce OpenAI message content into the plain strings the chat template expects.""" out: list[dict] = [] if SYSTEM_PROMPT and not (raw and raw[0].get("role") == "system"): out.append({"role": "system", "content": SYSTEM_PROMPT}) for m in raw: content = m.get("content", "") if isinstance(content, list): # OpenAI multimodal array form -- keep only the text parts. content = "".join( p.get("text", "") for p in content if isinstance(p, dict) and p.get("type") == "text" ) msg: dict = {"role": m.get("role", "user"), "content": content or ""} # Tool linkage must survive normalisation. An assistant turn that made a call # carries tool_calls; the following tool turn carries tool_call_id pointing back # at it. Rebuilding messages as {role, content} alone silently breaks that chain, # so the model can make a call but can never be told what the tool returned -- # the round trip dies on the second turn, which looks like the tool "not working". if m.get("tool_calls"): msg["tool_calls"] = m["tool_calls"] if m.get("tool_call_id"): msg["tool_call_id"] = m["tool_call_id"] if m.get("name"): msg["name"] = m["name"] out.append(msg) return out # -------------------------------------------------------------------------------------- # Tool-call parsing # -------------------------------------------------------------------------------------- _TOOL_CALL_RX = re.compile(r"\s*(\{.*?\})\s*", re.S) # Qwen2.5-Coder ignores its own template's instruction and emits an # XML form instead (observed verbatim, inside a ```xml fence): # x{} # The template DOES render the Hermes instruction correctly - the model simply # deviates - so this is a model-behaviour gap, not a prompting bug, and the parser # has to absorb it or the coder model can never call a tool. _FUNCTION_CALL_XML_RX = re.compile( r"\s*\s*([^<]+?)\s*\s*" r"\s*(.*?)\s*\s*", re.S) # A FOURTH form, observed on thox-coder-7b in production: format 3 wrapped in a markdown # code fence. # ```json # {"name": "generate_video", "arguments": {...}} # ``` # Qwen2.5-Coder is a CODER model and fencing JSON is precisely what it was trained to do, # so this is not an edge case for it -- it is the default. Without unwrapping, the # bare-JSON branch never fires, because the text starts with a backtick rather than a # brace, and the model reads as refusing tool use rather than as using another wrapper. _CODE_FENCE_RX = re.compile(r"^\s*```(?:[A-Za-z0-9_+-]*)\s*\n(.*?)\n?\s*```\s*$", re.S) def _strip_code_fence(text: str) -> str: """Unwrap a fence that surrounds the ENTIRE response. Deliberately anchored at both ends: a fence embedded in prose is a model TALKING about JSON, not calling something, and unwrapping that would turn an explanation into a spurious invocation. """ m = _CODE_FENCE_RX.match(text) return m.group(1).strip() if m else text def _extract_tool_calls(text: str) -> tuple[list[dict], str]: """Pull tool calls out of generated text -> (openai_tool_calls, remaining_content). Parsed in the PARENT process, deliberately. The @spaces.GPU worker keeps yielding plain ("delta"|"finish", str) pairs; putting structured objects through that boundary would mean pickling them across process lines for no benefit. THREE FORMATS, because the fleet is not one model family: * {"name":..., "arguments":{...}} Qwen2.5-Coder and Hermes-3 (Qwen adopted the Hermes convention). * x{...} Qwen2.5-Coder's actual observed output, despite its template instructing the tagged JSON form above. * a bare JSON object with name/arguments and no tags Llama-3.2 derivatives (thox-mini-3b) do not emit the Hermes tags. Supporting only the tagged form would leave one model in the fleet silently unable to call anything, which is worse than not shipping the feature - it would look like a model quality problem rather than a protocol gap. """ calls: list[dict] = [] remaining = text for m in _TOOL_CALL_RX.finditer(text): try: obj = json.loads(m.group(1)) except json.JSONDecodeError: continue calls.append(obj) remaining = remaining.replace(m.group(0), "") if not calls: for m in _FUNCTION_CALL_XML_RX.finditer(text): raw_args = (m.group(2) or "").strip() try: args = json.loads(raw_args) if raw_args else {} except json.JSONDecodeError: args = {} calls.append({"name": m.group(1).strip(), "arguments": args}) remaining = remaining.replace(m.group(0), "") if not calls: # Unwrap a whole-response code fence first -- see _strip_code_fence. stripped = _strip_code_fence(text).strip() # Only treat a bare object as a call when it actually looks like one; a model # answering a question ABOUT JSON must not be mistaken for invoking something. # Widening the bare-JSON path to fenced text widened the false-positive surface # too, so the shape test is tightened in the same change: a real call is # name+arguments (or name+parameters), or a bare name for a zero-argument tool. # {"name": "Alice", "age": 30} satisfies the old test and is not a call. if stripped.startswith("{") and '"name"' in stripped: try: obj = json.loads(stripped) looks_like_call = ( isinstance(obj, dict) and "name" in obj and ("arguments" in obj or "parameters" in obj or set(obj) == {"name"}) ) if looks_like_call: calls.append(obj) remaining = "" except json.JSONDecodeError: pass out: list[dict] = [] for i, c in enumerate(calls): args = c.get("arguments", c.get("parameters", {})) if not isinstance(args, str): args = json.dumps(args or {}) out.append({ "id": f"call_{uuid.uuid4().hex[:20]}", "type": "function", "index": i, "function": {"name": c.get("name", ""), "arguments": args}, }) return out, remaining.strip() # -------------------------------------------------------------------------------------- # OpenAI-compatible API # -------------------------------------------------------------------------------------- class ChatMessage(BaseModel): role: str content: Any = "" # Carried so a tool round trip can complete: the assistant turn that made the call, # and the tool turn that answers it, must keep their linkage. tool_calls: Any = None tool_call_id: str | None = None name: str | None = None # Ceilings are enforced by CLAMPING, not by Pydantic bounds. # # `Field(le=4096)` looks tidier but is an OpenAI incompatibility: FastAPI turns a # violation into `422 {"detail": [...]}`, while the OpenAI API returns # `400 {"error": {...}}`. Measured with the real `openai` SDK, `max_tokens=100000` raised # `UnprocessableEntityError` on all three template Spaces -- a client that reads # `error.message` sees nothing at all. Out-of-range sampling values are the caller being # sloppy, not the caller being wrong, so clamp them and answer. MAX_OUTPUT_TOKENS = 4096 def _clamp(value, low, high, fallback): try: return min(max(type(low)(value), low), high) except (TypeError, ValueError): return fallback class ChatRequest(BaseModel): model: str | None = None messages: list[ChatMessage] stream: bool = False max_tokens: int | None = None temperature: float = 0.7 top_p: float = 0.95 # Without these two fields Pydantic silently DROPS what the client sent, so the # model is never told a tool exists and every call looks like "I have no tools". tools: Any = None tool_choice: Any = None api = FastAPI(title=f"THOX ZeroGPU — {SERVED_MODEL_NAME}") @api.get("/health") def health() -> dict: # Open on purpose, and must stay free of anything secret. return { "status": "ok", "auth": "required" if thox_auth.is_configured() else "UNCONFIGURED", "model": SERVED_MODEL_NAME, "base": MODEL_ID, "adapter": ADAPTER_ID, } @api.get("/v1/models") def list_models(request: Request): # Gated too, deliberately. An unauthenticated /v1/models is how a bogus key gets # mistaken for a working one -- it answers 200 and nothing ever proves the key. denied = thox_auth.check(request) if denied: return denied return { "object": "list", "data": [{ "id": SERVED_MODEL_NAME, "object": "model", "created": 0, "owned_by": "thox", }], } @api.post("/v1/chat/completions") def chat_completions(req: ChatRequest, request: Request): denied = thox_auth.check(request) if denied: return denied messages = _normalise_messages([m.model_dump() for m in req.messages]) # tool_choice "none" means the caller explicitly wants no tool use this turn. tools = None if (req.tool_choice == "none") else (req.tools or None) max_new = _clamp(req.max_tokens or 512, 1, MAX_OUTPUT_TOKENS, 512) temperature = _clamp(req.temperature, 0.0, 2.0, 0.7) top_p = _clamp(req.top_p, 0.001, 1.0, 0.95) completion_id = f"chatcmpl-{uuid.uuid4().hex[:24]}" created = int(time.time()) if not req.stream: parts: list[str] = [] finish = "stop" for kind, payload in _stream_tokens(messages, max_new, temperature, top_p, tools): if kind == "delta": parts.append(payload) else: finish = payload text = "".join(parts) prompt_tokens = int(_encode(messages, tools)["input_ids"].shape[-1]) completion_tokens = len(TOKENIZER.encode(text)) message: dict = {"role": "assistant", "content": text} if tools: tool_calls, remaining = _extract_tool_calls(text) if tool_calls: # OpenAI contract: content is null when the turn is a tool call, and # finish_reason becomes "tool_calls". Clients branch on finish_reason to # decide whether to execute a tool, so reporting "stop" here means the # call is emitted and then never run. message = {"role": "assistant", "content": remaining or None, "tool_calls": tool_calls} finish = "tool_calls" return JSONResponse({ "id": completion_id, "object": "chat.completion", "created": created, "model": SERVED_MODEL_NAME, "choices": [{ "index": 0, "message": message, "finish_reason": finish, }], "usage": { "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "total_tokens": prompt_tokens + completion_tokens, }, }) def sse() -> Iterator[str]: base = { "id": completion_id, "object": "chat.completion.chunk", "created": created, "model": SERVED_MODEL_NAME, } # Open WebUI expects the role to arrive before any content delta. yield f"data: {json.dumps({**base, 'choices': [{'index': 0, 'delta': {'role': 'assistant'}, 'finish_reason': None}]})}\n\n" finish = "stop" # With tools in play the text must be BUFFERED rather than streamed through. # A tool call is only recognisable once its closing tag has arrived, so emitting # deltas as they appear would stream the raw {...} markup # into the user's chat window before we could tell it was a call. This costs the # typing effect only on tool-enabled turns; ordinary chat still streams. buffered: list[str] = [] try: for kind, payload in _stream_tokens(messages, max_new, temperature, top_p, tools): if kind == "finish": finish = payload continue if tools: buffered.append(payload) continue chunk = {**base, "choices": [{"index": 0, "delta": {"content": payload}, "finish_reason": None}]} yield f"data: {json.dumps(chunk)}\n\n" if tools: text = "".join(buffered) tool_calls, remaining = _extract_tool_calls(text) if tool_calls: delta: dict = {"tool_calls": tool_calls} if remaining: delta["content"] = remaining finish = "tool_calls" yield f"data: {json.dumps({**base, 'choices': [{'index': 0, 'delta': delta, 'finish_reason': None}]})}\n\n" elif text: yield f"data: {json.dumps({**base, 'choices': [{'index': 0, 'delta': {'content': text}, 'finish_reason': None}]})}\n\n" except Exception as exc: # noqa: BLE001 # Headers are already sent, so a failure cannot become an HTTP 500. Emit a # sanitised terminal event instead of tearing the connection down mid-frame, # which Open WebUI surfaces as a TransferEncodingError. print(f"[thox] generation failed: {type(exc).__name__}", flush=True) err = {**base, "choices": [{"index": 0, "delta": {}, "finish_reason": "error"}]} yield f"data: {json.dumps(err)}\n\n" yield "data: [DONE]\n\n" return done = {**base, "choices": [{"index": 0, "delta": {}, "finish_reason": finish}]} yield f"data: {json.dumps(done)}\n\n" yield "data: [DONE]\n\n" return StreamingResponse(sse(), media_type="text/event-stream", headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}) # -------------------------------------------------------------------------------------- # Gradio UI (mounted at /) # -------------------------------------------------------------------------------------- # Brand v2 theme + shared chrome. Canonical source: thox_theme.py at the repo # root, copied into each Space by build.py and guarded byte-for-byte by # ci/test_theme_parity.py. Do not inline a second copy here. from thox_theme import THOX_THEME, THOX_CSS, THOX_GREEN, THOX_ACCENT # noqa: E402 CSS = THOX_CSS def _ui_respond(message: str, history: list[dict]) -> Iterator[str]: convo = _normalise_messages([*history, {"role": "user", "content": message}]) acc = "" for kind, payload in _stream_tokens(convo, 512, 0.7, 0.95): if kind == "delta": acc += payload yield acc with gr.Blocks(css=CSS, title=f"THOX · {SERVED_MODEL_NAME}", theme=THOX_THEME) as demo: gr.HTML( f"""

THOX · {SERVED_MODEL_NAME}

{DESCRIPTION}

OpenAI-compatible endpoint: /v1/chat/completions · Your AI. Your Data. Your Rules.™

""" ) if CAVEAT: # Disclosed in the UI, not just the README: someone chatting here is exactly the # person who would otherwise take a confident wrong answer at face value. gr.HTML(f'
Known defect. {CAVEAT}
') gr.ChatInterface(fn=_ui_respond, type="messages") # Expose the FastAPI app for anyone importing this module (and for local `uvicorn app:app`). app = api if __name__ == "__main__": # LET GRADIO OWN THE SERVER. # # The obvious construction -- gr.mount_gradio_app(fastapi_app, demo) plus our own # uvicorn.run() -- does not survive on a Gradio SDK Space. The platform launches the # Blocks app itself, so our uvicorn came up as a redundant second server and was # SIGTERM'd about one second after binding ("Uvicorn running ..." immediately # followed by "Shutting down"). Binding the port the platform wanted instead failed # with "[Errno 98] address already in use", because the platform's own launch had # already taken it. Both failures are the same underlying mistake: competing for the # server rather than joining it. # # So Gradio launches (which is also what ZeroGPU expects), and the OpenAI routes are # grafted onto the FastAPI app Gradio builds. Starlette matches each request against # a plain list of routes, so inserting ahead of Gradio's catch-all is enough to make # /v1/* resolve to us while everything else stays Gradio's. demo.queue(max_size=32) _, local_url, _ = demo.launch(prevent_thread_lock=True, server_name="0.0.0.0", ssr_mode=False, show_api=False) print(f"[thox] gradio launched at {local_url}", flush=True) grafted = [] for route in api.routes: path = getattr(route, "path", "") if path.startswith("/v1") or path == "/health": demo.app.router.routes.insert(0, route) grafted.append(path) print(f"[thox] grafted OpenAI routes onto gradio app: {grafted}", flush=True) thox_auth.install(demo.app) # launch(prevent_thread_lock=True) returns immediately; block so PID 1 stays alive. demo.block_thread()