Port to ZeroGPU: lazy model loading and @spaces.GPU decorators

README.md

@@ -4,7 +4,8 @@ emoji: 🤡
 colorFrom: red
 colorTo: green
 sdk: gradio
-sdk_version: 5.16.0
+sdk_version: 4.31.0
+python_version: 3.10.13
 app_file: hugging_face/app.py
 pinned: false
 license: other

hugging_face/app.py

@@ -15,10 +15,11 @@ import cv2
 import torch
 import numpy as np
 import gradio as gr
+import spaces
  
 from tools.painter import mask_painter
 from tools.interact_tools import SamControler
-from tools.misc import get_device
+# get_device is NOT imported at module level to avoid CUDA init via torch.cuda.is_available()
 from tools.download_util import load_file_from_url
 
 from matanyone2_wrapper import matanyone2
@@ -37,10 +38,13 @@ def parse_augment():
     parser.add_argument('--mask_save', default=False)
     args = parser.parse_args()
     
+    # ZeroGPU: do NOT call get_device() (which calls torch.cuda.is_available()) at module level.
+    # It can trigger CUDA init in the main process. Default to 'cpu'; GPU functions
+    # determine the actual device at runtime inside @spaces.GPU-decorated functions.
     if not args.device:
-        args.device = str(get_device())
+        args.device = "cpu"
 
-    return args 
+    return args
 
 # SAM generator
 class MaskGenerator():
@@ -94,8 +98,7 @@ def get_frames_from_image(image_input, image_state):
         "fps": None
         }
     image_info = "Image Name: N/A,\nFPS: N/A,\nTotal Frames: {},\nImage Size:{}".format(len(frames), image_size)
-    model.samcontroler.sam_controler.reset_image() 
-    model.samcontroler.sam_controler.set_image(image_state["origin_images"][0])
+    # SAM loading and set_image are deferred to sam_refine() which runs under @spaces.GPU
     return image_state, image_info, image_state["origin_images"][0], \
                         gr.update(visible=True, maximum=10, value=10), gr.update(visible=False, maximum=len(frames), value=len(frames)), \
                         gr.update(visible=True), gr.update(visible=True), \
@@ -166,8 +169,7 @@ def get_frames_from_video(video_input, video_state):
         "audio": audio_path
         }
     video_info = "Video Name: {},\nFPS: {},\nTotal Frames: {},\nImage Size:{}".format(video_state["video_name"], round(video_state["fps"], 0), len(frames), image_size)
-    model.samcontroler.sam_controler.reset_image() 
-    model.samcontroler.sam_controler.set_image(video_state["origin_images"][0])
+    # SAM loading and set_image are deferred to sam_refine() which runs under @spaces.GPU
     return video_state, video_info, video_state["origin_images"][0], gr.update(visible=True, maximum=len(frames), value=1), gr.update(visible=False, maximum=len(frames), value=len(frames)), \
                         gr.update(visible=True), gr.update(visible=True), \
                         gr.update(visible=True), gr.update(visible=True),\
