Chnaged cuda

app.py

@@ -29,6 +29,7 @@ from diffusers.utils.import_utils import is_xformers_available
 from utils.flow_viz import flow_to_image
 from utils.utils import split_filename, image2arr, image2pil, ensure_dirname
 from huggingface_hub import login, hf_hub_download, snapshot_download
+import spaces
 
 
 output_dir_video = "./outputs/videos"
@@ -120,7 +121,7 @@ def init_models(pretrained_model_name_or_path, resume_from_checkpoint, weight_dt
     cmp = CMP_demo(
         './models/cmp/experiments/semiauto_annot/resnet50_vip+mpii_liteflow/config.yaml',
         42000
-    ).to(device)
+    )
     cmp.requires_grad_(False)
     
     # Freeze vae and image_encoder
@@ -497,8 +498,11 @@ class Drag:
 
         return viz_esti_flows
 
+    @spaces.GPU(enable_queue=True, duration=240)
     def run(self, first_frame_path, tracking_points, inference_batch_size, motion_brush_mask, motion_brush_viz, ctrl_scale):
-        
+        self.pipeline = self.pipeline.to("cuda:0")
+        self.cmp = self.cmp.to("cuda:0")
+
         original_width, original_height = self.width, self.height
 
         input_all_points = tracking_points.constructor_args['value']