Spaces:
Running on Zero
Running on Zero
| """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=<model.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))) | |
| 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"<tool_call>\s*(\{.*?\})\s*</tool_call>", re.S) | |
| # Qwen2.5-Coder ignores its own template's <tool_call> instruction and emits an | |
| # XML form instead (observed verbatim, inside a ```xml fence): | |
| # <response><function_call><name>x</name><arguments>{}</arguments></function_call></response> | |
| # 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"<function_call>\s*<name>\s*([^<]+?)\s*</name>\s*" | |
| r"<arguments>\s*(.*?)\s*</arguments>\s*</function_call>", 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: | |
| * <tool_call>{"name":..., "arguments":{...}}</tool_call> | |
| Qwen2.5-Coder and Hermes-3 (Qwen adopted the Hermes convention). | |
| * <function_call><name>x</name><arguments>{...}</arguments></function_call> | |
| 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}") | |
| 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, | |
| } | |
| 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", | |
| }], | |
| } | |
| 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 <tool_call>{...}</tool_call> 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"""<div id="thox-header"> | |
| <h1>THOX · {SERVED_MODEL_NAME}</h1> | |
| <p>{DESCRIPTION}</p> | |
| <p>OpenAI-compatible endpoint: <code>/v1/chat/completions</code> · | |
| Your AI. Your Data. Your Rules.™</p> | |
| </div>""" | |
| ) | |
| 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'<div id="thox-caveat"><strong>Known defect.</strong> {CAVEAT}</div>') | |
| 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() | |