GPU postprocessing pipeline + TF32 + conditional torch.compile

app.py

@@ -36,6 +36,17 @@ try:
 except ImportError:
     HAS_SPACES = False
 
+# GPU perf: TF32 tensor cores for FP32 postprocessing ops
+try:
+    import torch as _torch
+    _torch.set_float32_matmul_precision('high')
+    del _torch
+except ImportError:
+    pass
+# Persist torch.compile inductor cache across ZeroGPU sessions
+_inductor_cache = os.path.join(os.path.expanduser("~"), ".cache", "corridorkey", "inductor")
+os.environ.setdefault("TORCHINDUCTOR_CACHE_DIR", _inductor_cache)
+
 # Workaround: Gradio cache_examples bug with None outputs.
 _original_read_from_flag = gr.components.Component.read_from_flag
 def _patched_read_from_flag(self, payload):
@@ -116,7 +127,7 @@ def clean_matte(alpha_np, area_threshold=300, dilation=15, blur_size=5):
     is_3d = alpha_np.ndim == 3
     if is_3d:
         alpha_np = alpha_np[:, :, 0]
-    mask_8u = (alpha_np > 0.5).astype(np.uint8) * 255
+    mask_8u = (alpha_np > 0.25).astype(np.uint8) * 255
     num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask_8u, connectivity=8)
     valid = np.zeros(num_labels, dtype=bool)
     valid[1:] = stats[1:, cv2.CC_STAT_AREA] >= area_threshold
@@ -138,6 +149,60 @@ def create_checkerboard(w, h, checker_size=64, color1=0.15, color2=0.55):
 def premultiply(fg, alpha):
     return fg * alpha
 
