init

app.py

@@ -89,7 +89,7 @@ def get_sparseflow_and_mask_forward(
     return s_flow, mask
 
 
-@spaces.GPU(duration=300)
+@spaces.GPU(duration=200)
 def init_models(pretrained_model_name_or_path, resume_from_checkpoint, weight_dtype, device='cuda', enable_xformers_memory_efficient_attention=False, allow_tf32=False):
 
     from models.unet_spatio_temporal_condition_controlnet import UNetSpatioTemporalConditionControlNetModel
@@ -216,7 +216,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
 
@@ -488,7 +488,7 @@ class Drag:
 
         return viz_esti_flows
     
-    @spaces.GPU(duration=500)
+    @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

models/cmp/utils/common_utils.py

@@ -94,7 +94,6 @@ def accuracy(output, target, topk=(1,)):
     return res
 
 
-@spaces.GPU(duration=300)
 def load_state(path, model, optimizer=None):
     def map_func(storage, location):
         return storage.cuda()