ZeroGPU (#3)

- ZeroGPU (54ae3edca9b911d967abda51502a055270e80767)
- Update app.py (00ed98fea7ee505db889a8235c7657df0d1e9e62)

app.py

@@ -1,3 +1,5 @@
+import spaces
+
 import os
 import json
 import random
@@ -97,7 +99,9 @@ MODEL_CONTEXT: Dict[str, int] = {
     "CodeSoft/MetaDiffusion-150M-ChatBase": 5120,
 }
 
-DEVICE = "cpu"
+# ZeroGPU: CUDA is emulated at startup so models load onto cuda at module level;
+# real GPU is only mounted inside @spaces.GPU-decorated calls.
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 @dataclass
 class MetaDiffusionConfig:
@@ -551,12 +555,7 @@ def load_diffusion_manual():
         cfg = MetaDiffusionConfig(**valid)
         cfg.tie_word_embeddings = False
         mdl = MetaDiffusionLM(cfg).to(DEVICE)
-        try:
-            from safetensors.torch import load_file
-        except ImportError:
-            import subprocess, sys
-            subprocess.check_call([sys.executable, "-m", "pip", "install", "safetensors", "--quiet", "--break-system-packages"])
-            from safetensors.torch import load_file  # type: ignore
+        from safetensors.torch import load_file
         state = load_file(str(model_path), device="cpu")
         state = {k[len("model."):] if k.startswith("model.") else k: v for k, v in state.items()}
         missing, unexpected = mdl.load_state_dict(state, strict=False)
@@ -688,6 +687,7 @@ def build_inputs(tokenizer, model_id: str, prompt: str):
 def is_diffusion_model(model_id: str) -> bool:
     return "metadiffusion" in model_id.lower()
 
+@spaces.GPU(duration=120)
 def generate_for_model(model_id: str, prompt: str) -> str:
     ensure_models_loaded()
     if model_id not in models or model_id not in tokenizers:
@@ -1093,7 +1093,7 @@ def create_demo() -> gr.Blocks:
 # ---------------------------------------------------------------------------
 if __name__ == "__main__":
     print("=" * 60)
-    print("SLM Arena starting, attempting to load 4 models on CPU...")
+    print("SLM Arena starting, attempting to load 4 models...")
     print(f"Models: {MODEL_IDS}")
     print(f"Data dir: {get_data_dir().resolve()} (bucket /data if mounted)")
     print("=" * 60)
@@ -1113,4 +1113,4 @@ if __name__ == "__main__":
         logger.warning(f"Leaderboard preview failed: {e}")
     demo = create_demo()
     demo.queue(max_size=20)
-    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True, theme=gr.themes.Base(), css=CSS)
+    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True, theme=gr.themes.Base(), css=CSS)

requirements.txt

@@ -1,8 +1,8 @@
-torch==2.13.0
+torch==2.11.0
 transformers==5.14.1
 gradio==6.5.1
-pandas==3.0.3
+pandas==2.3.3
 safetensors==0.8.0
 huggingface_hub==1.22.0
 accelerate==1.14.0
-numpy==2.5.2
+numpy==2.2.6