+# ---------------------------------------------------------------------------
+# GPU postprocessing utilities (keep data on device, avoid CPU↔GPU transfers)
+# ---------------------------------------------------------------------------
+def despill_torch(image, strength, screen_channel=1):
+    """GPU despill on [B, 3, H, W] tensor."""
+    import torch
+    if strength <= 0.0:
+        return image
+    other_a, other_b = [i for i in (0, 1, 2) if i != screen_channel]
+    screen, a, b = image[:, screen_channel], image[:, other_a], image[:, other_b]
+    spill = torch.clamp(screen - (a + b) / 2.0, min=0.0)
+    out = [None, None, None]
+    out[screen_channel] = screen - spill
+    out[other_a] = a + spill * 0.5
+    out[other_b] = b + spill * 0.5
+    despilled = torch.stack(out, dim=1)
+    return image * (1.0 - strength) + despilled * strength if strength < 1.0 else despilled
+
+
+def _connected_components_gpu(mask, min_dist=4, max_iter=100):
+    """GPU flood-fill connected components on [B, 1, H, W] binary mask."""
+    import torch
+    import torch.nn.functional as F
+    bs, _, H, W = mask.shape
+    comp = (torch.randperm(bs * W * H, device=mask.device, dtype=torch.float32) + 1.1).view(mask.shape)
+    comp[mask != 1] = 0
+    k = 2 * min_dist + 1
+    for _ in range(max_iter):
+        comp[mask == 1] = F.max_pool2d(comp, k, stride=1, padding=min_dist)[mask == 1]
+    _, comp = torch.unique(comp, return_inverse=True)
+    return comp.view(mask.shape)
+
+
+def clean_matte_torch(alpha, area_threshold, dilation=25, blur_size=5):
+    """GPU clean matte on [B, 1, H, W] tensor. Removes small disconnected blobs."""
+    import torch
+    import torch.nn.functional as F
+    import torchvision.transforms.functional as TF
+    mask = (alpha > 0.25).float()
+    comp = _connected_components_gpu(mask, max_iter=max(area_threshold // 20, 5))
+    sizes = torch.bincount(comp.flatten())
+    big = torch.nonzero(sizes >= area_threshold).squeeze(-1)
+    big = big[big > 0]
+    cleaned = torch.zeros_like(mask)
+    if big.numel() > 0:
+        cleaned[torch.isin(comp, big)] = 1.0
+    if dilation > 0:
+        for _ in range(dilation // 2):
+            cleaned = F.max_pool2d(cleaned, 5, stride=1, padding=2)
+    if blur_size > 0:
+        cleaned = TF.gaussian_blur(cleaned, [blur_size * 2 + 1, blur_size * 2 + 1])
+    return alpha * cleaned
+
+
 # ---------------------------------------------------------------------------
 # Fast classical green-screen mask
 # ---------------------------------------------------------------------------
@@ -288,8 +353,14 @@ def _load_greenformer(img_size):
                 torch.backends.cuda.mem_efficient_sdp_enabled(),
                 torch.backends.cuda.math_sdp_enabled())
 
-    # Skip torch.compile on ZeroGPU — the 37s warmup eats too much of the 120s budget.
-    if not HAS_SPACES and sys.platform in ("linux", "win32"):
+    should_compile = False
+    if sys.platform in ("linux", "win32"):
+        if not HAS_SPACES:
+            should_compile = True
+        elif os.path.isdir(_inductor_cache) and os.listdir(_inductor_cache):
+            should_compile = True
+            logger.info("Warm ZeroGPU: using cached inductor kernels")
+    if should_compile:
         try:
             compiled = torch.compile(model)
             dummy = torch.zeros(1, 4, img_size, img_size, dtype=torch.float16, device="cuda")
@@ -301,7 +372,7 @@ def _load_greenformer(img_size):
             logger.warning("torch.compile() failed, using eager mode: %s", e)
             torch.cuda.empty_cache()
     else:
-        logger.info("Skipping torch.compile() (ZeroGPU: saving GPU time for inference)")
+        logger.info("Skipping torch.compile() (cold start, no cached kernels)")
 
     logger.info("GreenFormer loaded on CUDA (img_size=%d)", img_size)
     return model
@@ -359,61 +430,58 @@ def corridorkey_frame_onnx(session, image_f32, mask_f32, img_size,
 # ---------------------------------------------------------------------------
 def corridorkey_batch_pytorch(model, images_f32, masks_f32, img_size,
                               despill_strength=0.5, auto_despeckle=True, despeckle_size=400):
-    """PyTorch batched inference for multiple frames on GPU.
-
-    Args:
-        model: GreenFormer model on CUDA
-        images_f32: list of [H, W, 3] float32 numpy arrays (0-1, sRGB)
-        masks_f32: list of [H, W] float32 numpy arrays (0-1)
-        img_size: model input resolution (1024 or 2048)
-
-    Returns:
-        list of dicts with 'alpha' [H,W,1] and 'fg' [H,W,3]
-    """
+    """PyTorch batched inference with full GPU pipeline (preprocess + inference + postprocess on device)."""
     import torch
+    import torchvision.transforms.functional as TF
 
     batch_size = len(images_f32)
     if batch_size == 0:
         return []
 
-    # Store original sizes per frame
-    orig_sizes = [(img.shape[1], img.shape[0]) for img in images_f32]  # (w, h)
+    w, h = images_f32[0].shape[1], images_f32[0].shape[0]
+
+    # --- GPU Preprocessing in FP32 (avoids FP16 precision loss in normalize/resize) ---
+    batch_imgs = torch.stack([
+        torch.from_numpy(img.transpose(2, 0, 1)) for img in images_f32
+    ]).cuda()
+    batch_masks = torch.stack([
+        torch.from_numpy(m if m.ndim == 2 else m[:, :, 0]).unsqueeze(0) for m in masks_f32
+    ]).cuda()
 
-    # Preprocess: resize, normalize, concatenate into batch tensor
-    batch_inputs = []
-    for img, mask in zip(images_f32, masks_f32):
-        img_r = cv2.resize(img, (img_size, img_size))
-        mask_r = cv2.resize(mask, (img_size, img_size))[:, :, np.newaxis]
-        inp = np.concatenate([(img_r - IMAGENET_MEAN) / IMAGENET_STD, mask_r], axis=-1)
-        batch_inputs.append(inp.transpose(2, 0, 1))  # [4, H, W]
+    batch_imgs = TF.resize(batch_imgs, [img_size, img_size], antialias=False)
+    batch_masks = TF.resize(batch_masks, [img_size, img_size], antialias=False)
+    batch_imgs = TF.normalize(batch_imgs, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
 
-    batch_np = np.stack(batch_inputs, axis=0).astype(np.float32)  # [B, 4, H, W]
-    batch_tensor = torch.from_numpy(batch_np).cuda().half()  # FP16 input
+    inp = torch.cat([batch_imgs, batch_masks], dim=1).half()
+    del batch_imgs, batch_masks
 
-    # Forward pass — model is FP16, input is FP16, no autocast needed
+    # --- Forward Pass ---
     with torch.inference_mode():
-        out = model(batch_tensor)
+        out = model(inp)
+    del inp
+
+    # --- GPU Postprocessing (despill + clean_matte + resize stay on device) ---
+    alpha = out["alpha"].float()
+    fg = out["fg"].float()
 
-    # Extract results
-    alphas_gpu = out["alpha"].float().cpu().numpy()  # [B, 1, H, W]
-    fgs_gpu = out["fg"].float().cpu().numpy()         # [B, 3, H, W]
+    alpha = TF.resize(alpha, [h, w])
+    fg = TF.resize(fg, [h, w])
 
-    del batch_tensor
-    # Don't empty cache per batch - too expensive. Let PyTorch manage.
+    if auto_despeckle:
+        alpha = clean_matte_torch(alpha, area_threshold=int(despeckle_size), dilation=25, blur_size=5)
+    fg = despill_torch(fg, despill_strength)
+
+    # --- Single CPU transfer at the end ---
+    alpha_np = alpha.cpu().numpy()
+    fg_np = fg.cpu().numpy()
+    del alpha, fg
 
-    # Postprocess each frame
     results = []
     for i in range(batch_size):
-        w, h = orig_sizes[i]
-        alpha = cv2.resize(alphas_gpu[i].transpose(1, 2, 0), (w, h), interpolation=cv2.INTER_LANCZOS4)
-        fg = cv2.resize(fgs_gpu[i].transpose(1, 2, 0), (w, h), interpolation=cv2.INTER_LANCZOS4)
-        if alpha.ndim == 2:
-            alpha = alpha[:, :, np.newaxis]
-        if auto_despeckle:
-            alpha = clean_matte(alpha, area_threshold=despeckle_size, dilation=25, blur_size=5)
-        fg = despill(fg, green_limit_mode="average", strength=despill_strength)
-        results.append({"alpha": alpha, "fg": fg})
-
+        results.append({
+            "alpha": alpha_np[i].transpose(1, 2, 0),
+            "fg": fg_np[i].transpose(1, 2, 0),
+        })
     return results