Run on spaces gpu

app.py

@@ -5,6 +5,7 @@ import numpy as np
 import os
 from omegaconf import OmegaConf
 from PIL import Image 
+import spaces 
 
 from utils.app_utils import (
     remove_background, 
@@ -33,7 +34,6 @@ def main():
 
     if torch.cuda.is_available():
         device = "cuda:0"
-        torch.cuda.set_device(device)
         print("Found cuda")
     else:
         device = "cpu"
@@ -77,6 +77,7 @@ def main():
 
     ply_out_path = f'./mesh.ply'
 
+    @spaces.GPU()
     def reconstruct_and_export(image):
         """
         Passes image through model, outputs reconstruction in form of a dict of tensors.

requirements.txt

@@ -8,4 +8,5 @@ moviepy
 markupsafe==2.0.1
 gradio==4.25.0
 rembg
+spaces
 # git+https://github.com/graphdeco-inria/diff-gaussian-rasterization