@@ -181,22 +183,14 @@ def select_video_template(image_selection_slider, video_state, interactive_state
 
     image_selection_slider -= 1
     video_state["select_frame_number"] = image_selection_slider
-
-    # once select a new template frame, set the image in sam
-    model.samcontroler.sam_controler.reset_image()
-    model.samcontroler.sam_controler.set_image(video_state["origin_images"][image_selection_slider])
-
+    # SAM set_image is deferred to sam_refine() which runs under @spaces.GPU
     return video_state["painted_images"][image_selection_slider], video_state, interactive_state
 
 def select_image_template(image_selection_slider, video_state, interactive_state):
 
     image_selection_slider = 0 # fixed for image
     video_state["select_frame_number"] = image_selection_slider
-
-    # once select a new template frame, set the image in sam
-    model.samcontroler.sam_controler.reset_image()
-    model.samcontroler.sam_controler.set_image(video_state["origin_images"][image_selection_slider])
-
+    # SAM set_image is deferred to sam_refine() which runs under @spaces.GPU
     return video_state["painted_images"][image_selection_slider], video_state, interactive_state
 
 # set the tracking end frame
@@ -206,6 +200,7 @@ def get_end_number(track_pause_number_slider, video_state, interactive_state):
     return video_state["painted_images"][track_pause_number_slider],interactive_state
 
 # use sam to get the mask
+@spaces.GPU(duration=60)
 def sam_refine(video_state, point_prompt, click_state, interactive_state, evt:gr.SelectData):
     """
     Args:
@@ -221,6 +216,7 @@ def sam_refine(video_state, point_prompt, click_state, interactive_state, evt:gr
         interactive_state["negative_click_times"] += 1
     
     # prompt for sam model
+    ensure_sam_on_cuda()
     model.samcontroler.sam_controler.reset_image()
     model.samcontroler.sam_controler.set_image(video_state["origin_images"][video_state["select_frame_number"]])
     prompt = get_prompt(click_state=click_state, click_input=coordinate)
@@ -269,6 +265,7 @@ def show_mask(video_state, interactive_state, mask_dropdown):
         return select_frame
 
 # image matting
+@spaces.GPU(duration=120)
 def image_matting(video_state, interactive_state, mask_dropdown, erode_kernel_size, dilate_kernel_size, refine_iter, model_selection):
     # Load model if not already loaded
     try:
@@ -308,6 +305,7 @@ def image_matting(video_state, interactive_state, mask_dropdown, erode_kernel_si
     return foreground_output, alpha_output
 
 # video matting
+@spaces.GPU(duration=300)
 def video_matting(video_state, interactive_state, mask_dropdown, erode_kernel_size, dilate_kernel_size, model_selection):
     # Load model if not already loaded
     try:
@@ -368,34 +366,40 @@ def add_audio_to_video(video_path, audio_path, output_path):
 
 
 def generate_video_from_frames(frames, output_path, fps=30, gray2rgb=False, audio_path=""):
-    """
-    Generates a video from a list of frames.
-    
-    Args:
-        frames (list of numpy arrays): The frames to include in the video.
-        output_path (str): The path to save the generated video.
-        fps (int, optional): The frame rate of the output video. Defaults to 30.
-    """
-    frames = torch.from_numpy(np.asarray(frames))
-    _, h, w, _ = frames.shape
+    frames = np.asarray(frames)
+
     if gray2rgb:
         frames = np.repeat(frames, 3, axis=3)
 
+    _, h, w, _ = frames.shape
+    h = h // 2 * 2
+    w = w // 2 * 2
+
+    if frames.shape[1] != h or frames.shape[2] != w:
+        frames = np.asarray([
+            cv2.resize(frame, (w, h), interpolation=cv2.INTER_LINEAR)
+            for frame in frames
+        ])
+
     if not os.path.exists(os.path.dirname(output_path)):
         os.makedirs(os.path.dirname(output_path))
+
     video_temp_path = output_path.replace(".mp4", "_temp.mp4")
-    
-    # resize back to ensure input resolution
-    imageio.mimwrite(video_temp_path, frames, fps=fps, quality=7, 
-                     codec='libx264', ffmpeg_params=["-vf", f"scale={w}:{h}"])
-    
-    # add audio to video if audio path exists
+
+    imageio.mimwrite(
+        video_temp_path,
+        frames,
+        fps=fps,
+        quality=7,
+        codec="libx264",
+        macro_block_size=1
+    )
+
     if audio_path != "" and os.path.exists(audio_path):
-        output_path = add_audio_to_video(video_temp_path, audio_path, output_path)    
+        output_path = add_audio_to_video(video_temp_path, audio_path, output_path)
         os.remove(video_temp_path)
         return output_path
-    else:
-        return video_temp_path
+    return video_temp_path
 
 # reset all states for a new input
 def restart():
@@ -434,11 +438,10 @@ sam_checkpoint_url_dict = {
 }
 checkpoint_folder = os.path.join('/home/user/app/', 'pretrained_models')
 
-sam_checkpoint = load_file_from_url(sam_checkpoint_url_dict[args.sam_model_type], checkpoint_folder)
-# initialize sams
-model = MaskGenerator(sam_checkpoint, args)
+# ZeroGPU: do NOT download or load models at module level.
+# All model loading is deferred to the first GPU function call.
+model = None  # SAM MaskGenerator — lazily initialized
 
-# initialize matanyone - lazy loading
 # Model display names to file names mapping
 model_display_to_file = {
     "MatAnyone": "matanyone.pth",
@@ -451,71 +454,79 @@ model_urls = {
     "matanyone2.pth": "https://github.com/pq-yang/MatAnyone2/releases/download/v1.0.0/matanyone2.pth"
 }
 
-# Model paths - download models using load_file_from_url
-model_paths = {
-    "matanyone.pth": load_file_from_url(model_urls["matanyone.pth"], checkpoint_folder),
-    "matanyone2.pth": load_file_from_url(model_urls["matanyone2.pth"], checkpoint_folder)
-}
+# MatAnyone model file paths — filled lazily on first download
+model_paths = {}
 
-# Cache for loaded models (lazy loading)
+# Cache for loaded MatAnyone models (lazy loading)
 loaded_models = {}
 
+# All supported models (for the UI) — always show both options
+available_models = ["MatAnyone 2", "MatAnyone"]
+default_model = "MatAnyone 2"
+
+
+def ensure_sam_loaded():
+    """Download SAM checkpoint and init MaskGenerator on CPU (safe to call outside GPU context)."""
+    global model
+    if model is None:
+        sam_checkpoint = load_file_from_url(sam_checkpoint_url_dict[args.sam_model_type], checkpoint_folder)
+        # Always load on CPU here — CUDA placement happens in ensure_sam_on_cuda(),
+        # which is only ever called from within a @spaces.GPU-decorated function.
+        import copy
+        cpu_args = copy.copy(args)
+        cpu_args.device = "cpu"
+        model = MaskGenerator(sam_checkpoint, cpu_args)
+
+
+def ensure_sam_on_cuda():
+    """Move SAM to CUDA. Must only be called inside a @spaces.GPU-decorated function."""
+    ensure_sam_loaded()
+    cuda_device = "cuda" if torch.cuda.is_available() else "cpu"
+    model.samcontroler.sam_controler.predictor.model.to(cuda_device)
+    model.samcontroler.sam_controler.device = cuda_device
+    model.samcontroler.sam_controler.torch_dtype = torch.float16 if cuda_device == "cuda" else torch.float32
+
+
+def _ensure_matanyone_downloaded(model_file):
+    """Download the MatAnyone checkpoint if not already present."""
+    if model_file not in model_paths:
+        model_paths[model_file] = load_file_from_url(model_urls[model_file], checkpoint_folder)
+    return model_paths[model_file]
+
+
 def load_model(display_name):
-    """Load a model if not already loaded"""
-    # Convert display name to file name
+    """Download (if needed) and load a MatAnyone model. Cached after first load."""
+    # Map display name to file name
     if display_name in model_display_to_file:
         model_file = model_display_to_file[display_name]
-    elif display_name in model_paths:
-        # Also support direct file name for backward compatibility
+    elif display_name in model_urls:
         model_file = display_name
     else:
         raise ValueError(f"Unknown model: {display_name}")
-    
+
     if model_file in loaded_models:
         return loaded_models[model_file]
-    
-    if model_file not in model_paths:
-        raise ValueError(f"Unknown model file: {model_file}")
-    
-    ckpt_path = model_paths[model_file]
-    if not os.path.exists(ckpt_path):
-        raise FileNotFoundError(f"Model file not found: {ckpt_path}")
-    
+
+    ckpt_path = _ensure_matanyone_downloaded(model_file)
+
     # Clear Hydra instance if already initialized (to allow loading different models)
     try:
         GlobalHydra.instance().clear()
-    except:
-        pass  # If Hydra is not initialized, this is fine
-    
-    print(f"Loading model: {display_name} ({model_file})...")
-    model = get_matanyone2_model(ckpt_path, args.device)
-    model = model.to(args.device).eval()
-    loaded_models[model_file] = model
+    except Exception:
+        pass
+
+    device = "cuda" if torch.cuda.is_available() else args.device
+    print(f"Loading model: {display_name} ({model_file}) on {device}...")
+    loaded_mat_model = get_matanyone2_model(ckpt_path, device)
+    loaded_mat_model = loaded_mat_model.to(device).eval()
+    loaded_models[model_file] = loaded_mat_model
     print(f"Model {display_name} loaded successfully.")
-    return model
-
-# Get available model choices for the UI (check if files exist)
-# Order: MatAnyone 2 first, then MatAnyone
-available_models = []
-# Check MatAnyone 2 first
-if "MatAnyone 2" in model_display_to_file:
-    file_name = model_display_to_file["MatAnyone 2"]
-    if file_name in model_paths and os.path.exists(model_paths[file_name]):
-        available_models.append("MatAnyone 2")
-# Then check MatAnyone
-if "MatAnyone" in model_display_to_file:
-    file_name = model_display_to_file["MatAnyone"]
-    if file_name in model_paths and os.path.exists(model_paths[file_name]):
-        available_models.append("MatAnyone")
-
-if not available_models:
-    raise RuntimeError("No models are available! Please ensure at least one model file exists in ../pretrained_models/")
-default_model = "MatAnyone 2" if "MatAnyone 2" in available_models else available_models[0]
+    return loaded_mat_model
 
 # download test samples
 test_sample_path = os.path.join('/home/user/app/hugging_face/', "test_sample/")
-load_file_from_url('https://github.com/pq-yang/MatAnyone2/releases/download/media/test-sample-0-1080p.mp4', test_sample_path)
-load_file_from_url('https://github.com/pq-yang/MatAnyone2/releases/download/media/test-sample-1-1080p.mp4', test_sample_path)
+load_file_from_url('https://github.com/pq-yang/MatAnyone2/releases/download/media/test-sample-0-720p.mp4', test_sample_path)
+load_file_from_url('https://github.com/pq-yang/MatAnyone2/releases/download/media/test-sample-1-720p.mp4', test_sample_path)
 load_file_from_url('https://github.com/pq-yang/MatAnyone2/releases/download/media/test-sample-2-720p.mp4', test_sample_path)
 load_file_from_url('https://github.com/pq-yang/MatAnyone2/releases/download/media/test-sample-3-720p.mp4', test_sample_path)
 load_file_from_url('https://github.com/pq-yang/MatAnyone2/releases/download/media/test-sample-4-720p.mp4', test_sample_path)
@@ -562,7 +573,7 @@ If our work is useful for your research, please consider citing:
 @InProceedings{yang2025matanyone,
      title     = {{MatAnyone}: Stable Video Matting with Consistent Memory Propagation},
      author    = {Yang, Peiqing and Zhou, Shangchen and Zhao, Jixin and Tao, Qingyi and Loy, Chen Change},
-     booktitle = {arXiv preprint arXiv:2501.14677},
+     booktitle = {CVPR},
      year      = {2025}
 }
 ```
@@ -632,7 +643,7 @@ body {
     text-align: center;
     padding: 0;
     margin: 0;
-    height: 5vh;
+    height: 2vh;
     width: 80vw;
     font-family: "Sarpanch", sans-serif;
     font-weight: 60;
@@ -878,7 +889,7 @@ with gr.Blocks(theme=gr.themes.Monochrome(), css=my_custom_css) as demo:
             gr.Markdown("---")
             gr.Markdown("## Examples")
             gr.Examples(
-                examples=[os.path.join(os.path.dirname(__file__), "./test_sample/", test_sample) for test_sample in ["test-sample-0-1080p.mp4", "test-sample-1-1080p.mp4", "test-sample-2-720p.mp4", "test-sample-3-720p.mp4", "test-sample-4-720p.mp4", "test-sample-5-720p.mp4"]],
+                examples=[os.path.join(os.path.dirname(__file__), "./test_sample/", test_sample) for test_sample in ["test-sample-0-720p.mp4", "test-sample-1-720p.mp4", "test-sample-2-720p.mp4", "test-sample-3-720p.mp4", "test-sample-4-720p.mp4", "test-sample-5-720p.mp4"]],
                 inputs=[video_input],
             )
 

hugging_face/matanyone2_wrapper.py

@@ -7,8 +7,6 @@ import random
 import cv2
 from matanyone2.utils.device import get_default_device, safe_autocast_decorator
 
-device = get_default_device()
-
 def gen_dilate(alpha, min_kernel_size, max_kernel_size): 
     kernel_size = random.randint(min_kernel_size, max_kernel_size)
     kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size,kernel_size))
@@ -34,6 +32,7 @@ def matanyone2(processor, frames_np, mask, r_erode=0, r_dilate=0, n_warmup=10):
         com: [(H,W,C)]*n, uint8
         pha: [(H,W,C)]*n, uint8
     """
+    device = get_default_device()
 
     # print(f'===== [r_erode] {r_erode}; [r_dilate] {r_dilate} =====')
     bgr = (np.array([120, 255, 155], dtype=np.float32)/255).reshape((1, 1, 3))

requirements.txt

@@ -1,3 +1,4 @@
+spaces
 progressbar2
 gdown >= 4.7.1
 gitpython >= 3.1
@@ -6,7 +7,7 @@ hickle >= 5.0
 tensorboard >= 2.11
 numpy >= 1.21
 git+https://github.com/facebookresearch/segment-anything.git
-gradio==4.31.0
+# gradio==4.31.0
 fastapi==0.111.0
 pydantic==2.7.1
 opencv-python >= 4.8
@@ -33,4 +34,5 @@ pyqtdarktheme
 imageio == 2.25.0
 imageio[ffmpeg]
 ffmpeg-python
-safetensors
+safetensors
+huggingface_hub < 1.0