safetensors loading, Phase 0 4x faster (uint8), total time in status

- Prefer .safetensors over .pth (faster load, no pickle)
- Phase 0: save uint8 RGB (25MB/frame) instead of float32 (99MB/frame) — 4x less disk I/O
- Status now shows total wall time

app.py

@@ -68,8 +68,9 @@ CORRIDORKEY_MODELS = {
     "1024": os.path.join(MODELS_DIR, "corridorkey_1024.onnx"),
     "2048": os.path.join(MODELS_DIR, "corridorkey_2048.onnx"),
 }
-CORRIDORKEY_PTH_REPO = "nikopueringer/CorridorKey_v1.0"
-CORRIDORKEY_PTH_FILE = "CorridorKey_v1.0.pth"
+CORRIDORKEY_REPO = "nikopueringer/CorridorKey_v1.0"
+CORRIDORKEY_SAFETENSORS = "CorridorKey_v1.0.safetensors"
+CORRIDORKEY_PTH_FALLBACK = "CorridorKey_v1.0.pth"
 IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32).reshape(1, 1, 3)
 IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32).reshape(1, 1, 3)
 MAX_DURATION_CPU = 5
@@ -82,17 +83,21 @@ HAS_CUDA = "CUDAExecutionProvider" in ort.get_available_providers()
 # ---------------------------------------------------------------------------
 logger.info("Preloading model files at startup...")
 _preloaded_birefnet_path = None
-_preloaded_pth_path = None
+_preloaded_ckpt_path = None
 try:
     _preloaded_birefnet_path = hf_hub_download(repo_id=BIREFNET_REPO, filename=BIREFNET_FILE)
     logger.info("BiRefNet cached: %s", _preloaded_birefnet_path)
 except Exception as e:
     logger.warning("BiRefNet preload failed (will retry later): %s", e)
 try:
-    _preloaded_pth_path = hf_hub_download(repo_id=CORRIDORKEY_PTH_REPO, filename=CORRIDORKEY_PTH_FILE)
-    logger.info("CorridorKey.pth cached: %s", _preloaded_pth_path)
-except Exception as e:
-    logger.warning("CorridorKey.pth preload failed (will retry later): %s", e)
+    _preloaded_ckpt_path = hf_hub_download(repo_id=CORRIDORKEY_REPO, filename=CORRIDORKEY_SAFETENSORS)
+    logger.info("CorridorKey.safetensors cached: %s", _preloaded_ckpt_path)
+except Exception:
+    try:
+        _preloaded_ckpt_path = hf_hub_download(repo_id=CORRIDORKEY_REPO, filename=CORRIDORKEY_PTH_FALLBACK)
+        logger.info("CorridorKey.pth fallback cached: %s", _preloaded_ckpt_path)
+    except Exception as e:
+        logger.warning("CorridorKey preload failed (will retry later): %s", e)
 
 # Batch sizes for GPU inference (conservative for H200 80GB)
 GPU_BATCH_SIZES = {"1024": 32, "2048": 16}  # 2048 uses only 5.7GB/batch=2, so 16 easily fits in 69.8GB
@@ -306,7 +311,13 @@ def _load_greenformer(img_size):
     import torch.nn.functional as F
     from CorridorKeyModule.core.model_transformer import GreenFormer
 
-    checkpoint_path = _preloaded_pth_path or hf_hub_download(repo_id=CORRIDORKEY_PTH_REPO, filename=CORRIDORKEY_PTH_FILE)
+    if _preloaded_ckpt_path:
+        checkpoint_path = _preloaded_ckpt_path
+    else:
+        try:
+            checkpoint_path = hf_hub_download(repo_id=CORRIDORKEY_REPO, filename=CORRIDORKEY_SAFETENSORS)
+        except Exception:
+            checkpoint_path = hf_hub_download(repo_id=CORRIDORKEY_REPO, filename=CORRIDORKEY_PTH_FALLBACK)
     logger.info("Using checkpoint: %s", checkpoint_path)
 
     logger.info("Initializing GreenFormer (img_size=%d)...", img_size)
@@ -316,9 +327,13 @@ def _load_greenformer(img_size):
         use_refiner=True,
     )
 
-    # Load weights
-    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
-    state_dict = checkpoint.get("state_dict", checkpoint)
+    # Load weights (safetensors preferred, .pth fallback)
+    if checkpoint_path.endswith(".safetensors"):
+        from safetensors.torch import load_file
+        state_dict = load_file(checkpoint_path, device="cpu")
+    else:
+        checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
+        state_dict = checkpoint.get("state_dict", checkpoint)
 
     # Fix compiled model prefix & handle PosEmbed mismatch
     new_state_dict = {}
@@ -643,15 +658,13 @@ def _gpu_phase(video_path, resolution, despill_val, mask_mode,
                         break
 
                     if use_precomputed:
-                        frame_f32 = np.load(os.path.join(precompute_dir, f"frame_{frame_idx:05d}.npy"))
+                        frame_rgb = np.load(os.path.join(precompute_dir, f"frame_{frame_idx:05d}.npy"))
+                        frame_f32 = frame_rgb.astype(np.float32) / 255.0
                         mask_path = os.path.join(precompute_dir, f"mask_{frame_idx:05d}.npy")
                         if os.path.exists(mask_path):
                             mask = np.load(mask_path)
                             fast_n += 1
                         else:
-                            # BiRefNet fallback — load original RGB, run on GPU
-                            rgb_path = os.path.join(precompute_dir, f"rgb_{frame_idx:05d}.npy")
-                            frame_rgb = np.load(rgb_path)
                             tm = time.time()
                             mask = birefnet_frame(birefnet, frame_rgb)
                             t_mask += time.time() - tm
@@ -800,6 +813,7 @@ def _gpu_phase(video_path, resolution, despill_val, mask_mode,
 def process_video(video_path, resolution, despill_val, mask_mode,
                   auto_despeckle, despeckle_size, progress=gr.Progress()):
     """Orchestrator: precompute fast masks (CPU) → GPU inference → CPU I/O."""
+    t_total = time.time()
     if video_path is None:
         raise gr.Error("Please upload a video.")
 
@@ -830,12 +844,9 @@ def process_video(video_path, resolution, despill_val, mask_mode,
         else:
             mask = None
             needs_birefnet = True
-        # Save as compressed numpy (fast to load, no serialization overhead)
-        np.save(os.path.join(precompute_dir, f"frame_{frame_count:05d}.npy"), frame_f32)
+        np.save(os.path.join(precompute_dir, f"frame_{frame_count:05d}.npy"), frame_rgb)
         if mask is not None:
             np.save(os.path.join(precompute_dir, f"mask_{frame_count:05d}.npy"), mask)
-        if mask is None:
-            np.save(os.path.join(precompute_dir, f"rgb_{frame_count:05d}.npy"), frame_rgb)
         frame_count += 1
     cap.release()
     logger.info("[Phase 0] %d frames saved to %s in %.1fs (needs_birefnet=%s)",
@@ -886,9 +897,11 @@ def process_video(video_path, resolution, despill_val, mask_mode,
         n = len(frame_times)
         avg = np.mean(frame_times) if frame_times else 0
         engine = "PyTorch GPU" if use_gpu else "ONNX CPU"
+        wall = time.time() - t_total
         status = (f"Processed {n} frames ({w}x{h}) at {resolution}px | "
                   f"{avg:.2f}s/frame | {engine}" +
-                  (f" batch={batch_size}" if use_gpu else ""))
+                  (f" batch={batch_size}" if use_gpu else "") +
+                  f" | {wall:.0f}s total")
 
         return (
             comp_video if os.path.exists(comp_video) else None,