move model to cpu when not using ZeroGPU

app.py

@@ -70,26 +70,9 @@ examples = [
 
 OBJ_ID = 0
 
-
-@spaces.GPU
-def get_predictor(session_state):
-    if "predictor" not in session_state:
-        sam2_checkpoint = "checkpoints/edgetam.pt"
-        model_cfg = "edgetam.yaml"
-        predictor = build_sam2_video_predictor(
-            model_cfg, sam2_checkpoint, device="cuda"
-        )
-        print("predictor loaded")
-
-        # use bfloat16 for the entire demo
-        torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()
-        if torch.cuda.get_device_properties(0).major >= 8:
-            # turn on tfloat32 for Ampere GPUs (https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices)
-            torch.backends.cuda.matmul.allow_tf32 = True
-            torch.backends.cudnn.allow_tf32 = True
-
-        session_state["predictor"] = predictor
-    return session_state["predictor"]
+sam2_checkpoint = "checkpoints/edgetam.pt"
+model_cfg = "edgetam.yaml"
+predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint, device="cpu")
 
 
 def get_video_fps(video_path):
@@ -106,10 +89,8 @@ def get_video_fps(video_path):
     return fps
 
 
-@spaces.GPU
 def reset(session_state):
-    predictor = get_predictor(session_state)
-    predictor.to("cuda")
+    predictor.to("cpu")
     session_state["input_points"] = []
     session_state["input_labels"] = []
     if session_state["inference_state"] is not None:
@@ -127,10 +108,8 @@ def reset(session_state):
     )
 
 
-@spaces.GPU
 def clear_points(session_state):
-    predictor = get_predictor(session_state)
-    predictor.to("cuda")
+    predictor.to("cpu")
     session_state["input_points"] = []
     session_state["input_labels"] = []
     if session_state["inference_state"]["tracking_has_started"]:
@@ -143,10 +122,8 @@ def clear_points(session_state):
     )
 
 
-@spaces.GPU
 def preprocess_video_in(video_path, session_state):
