changes for zeroGPU

app.py

@@ -1,11 +1,15 @@
+import spaces
 import gradio as gr
 import numpy as np
+impot torch
 
 from chrislib.general import uninvert, invert, view, view_scale
 
 from intrinsic.pipeline import load_models, run_pipeline
 
-intrinsic_models = load_models('v2', device='cpu')
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+intrinsic_models = load_models('v2', device=DEVICE)
 
 def generate_pipeline(models):
 
@@ -17,11 +21,12 @@ def generate_pipeline(models):
 
 pipeline_func = generate_pipeline(intrinsic_models)
 
+@spaces.GPU
 def process_image(image):
     print(image.shape)
     image = image.astype(np.single) / 255.
 
-    result = pipeline_func(image, device='cpu')
+    result = pipeline_func(image, device=DEVICE)
     
     return [view(result['hr_alb']), 1 - invert(result['dif_shd']), view_scale(result['pos_res'])]