Added gpu

app.py

@@ -236,6 +236,7 @@ class Drag:
         self.width = width
         self.model_length = model_length
 
+    @spaces.GPU(enable_queue=True, duration=240)
     def get_cmp_flow(self, frames, sparse_optical_flow, mask, brush_mask=None):
 
         '''
@@ -243,6 +244,8 @@ class Drag:
             sparse_optical_flow: [b, 13, 2, 384, 384] (-384, 384) tensor
             mask: [b, 13, 2, 384, 384] {0, 1} tensor
         '''
+        if not self.cmp or not self.pipeline:
+            self.run_model_init()
 
         b, t, c, h, w = frames.shape
         assert h == 384 and w == 384
@@ -486,9 +489,8 @@ class Drag:
         viz_esti_flows = flow_to_image(controlnet_flow)  # [h, w, c]
 
         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):
+    
+    def run_model_init(self):
         svd_ckpt = "./ckpts/stable-video-diffusion-img2vid-xt-1-1"
         mofa_ckpt = "./ckpts/controlnet/ckpts/controlnet"
     
@@ -502,6 +504,11 @@ class Drag:
         self.pipeline = self.pipeline.to("cuda:0")
         self.cmp = self.cmp.to("cuda:0")
 
+    @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):
+        if not self.cmp or not self.pipeline:
+            self.run_model_init()
+
         original_width, original_height = self.width, self.height
 
         input_all_points = tracking_points.constructor_args['value']