Update app.py

app.py

@@ -4,7 +4,7 @@ import shutil
 import uuid
 import torch
 import random
-
+import spaces
 import gradio as gr
 import numpy as np
 
@@ -16,7 +16,7 @@ from pipeline_flux_control_removal import FluxControlRemovalPipeline
 
 torch.set_grad_enabled(False)
 os.environ['GRADIO_TEMP_DIR'] = './tmp'
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = "cuda"
 print(device)
 image_path = mask_path = None
 image_examples = [...] 
@@ -52,7 +52,7 @@ image_examples = [
     ]
     
 ]
-
+@spaces.GPU
 def load_model(base_model_path, lora_path):
     global pipe
     transformer = FluxTransformer2DModel.from_pretrained(base_model_path, subfolder='transformer', torch_dtype=torch.bfloat16)
@@ -86,7 +86,8 @@ def load_model(base_model_path, lora_path):
     gr.Info(str(f"Inject LoRA: {lora_path}"))
     pipe.load_lora_weights(lora_path, weight_name="pytorch_lora_weights.safetensors")
     gr.Info(str(f"Model loading: {int((100 / 100) * 100)}%"))
-
+    
+@spaces.GPU
 def set_seed(seed):
     torch.manual_seed(seed)
     torch.cuda.manual_seed(seed)
@@ -94,7 +95,7 @@ def set_seed(seed):
     np.random.seed(seed)
     random.seed(seed)
 
-    
+@spaces.GPU
 def predict(
     input_image,
     prompt,