init

app.py

@@ -216,7 +216,7 @@ def visualize_drag_v2(background_image_path, splited_tracks, width, height):
 
 class Drag:
     @spaces.GPU(duration=200)
-    def __init__(self, device, height, width, model_length):
+    def __init__(self, device, height, width):
         self.device = device
 
         svd_ckpt = "ckpts/stable-video-diffusion-img2vid-xt-1-1"
@@ -234,7 +234,6 @@ class Drag:
 
         self.height = height
         self.width = width
-        self.model_length = model_length
 
     def get_cmp_flow(self, frames, sparse_optical_flow, mask, brush_mask=None):
 
@@ -289,7 +288,7 @@ class Drag:
         '''
 
         seed = 42
-        num_frames = self.model_length
+        num_frames = 25
         
         set_seed(seed)
 
@@ -389,8 +388,8 @@ class Drag:
         new_resized_all_points = []
         new_resized_all_points_384 = []
         for tnum in range(len(resized_all_points)):
-            new_resized_all_points.append(interpolate_trajectory(input_all_points[tnum], self.model_length))
-            new_resized_all_points_384.append(interpolate_trajectory(resized_all_points_384[tnum], self.model_length))
+            new_resized_all_points.append(interpolate_trajectory(input_all_points[tnum], 25))
+            new_resized_all_points_384.append(interpolate_trajectory(resized_all_points_384[tnum], 25))
 
         resized_all_points = np.array(new_resized_all_points)
         resized_all_points_384 = np.array(new_resized_all_points_384)
@@ -408,24 +407,24 @@ class Drag:
             input_drag_384_inmask, input_mask_384_inmask = \
                 get_sparseflow_and_mask_forward(
                     resized_all_points_384_inmask, 
-                    self.model_length - 1, 384, 384
+                    25 - 1, 384, 384
                 )
         else:
             input_drag_384_inmask, input_mask_384_inmask = \
-                np.zeros((self.model_length - 1, 384, 384, 2)), \
-                    np.zeros((self.model_length - 1, 384, 384))
+                np.zeros((25 - 1, 384, 384, 2)), \
+                    np.zeros((25 - 1, 384, 384))
         
         if resized_all_points_384_outmask.shape[0] != 0:
             out_mask_flag = True
             input_drag_384_outmask, input_mask_384_outmask = \
                 get_sparseflow_and_mask_forward(
                     resized_all_points_384_outmask, 
-                    self.model_length - 1, 384, 384
+                    25 - 1, 384, 384
                 )
         else:
             input_drag_384_outmask, input_mask_384_outmask = \
-                np.zeros((self.model_length - 1, 384, 384, 2)), \
-                    np.zeros((self.model_length - 1, 384, 384))
+                np.zeros((25 - 1, 384, 384, 2)), \
+                    np.zeros((25 - 1, 384, 384))
 
         input_drag_384_inmask = torch.from_numpy(input_drag_384_inmask).unsqueeze(0).to(self.device)  # [1, 13, h, w, 2]
         input_mask_384_inmask = torch.from_numpy(input_mask_384_inmask).unsqueeze(0).to(self.device)  # [1, 13, h, w]
@@ -441,7 +440,7 @@ class Drag:
         input_first_frame = repeat(first_frames_transform(input_first_frame), 'c h w -> b c h w', b=1).to(self.device)
 
         seed = 42
-        num_frames = self.model_length
+        num_frames = 25
         
         set_seed(seed)
 
@@ -500,8 +499,8 @@ class Drag:
         new_resized_all_points = []
         new_resized_all_points_384 = []
         for tnum in range(len(resized_all_points)):
-            new_resized_all_points.append(interpolate_trajectory(input_all_points[tnum], self.model_length))
-            new_resized_all_points_384.append(interpolate_trajectory(resized_all_points_384[tnum], self.model_length))
+            new_resized_all_points.append(interpolate_trajectory(input_all_points[tnum], 25))
+            new_resized_all_points_384.append(interpolate_trajectory(resized_all_points_384[tnum], 25))
 
         resized_all_points = np.array(new_resized_all_points)
         resized_all_points_384 = np.array(new_resized_all_points_384)
@@ -519,24 +518,24 @@ class Drag:
             input_drag_384_inmask, input_mask_384_inmask = \
                 get_sparseflow_and_mask_forward(
                     resized_all_points_384_inmask, 
-                    self.model_length - 1, 384, 384
+                    25 - 1, 384, 384
                 )
         else:
             input_drag_384_inmask, input_mask_384_inmask = \
-                np.zeros((self.model_length - 1, 384, 384, 2)), \
-                    np.zeros((self.model_length - 1, 384, 384))
+                np.zeros((25 - 1, 384, 384, 2)), \
+                    np.zeros((25 - 1, 384, 384))
         
         if resized_all_points_384_outmask.shape[0] != 0:
             out_mask_flag = True
             input_drag_384_outmask, input_mask_384_outmask = \
                 get_sparseflow_and_mask_forward(
                     resized_all_points_384_outmask, 
-                    self.model_length - 1, 384, 384
+                    25 - 1, 384, 384
                 )
         else:
             input_drag_384_outmask, input_mask_384_outmask = \
-                np.zeros((self.model_length - 1, 384, 384, 2)), \
-                    np.zeros((self.model_length - 1, 384, 384))
+                np.zeros((25 - 1, 384, 384, 2)), \
+                    np.zeros((25 - 1, 384, 384))
 
         input_drag_384_inmask = torch.from_numpy(input_drag_384_inmask).unsqueeze(0)  # [1, 13, h, w, 2]
         input_mask_384_inmask = torch.from_numpy(input_mask_384_inmask).unsqueeze(0)  # [1, 13, h, w]
@@ -657,7 +656,7 @@ with gr.Blocks() as demo:
     )
 
     target_size = 512
-    DragNUWA_net = Drag("cuda:0", target_size, target_size, 25)
+    DragNUWA_net = Drag("cuda:0", target_size, target_size)
     first_frame_path = gr.State()
     tracking_points = gr.State([])
     motion_brush_points = gr.State([])