Update app.py

app.py

@@ -160,9 +160,10 @@ def clear_clicks(video_state):
 def load_uploaded_mask_video(video_path, mask_video_path, n_frames, video_state):
     """Load a user-supplied mask video (white = object) and use it directly,
     skipping click segmentation and SAM2 tracking."""
-    src = video_state["video_path"] or video_path
-    if mask_video_path is None or src is None:
-        return video_state, None
+    src = video_state.get("video_path") or video_path
+    if src is None or mask_video_path is None:
+        raise gr.Error("Upload both the source video and the mask video first.")
+    video_state["video_path"] = src   # <-- persist for the Remove step
 
     # --- load source frames (same logic as track_video) ---
     vr = VideoReader(src, ctx=cpu(0))
@@ -203,7 +204,13 @@ def load_uploaded_mask_video(video_path, mask_video_path, n_frames, video_state)
     return video_state, Image.fromarray(preview)
 
 def get_source_info(video_state):
-    vr = VideoReader(video_state["video_path"], ctx=cpu(0))
+    path = video_state.get("video_path")
+    if path is None:
+        raise gr.Error(
+            "No source video found. Upload a video, then click "
+            "'Extract First Frame' or 'Use Uploaded Mask' before Remove."
+        )
+    vr = VideoReader(path, ctx=cpu(0))
     fps = float(vr.get_avg_fps())
     h, w = vr[0].shape[:2]
     del vr
@@ -234,6 +241,8 @@ 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.get("video_path") is None:
+        video_state["video_path"] = video_path
     if video_state["origin_images"] is None or video_state["masks"] is None:
         return None
 
@@ -484,7 +493,7 @@ with gr.Blocks() as demo:
         remove_video = gr.Video(label="Remove Results", elem_id="my-video")
         remove_btn.click(
             inference_and_return_video,
-            inputs=[dilation_slider, inference_steps_slider, video_state],
+            inputs=[dilation_slider, inference_steps_slider, video_input, video_state],
             outputs=remove_video
         )