Configure ZeroGPU runtime

README.md

@@ -17,9 +17,9 @@ models:
 
 # tinyLLM 0.51B Demo
 
-This CPU Space loads the public tinyLLM SFT model and switches the ARC GRPO or
-IFEval OPD LoRA without keeping multiple base models in memory.
+This free ZeroGPU Space loads the public tinyLLM SFT model and switches the ARC
+GRPO or IFEval OPD LoRA without keeping multiple base models in memory.
 
-The first response is slower because the 0.51B checkpoint is downloaded and
-loaded on demand. The full multimodal demo remains available in the GitHub
-repository; the free CPU Space focuses on text and adapter switching.
+The first start is slower while the 0.51B checkpoint is downloaded. The full
+multimodal demo remains available in the GitHub repository; this small hosted
+Space focuses on text and adapter switching.

app.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 import threading
 
 import gradio as gr
+import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
@@ -29,8 +30,8 @@ def load_runtime():
             _model = AutoModelForCausalLM.from_pretrained(
                 BASE_REPO,
                 trust_remote_code=True,
-                dtype=torch.float32,
-            ).eval()
+                dtype=torch.bfloat16,
+            ).to("cuda").eval()
         return _model, _tokenizer
 
 
@@ -47,6 +48,7 @@ def select_adapter(model, mode: str) -> None:
         model.activate_single_lora(name)
 
 
+@spaces.GPU(duration=120)
 def respond(message: str, history: list[dict], mode: str, max_new_tokens: int):
     if not message.strip():
         yield "请输入一个问题。"
@@ -65,7 +67,10 @@ def respond(message: str, history: list[dict], mode: str, max_new_tokens: int):
             tokenize=False,
             add_generation_prompt=True,
         )
-        inputs = tokenizer(prompt, return_tensors="pt")
+        inputs = {
+            key: value.to(next(model.parameters()).device)
+            for key, value in tokenizer(prompt, return_tensors="pt").items()
+        }
         streamer = TextIteratorStreamer(
             tokenizer,
             skip_prompt=True,
@@ -89,6 +94,11 @@ def respond(message: str, history: list[dict], mode: str, max_new_tokens: int):
         worker.join()
 
 
+# ZeroGPU provides a CUDA emulation context during startup, so the weights can
+# stay GPU-resident and are ready when the decorated function receives a slot.
+load_runtime()
+
+
 with gr.Blocks(title="tinyLLM 0.51B") as demo:
     gr.Markdown(
         "# tinyLLM 0.51B\n"