Update app.py

app.py

@@ -19,6 +19,7 @@ aura_sr = AuraSR.from_pretrained("fal/AuraSR-v2")
 # Restore original torch.load
 torch.load = original_load
 
+@spaces.GPU
 def process_image(input_image):
     if input_image is None:
         raise gr.Error("Please provide an image to upscale.")
@@ -26,8 +27,7 @@ def process_image(input_image):
     # Convert to PIL Image for resizing
     pil_image = Image.fromarray(input_image)
 
-    # Upscale the image using AuraSR
-    upscaled_image = process_image_on_gpu(pil_image)
+    upscaled_image = aura_sr.upscale_4x(pil_image)
 
     # Convert result to numpy array if it's not already
     result_array = np.array(upscaled_image)
@@ -35,10 +35,6 @@ def process_image(input_image):
     print(input_image, result_array)
 
     return [input_image, result_array]
-
-@spaces.GPU
-def process_image_on_gpu(pil_image):
-    return aura_sr.upscale_4x(pil_image)
     
 title = """<h1 align="center">AuraSR</h1>
 <p><center>Upscales your images to x4</center></p>