spaces gpu decorator

app.py

@@ -17,7 +17,6 @@ from diffusers import DiffusionPipeline
 
 
 # Global variables for model and scheduler
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = None
 ddim_scheduler = None
 model_type = None
@@ -42,9 +41,12 @@ def convert_to_np(torch_image):
     return ((torch_image.detach().clamp(-1, 1).cpu().numpy().transpose(1, 2, 0) + 1) * 127.5).astype(np.uint8)
 
 
+@spaces.GPU
 def generate_noisy_image(image_choice, noise_sigma, operator_key):
     """Generate the noisy image and store necessary data for restoration."""
     # Map image choice to path
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
     image_paths = {
         "CelebA HQ 1": "sample_images/celebhq_29999.jpg",
         "CelebA HQ 2": "sample_images/celebhq_00001.jpg",
@@ -86,6 +88,8 @@ def run_restoration(data, T, K):
     """Run the restoration process and return the restored image."""
     global model, ddim_scheduler, model_type
 
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
     # Extract stored data
     noisy_measurement = data['noisy_measurement']
     operator = data['operator']