Updated app.py

app.py

@@ -2,10 +2,9 @@ import gradio as gr
 import subprocess
 import os
 import shutil
-import os
 from huggingface_hub import hf_hub_download
 import torch
-
+import spaces  # Import spaces for GPU decoration
 
 # Define paths
 MODEL_DIR = "./model"  # Local directory to store the downloaded model
@@ -26,6 +25,7 @@ def download_model():
         print("Dataset004_WML downloaded and extracted.")
 
 # Function to run nnUNet inference
+@spaces.GPU  # Decorate the function to allocate GPU for its execution
 def run_nnunet_predict(nifti_file):
     # Prepare directories
     os.makedirs(INPUT_DIR, exist_ok=True)
@@ -49,19 +49,20 @@ def run_nnunet_predict(nifti_file):
         "-o", OUTPUT_DIR,
         "-d", "004",                  # Dataset ID
         "-c", "3d_fullres",           # Configuration
-        "-tr", "nnUNetTrainer_8000epochs",			
+        "-tr", "nnUNetTrainer_8000epochs",
+        "--device", "cuda:0"  # Explicitly use GPU 0
     ]
-    print("Files in /tmp/output:")
-    print(os.listdir(OUTPUT_DIR))
     try:
         subprocess.run(command, check=True)
         # Get the output file
         output_file = os.path.join(OUTPUT_DIR, "image.nii.gz")
-        return output_file
+        if os.path.exists(output_file):
+            return output_file
+        else:
+            return "Error: Output file not found."
     except subprocess.CalledProcessError as e:
         return f"Error: {e}"
 
-
 # Gradio Interface
 interface = gr.Interface(
     fn=run_nnunet_predict,
@@ -84,3 +85,4 @@ download_model()
 # Launch the app
 if __name__ == "__main__":
     interface.launch()
+