Update app.py

app.py

@@ -188,8 +188,6 @@ def inference_and_return_video(dilation_iterations, num_inference_steps, video_s
 
     images = np.array(images)
     masks = np.array(masks)
-    print(f"line 191 images shape:{images.shape},masks shape:{masks.shape}")
-    #line 191 images shape:(1, 1024, 1820, 3),masks shape:(1, 1024, 1820), which should be (16, 1024, 1820, 3) and (16, 1024, 1820, 3)
     img_tensor, mask_tensor = preprocess_for_removal(images, masks)
     mask_tensor = mask_tensor[:,:,:,:1]
 
@@ -200,6 +198,7 @@ def inference_and_return_video(dilation_iterations, num_inference_steps, video_s
         height = 832
         width = 480
 
+    pipe=pipe.to("cuda")
     with torch.no_grad():
         out = pipe(
                 images=img_tensor,
@@ -220,7 +219,7 @@ def inference_and_return_video(dilation_iterations, num_inference_steps, video_s
     clip.write_videofile(video_file, codec='libx264', audio=False, verbose=False, logger=None)
     return video_file
 
-@spaces.GPU(duration=60)
+@spaces.GPU(duration=200)
 def track_video(n_frames,video_state):
     input_points = video_state["input_points"]
     input_labels = video_state["input_labels"]
@@ -271,16 +270,11 @@ def track_video(n_frames,video_state):
             mask += out_mask
         mask = np.clip(mask, 0, 1)
         mask = cv2.resize(mask, (W_, H_))
-        print(f"line 275 mask shape:{mask.shape}")
         mask_frames.append(mask)
-        print(f"line 277 len(mask_frames)={len(mask_frames)}")
         painted = (1 - mask * 0.5) * frame + mask * 0.5 * color
         painted = np.uint8(np.clip(painted * 255, 0, 255))
         output_frames.append(painted)
-        print(f"line 281 len(output_frames)={len(output_frames)}, painted shape:{painted.shape}")
     video_state["masks"] =mask_frames
-    print(f'line 283 len video_state["masks"]:{len(video_state["masks"])}')
-    print(f'line 284 video_state["masks"][0].shape:{video_state["masks"][0].shape}')
     video_file = f"/tmp/{time.time()}-{random.random()}-tracked_output.mp4"
     clip = ImageSequenceClip(output_frames, fps=15)
     clip.write_videofile(video_file, codec='libx264', audio=False, verbose=False, logger=None)