Update app.py

app.py

@@ -202,6 +202,13 @@ def load_uploaded_mask_video(video_path, mask_video_path, n_frames, video_state)
     preview = np.uint8(np.clip(painted * 255, 0, 255))
     return video_state, Image.fromarray(preview)
 
+def get_source_info(video_state):
+    vr = VideoReader(video_state["video_path"], ctx=cpu(0))
+    fps = float(vr.get_avg_fps())
+    h, w = vr[0].shape[:2]
+    del vr
+    return (fps if fps and fps > 0 else 15.0), (w, h)
+
 def preprocess_for_removal(images, masks):
     out_images = []
     out_masks = []
@@ -223,45 +230,69 @@ def preprocess_for_removal(images, masks):
     arr_masks = np.stack(out_masks)
     return torch.from_numpy(arr_images).half(), torch.from_numpy(arr_masks).half()
 
-@spaces.GPU(duration=50)
+MAX_PROC_SIDE = None  # None = native resolution; set e.g. 1280 if you hit OOM
+
+@spaces.GPU(duration=300)  # 1080p is heavy — give it time
 def inference_and_return_video(dilation_iterations, num_inference_steps, video_state):
     if video_state["origin_images"] is None or video_state["masks"] is None:
         return None
-    images = video_state["origin_images"]
+
     masks = video_state["masks"]
+    fps, (orig_w, orig_h) = get_source_info(video_state)  # helper from before
 
-    images = np.array(images)
-    masks = np.array(masks)
-    img_tensor, mask_tensor = preprocess_for_removal(images, masks)
-    img_tensor=img_tensor.to("cuda")
-    mask_tensor = mask_tensor[:,:,:,:1].to("cuda")
-
-    if mask_tensor.shape[1] < mask_tensor.shape[2]:
-        height = 480
-        width = 832
-    else:
-        height = 832
-        width = 480
+    # --- load full-res frames straight from the source video ---
+    vr = VideoReader(video_state["video_path"], ctx=cpu(0))
+    images = [vr[i].asnumpy() for i in range(min(len(masks), len(vr)))]
+    del vr
+    n_frames = len(images)
+
+    # --- processing size: native (optionally capped), rounded UP to a multiple of 16 ---
+    proc_w, proc_h = orig_w, orig_h
+    if MAX_PROC_SIDE and max(proc_w, proc_h) > MAX_PROC_SIDE:
+        s = MAX_PROC_SIDE / max(proc_w, proc_h)
+        proc_w, proc_h = round(proc_w * s), round(proc_h * s)
+    W16 = (proc_w + 15) // 16 * 16   # 1920 -> 1920
+    H16 = (proc_h + 15) // 16 * 16   # 1080 -> 1088
+    pad_w, pad_h = W16 - proc_w, H16 - proc_h
+
+    proc_images, proc_masks = [], []
+    for img, msk in zip(images, masks):
+        img = cv2.resize(img, (proc_w, proc_h), interpolation=cv2.INTER_LINEAR)
+        msk = cv2.resize(msk, (proc_w, proc_h), interpolation=cv2.INTER_NEAREST)
+        img = img.astype(np.float32) / 127.5 - 1.0          # [-1, 1]
+        msk = (msk > 0.5).astype(np.float32)
+        if msk.ndim == 2:
+            msk = msk[..., None]
+        if pad_h or pad_w:  # reflect-pad image, zero-pad mask (background)
+            img = cv2.copyMakeBorder(img, 0, pad_h, 0, pad_w, cv2.BORDER_REFLECT)
+            msk = cv2.copyMakeBorder(msk, 0, pad_h, 0, pad_w, cv2.BORDER_CONSTANT, value=0)
+        proc_images.append(img)
+        proc_masks.append(msk)
+
+    img_tensor = torch.from_numpy(np.stack(proc_images)).half().to("cuda")
+    mask_tensor = torch.from_numpy(np.stack(proc_masks)).half()[..., :1].to("cuda")
 
     pipe.to("cuda")
     with torch.no_grad():
         out = pipe(
-                images=img_tensor,
-                masks=mask_tensor,
-                num_frames=mask_tensor.shape[0],
-                height=height,
-                width=width,
-                num_inference_steps=int(num_inference_steps),
-                generator=torch.Generator(device=device).manual_seed(random_seed),
-                iterations=int(dilation_iterations)
+            images=img_tensor,
+            masks=mask_tensor,
+            num_frames=n_frames,
+            height=H16,
+            width=W16,
+            num_inference_steps=int(num_inference_steps),
+            generator=torch.Generator(device=device).manual_seed(random_seed),
+            iterations=int(dilation_iterations),
         ).frames[0]
 
-        out = np.uint8(out * 255)
-        output_frames = [img for img in out]
+    out = np.uint8(out * 255)
+    out = out[:, :proc_h, :proc_w]                       # strip the padding back off
+    if (proc_w, proc_h) != (orig_w, orig_h):             # only if MAX_PROC_SIDE kicked in
+        out = [cv2.resize(f, (orig_w, orig_h), interpolation=cv2.INTER_CUBIC) for f in out]
 
     video_file = f"/tmp/{time.time()}-{random.random()}-removed_output.mp4"
-    clip = ImageSequenceClip(output_frames, fps=15)
-    clip.write_videofile(video_file, codec='libx264', audio=False, verbose=False, logger=None)
+    clip = ImageSequenceClip(list(out), fps=fps)
+    clip.write_videofile(video_file, codec="libx264", audio=False, verbose=False, logger=None)
     return video_file
 
 @spaces.GPU(duration=40)
@@ -436,10 +467,19 @@ with gr.Blocks() as demo:
             track_btn = gr.Button("Tracking")
         video_output = gr.Video(label="Tracking Result", elem_id="my-video")
 
+        mask_video_input = gr.Video(label="(Optional) Upload Mask Video — white = object to remove")
+        use_mask_btn = gr.Button("Use Uploaded Mask (skips clicking & tracking)", elem_id="my-btn")
+        
+        use_mask_btn.click(
+            load_uploaded_mask_video,
+            inputs=[video_input, mask_video_input, n_frames_slider, video_state],
+            outputs=[video_state, image_output],
+        )
+
         with gr.Column(elem_id="my-btn"):
             dilation_slider = gr.Slider(minimum=1, maximum=20, value=6, step=1, label="Mask Dilation")
             inference_steps_slider = gr.Slider(minimum=1, maximum=100, value=6, step=1, label="Num Inference Steps")
-
+          
         remove_btn = gr.Button("Remove", elem_id="my-btn")
         remove_video = gr.Video(label="Remove Results", elem_id="my-video")
         remove_btn.click(
@@ -447,15 +487,6 @@ with gr.Blocks() as demo:
             inputs=[dilation_slider, inference_steps_slider, video_state],
             outputs=remove_video
         )
-
-        mask_video_input = gr.Video(label="(Optional) Upload Mask Video — white = object to remove")
-        use_mask_btn = gr.Button("Use Uploaded Mask (skips clicking & tracking)", elem_id="my-btn")
-        
-        use_mask_btn.click(
-            load_uploaded_mask_video,
-            inputs=[video_input, mask_video_input, n_frames_slider, video_state],
-            outputs=[video_state, image_output],
-        )
         
         get_info_btn.click(get_video_info, inputs=[video_input, video_state], \
                        outputs=image_output)