-    predictor = get_predictor(session_state)
-    predictor.to("cuda")
+    predictor.to("cpu")
     if video_path is None:
         return (
             gr.update(open=True),  # video_in_drawer
@@ -210,59 +187,62 @@ def segment_with_points(
     session_state,
     evt: gr.SelectData,
 ):
-    predictor = get_predictor(session_state)
-    predictor.to("cuda")
-    session_state["input_points"].append(evt.index)
-    print(f"TRACKING INPUT POINT: {session_state['input_points']}")
-
-    if point_type == "include":
-        session_state["input_labels"].append(1)
-    elif point_type == "exclude":
-        session_state["input_labels"].append(0)
-    print(f"TRACKING INPUT LABEL: {session_state['input_labels']}")
-
-    # Open the image and get its dimensions
-    transparent_background = Image.fromarray(session_state["first_frame"]).convert(
-        "RGBA"
-    )
-    w, h = transparent_background.size
+    if torch.cuda.get_device_properties(0).major >= 8:
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+        predictor.to("cuda")
+        session_state["input_points"].append(evt.index)
+        print(f"TRACKING INPUT POINT: {session_state['input_points']}")
+
+        if point_type == "include":
+            session_state["input_labels"].append(1)
+        elif point_type == "exclude":
+            session_state["input_labels"].append(0)
+        print(f"TRACKING INPUT LABEL: {session_state['input_labels']}")
+
+        # Open the image and get its dimensions
+        transparent_background = Image.fromarray(session_state["first_frame"]).convert(
+            "RGBA"
+        )
+        w, h = transparent_background.size
 
-    # Define the circle radius as a fraction of the smaller dimension
-    fraction = 0.01  # You can adjust this value as needed
-    radius = int(fraction * min(w, h))
+        # Define the circle radius as a fraction of the smaller dimension
+        fraction = 0.01  # You can adjust this value as needed
+        radius = int(fraction * min(w, h))
 
-    # Create a transparent layer to draw on
-    transparent_layer = np.zeros((h, w, 4), dtype=np.uint8)
+        # Create a transparent layer to draw on
+        transparent_layer = np.zeros((h, w, 4), dtype=np.uint8)
 
-    for index, track in enumerate(session_state["input_points"]):
-        if session_state["input_labels"][index] == 1:
-            cv2.circle(transparent_layer, track, radius, (0, 255, 0, 255), -1)
-        else:
-            cv2.circle(transparent_layer, track, radius, (255, 0, 0, 255), -1)
+        for index, track in enumerate(session_state["input_points"]):
+            if session_state["input_labels"][index] == 1:
+                cv2.circle(transparent_layer, track, radius, (0, 255, 0, 255), -1)
+            else:
+                cv2.circle(transparent_layer, track, radius, (255, 0, 0, 255), -1)
 
-    # Convert the transparent layer back to an image
-    transparent_layer = Image.fromarray(transparent_layer, "RGBA")
-    selected_point_map = Image.alpha_composite(
-        transparent_background, transparent_layer
-    )
+        # Convert the transparent layer back to an image
+        transparent_layer = Image.fromarray(transparent_layer, "RGBA")
+        selected_point_map = Image.alpha_composite(
+            transparent_background, transparent_layer
+        )
 
-    # Let's add a positive click at (x, y) = (210, 350) to get started
-    points = np.array(session_state["input_points"], dtype=np.float32)
-    # for labels, `1` means positive click and `0` means negative click
-    labels = np.array(session_state["input_labels"], np.int32)
-    _, _, out_mask_logits = predictor.add_new_points(
-        inference_state=session_state["inference_state"],
-        frame_idx=0,
-        obj_id=OBJ_ID,
-        points=points,
-        labels=labels,
-    )
+        # Let's add a positive click at (x, y) = (210, 350) to get started
+        points = np.array(session_state["input_points"], dtype=np.float32)
+        # for labels, `1` means positive click and `0` means negative click
+        labels = np.array(session_state["input_labels"], np.int32)
+        _, _, out_mask_logits = predictor.add_new_points(
+            inference_state=session_state["inference_state"],
+            frame_idx=0,
+            obj_id=OBJ_ID,
+            points=points,
+            labels=labels,
+        )
 
-    mask_image = show_mask((out_mask_logits[0] > 0.0).cpu().numpy())
-    first_frame_output = Image.alpha_composite(transparent_background, mask_image)
+        mask_image = show_mask((out_mask_logits[0] > 0.0).cpu().numpy())
+        first_frame_output = Image.alpha_composite(transparent_background, mask_image)
 
-    torch.cuda.empty_cache()
-    return selected_point_map, first_frame_output, session_state
+        torch.cuda.empty_cache()
+        return selected_point_map, first_frame_output, session_state
 
 
 def show_mask(mask, obj_id=None, random_color=False, convert_to_image=True):
@@ -285,61 +265,68 @@ def propagate_to_all(
     video_in,
     session_state,
 ):
-    predictor = get_predictor(session_state)
     predictor.to("cuda")
-    if (
-        len(session_state["input_points"]) == 0
-        or video_in is None
-        or session_state["inference_state"] is None
-    ):
-        return (
-            None,
-            session_state,
+    if torch.cuda.get_device_properties(0).major >= 8:
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+        if (
+            len(session_state["input_points"]) == 0
+            or video_in is None
+            or session_state["inference_state"] is None
+        ):
+            return (
+                None,
+                session_state,
+            )
+
+        # run propagation throughout the video and collect the results in a dict
+        video_segments = (
+            {}
+        )  # video_segments contains the per-frame segmentation results
+        print("starting propagate_in_video")
+        for out_frame_idx, out_obj_ids, out_mask_logits in predictor.propagate_in_video(
+            session_state["inference_state"]
+        ):
+            video_segments[out_frame_idx] = {
+                out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()
+                for i, out_obj_id in enumerate(out_obj_ids)
+            }
+
+        # obtain the segmentation results every few frames
+        vis_frame_stride = 1
+
+        output_frames = []
+        for out_frame_idx in range(0, len(video_segments), vis_frame_stride):
+            transparent_background = Image.fromarray(
+                session_state["all_frames"][out_frame_idx]
+            ).convert("RGBA")
+            out_mask = video_segments[out_frame_idx][OBJ_ID]
+            mask_image = show_mask(out_mask)
+            output_frame = Image.alpha_composite(transparent_background, mask_image)
+            output_frame = np.array(output_frame)
+            output_frames.append(output_frame)
+
+        torch.cuda.empty_cache()
+
+        # Create a video clip from the image sequence
+        original_fps = get_video_fps(video_in)
+        fps = original_fps  # Frames per second
+        clip = ImageSequenceClip(output_frames, fps=fps)
+        # Write the result to a file
+        unique_id = datetime.now().strftime("%Y%m%d%H%M%S")
+        final_vid_output_path = f"output_video_{unique_id}.mp4"
+        final_vid_output_path = os.path.join(
+            tempfile.gettempdir(), final_vid_output_path
         )
 
-    # run propagation throughout the video and collect the results in a dict
-    video_segments = {}  # video_segments contains the per-frame segmentation results
-    print("starting propagate_in_video")
-    for out_frame_idx, out_obj_ids, out_mask_logits in predictor.propagate_in_video(
-        session_state["inference_state"]
-    ):
-        video_segments[out_frame_idx] = {
-            out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()
-            for i, out_obj_id in enumerate(out_obj_ids)
-        }
-
-    # obtain the segmentation results every few frames
-    vis_frame_stride = 1
-
-    output_frames = []
-    for out_frame_idx in range(0, len(video_segments), vis_frame_stride):
-        transparent_background = Image.fromarray(
-            session_state["all_frames"][out_frame_idx]
-        ).convert("RGBA")
-        out_mask = video_segments[out_frame_idx][OBJ_ID]
-        mask_image = show_mask(out_mask)
-        output_frame = Image.alpha_composite(transparent_background, mask_image)
-        output_frame = np.array(output_frame)
-        output_frames.append(output_frame)
-
-    torch.cuda.empty_cache()
-
-    # Create a video clip from the image sequence
-    original_fps = get_video_fps(video_in)
-    fps = original_fps  # Frames per second
-    clip = ImageSequenceClip(output_frames, fps=fps)
-    # Write the result to a file
-    unique_id = datetime.now().strftime("%Y%m%d%H%M%S")
-    final_vid_output_path = f"output_video_{unique_id}.mp4"
-    final_vid_output_path = os.path.join(tempfile.gettempdir(), final_vid_output_path)
-
-    # Write the result to a file
-    clip.write_videofile(final_vid_output_path, codec="libx264")
+        # Write the result to a file
+        clip.write_videofile(final_vid_output_path, codec="libx264")
 
-    return (
-        gr.update(value=final_vid_output_path),
-        session_state,
-    )
+        return (
+            gr.update(value=final_vid_output_path),
+            session_state,
+        )
 
 
 def update_ui():