Enable ZeroGPU inference

app.py

@@ -9,15 +9,16 @@ import numpy as np
 import gymnasium as gym
 import imageio
 import gradio as gr
+import spaces
 from huggingface_hub import hf_hub_download
 
-# 1. Policy Network Architecture
+
 def layer_init(layer, std=np.sqrt(2), bias_const=0.0):
     nn.init.orthogonal_(layer.weight, std)
     nn.init.constant_(layer.bias, bias_const)
     return layer
 
-# Reconstruct pure Sequential actor
+
 def get_actor_network(state_dim=8, action_dim=4):
     actor = nn.Sequential(
         layer_init(nn.Linear(state_dim, 64)),
@@ -28,7 +29,8 @@ def get_actor_network(state_dim=8, action_dim=4):
     )
     return actor
 
-# 2. Inference & Rendering
+
+@spaces.GPU(duration=60)
 def simulate_agent(stage_selection):
     weight_mapping = {
         "Stage 1: Baseline": "1_baseline.pth",
@@ -39,20 +41,16 @@ def simulate_agent(stage_selection):
     filename = weight_mapping.get(stage_selection)
     repo_id = "ben-dlwlrma/Representation-Over-Routing" 
     
-    # Download weights from HF Hub
     try:
         weights_path = hf_hub_download(repo_id=repo_id, filename=filename)
     except Exception as e:
         raise gr.Error(f"Weight download failed. Error: {str(e)}")
 
-    # Initialize env
     env = gym.make("LunarLander-v3", render_mode="rgb_array")
     
-    # Initialize model on CPU
-    device = torch.device("cpu")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     actor = get_actor_network(state_dim=8, action_dim=4).to(device)
     
-    # Load weights
     try:
         actor.load_state_dict(torch.load(weights_path, map_location=device, weights_only=True))
         actor.eval()
@@ -85,7 +83,6 @@ def simulate_agent(stage_selection):
 
     env.close()
 
-    # Export to MP4
     video_filename = "eval_output.mp4"
     fps = 30
     try:
@@ -95,7 +92,7 @@ def simulate_agent(stage_selection):
 
     return video_filename
 
-# 3. Gradio Web UI
+
 with gr.Blocks(title="Representation over Routing", theme=gr.themes.Base()) as demo:
     gr.Markdown("## Representation over Routing")
     gr.Markdown("Multi-timescale RL evaluation environment. Select an ablation stage to visualize policy behavior.")