Improve ZeroGPU compatibility

app.py

@@ -18,23 +18,46 @@ import gradio as gr
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 yaml_args = OmegaConf.load(f"{REPO_ROOT}/ckpt/model_config.yaml")
-pipeline = None
 
+if not os.path.exists(f"{REPO_ROOT}/ckpt/model.safetensors"):
+    subprocess.run(["bash", f"{REPO_ROOT}/infer_bash/download_ckpt.sh"], check=True)
+pipeline = load_model(f"{REPO_ROOT}/ckpt", yaml_args)
 
-@spaces.GPU
+
+MAX_FRAMES = 300
+
+
+def read_video_limited(video_path, max_frames):
+    """Read up to max_frames from a video without loading the entire file."""
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        raise gr.Error(f"Cannot open video: {video_path}")
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    if total_frames > max_frames:
+        gr.Warning(
+            f"Video has {total_frames} frames, processing only the first {max_frames}."
+        )
+    frames = []
+    while len(frames) < max_frames:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+    cap.release()
+    video_np = np.stack(frames)
+    video_tensor = (
+        torch.from_numpy(video_np).permute(0, 3, 1, 2).float() / 255.0
+    )
+    return video_tensor.unsqueeze(0), fps
+
+
+@spaces.GPU(duration=90)
 def fn(input_video):
-    global pipeline, yaml_args, device
-    if pipeline is None:
-        if not os.path.exists(f"{REPO_ROOT}/ckpt/model.safetensors"):
-            subprocess.run(["bash", f"{REPO_ROOT}/infer_bash/download_ckpt.sh"], check=True)
-        pipeline = load_model(f"{REPO_ROOT}/ckpt", yaml_args)
-    
-    input_video_basename = os.path.basename(input_video)
-    input_tensor, orig_size, origin_fps = load_video_data(Namespace(
-        input_video=input_video,
-        height=480,
-        width=640,
-    ))
+    input_tensor, origin_fps = read_video_limited(input_video, MAX_FRAMES)
+    input_tensor, orig_size = resize_for_training_scale(
+        input_tensor, 480, 640
+    )
     depth = predict_depth(pipeline, input_tensor, orig_size, Namespace(
         window_size=81,
         overlap=21