init

app.py

@@ -215,7 +215,7 @@ def visualize_drag_v2(background_image_path, splited_tracks, width, height):
 
 
 class Drag:
-    @spaces.GPU(duration=300)
+    @spaces.GPU(duration=200)
     def __init__(self, device, height, width, model_length):
         self.device = device
 
@@ -371,7 +371,8 @@ class Drag:
         outputs['flows'] = torch.from_numpy(total_nps).cuda().permute(0, 3, 1, 2).unsqueeze(0) / 255.
 
         return outputs
-
+    
+    @spaces.GPU
     @torch.no_grad()
     def get_cmp_flow_from_tracking_points(self, tracking_points, motion_brush_mask, first_frame_path):
 
@@ -487,7 +488,7 @@ class Drag:
 
         return viz_esti_flows
     
-    @spaces.GPU(duration=300)
+    @spaces.GPU(duration=200)
     def run(self, first_frame_path, tracking_points, inference_batch_size, motion_brush_mask, motion_brush_viz, ctrl_scale):
         
         original_width, original_height = self.width, self.height