Update app.py

app.py

@@ -3,7 +3,12 @@ import torch
 import subprocess
 import os
 import time
+from huggingface_hub import SpaceStage
+from huggingface_hub.utils import RepositoryNotFoundError
+from huggingface_hub.spaces import get_space_stage
+import spaces
 
+@spaces.GPU
 def check_gpu():
     results = []
     
@@ -50,6 +55,7 @@ def check_gpu():
     
     return "\n".join(results)
 
+@spaces.GPU
 def test_memory_allocation():
     try:
         # See how much GPU memory we can allocate
@@ -76,11 +82,30 @@ def test_memory_allocation():
     except Exception as e:
         return f"Memory test failed: {e}"
 
+def get_space_status():
+    try:
+        stage = get_space_stage()
+        if stage == SpaceStage.RUNNING:
+            status = "Space is running"
+        elif stage == SpaceStage.BUILDING:
+            status = "Space is building"
+        else:
+            status = f"Space is in stage: {stage}"
+    except RepositoryNotFoundError:
+        status = "Not running in a Space"
+    except Exception as e:
+        status = f"Error getting Space stage: {e}"
+    
+    return status
+
 # Create the Gradio interface
 with gr.Blocks(title="GPU Test") as demo:
     gr.Markdown("# GPU Availability Test")
     gr.Markdown("This app checks if GPU/CUDA is available and working in this Hugging Face Space")
     
+    # Display Space status
+    status_text = gr.Markdown(f"**Space Status**: {get_space_status()}")
+    
     with gr.Tab("Basic GPU Test"):
         check_btn = gr.Button("Check GPU Status", variant="primary")
         output = gr.Textbox(label="Results", lines=20)