changed to mmcv-lite

app.py

@@ -120,6 +120,15 @@ def ensure_egoforce_repo() -> Path:
 def patch_upstream_gradio_for_zerogpu(demo_entrypoint: Path) -> None:
     source = demo_entrypoint.read_text(encoding="utf-8")
 
+    if "from egoforce_runtime_patches import apply_runtime_patches\n" not in source:
+        if "import torch\n" not in source:
+            raise RuntimeError(f"Could not insert runtime patches in {demo_entrypoint}")
+        source = source.replace(
+            "import torch\n",
+            "import torch\nfrom egoforce_runtime_patches import apply_runtime_patches\napply_runtime_patches()\n",
+            1,
+        )
+
     if "import spaces\n" not in source:
         if "import torch\n" not in source:
             raise RuntimeError(f"Could not insert ZeroGPU import in {demo_entrypoint}")
@@ -164,12 +173,11 @@ def pip_install(requirement: str, *extra_args: str) -> None:
 def ensure_runtime_python_packages(repo_root: Path) -> None:
     datapipes_path = repo_root / "thirdparty" / "datapipes"
     install_plan = [
-        ("mmcv", "mmcv==2.1.0", ("--no-build-isolation", "--no-deps")),
         ("anycalib", "git+https://github.com/javrtg/AnyCalib.git", ("--no-build-isolation",)),
         ("chumpy", "git+https://github.com/mattloper/chumpy.git", ("--no-build-isolation",)),
         ("pytorch3d", "git+https://github.com/facebookresearch/pytorch3d.git", ("--no-build-isolation",)),
         ("datapipes", str(datapipes_path), ()),
-        ("mmdet", str(repo_root / "thirdparty" / "mmdetection"), ("--no-build-isolation",)),
+        ("mmdet", str(repo_root / "thirdparty" / "mmdetection"), ("--no-build-isolation", "--no-deps")),
     ]
 
     for module_name, requirement, extra_args in install_plan:

egoforce_runtime_patches.py

@@ -0,0 +1,119 @@
+from __future__ import annotations
+
+import importlib
+import sys
+import types
+from typing import Any
+
+
+def _torchvision_nms():
+    from torchvision.ops import nms
+
+    return nms
+
+
+def _nms(
+    boxes: Any,
+    scores: Any,
+    iou_threshold: float,
+    offset: int = 0,
+    score_threshold: float = 0,
+    max_num: int = -1,
+) -> tuple[Any, Any]:
+    import torch
+
+    if boxes.numel() == 0 or scores.numel() == 0:
+        keep = torch.empty((0,), dtype=torch.long, device=scores.device)
+        dets = torch.cat((boxes.reshape(0, boxes.shape[-1]), scores.reshape(0, 1)), dim=1)
+        return dets, keep
+
+    if score_threshold > 0:
+        valid = scores > score_threshold
+        original_indices = torch.nonzero(valid, as_tuple=False).squeeze(1)
+        filtered_boxes = boxes[valid]
+        filtered_scores = scores[valid]
+    else:
+        original_indices = torch.arange(scores.numel(), device=scores.device)
+        filtered_boxes = boxes
+        filtered_scores = scores
+
+    keep_local = _torchvision_nms()(filtered_boxes, filtered_scores, float(iou_threshold))
+    if max_num > 0:
+        keep_local = keep_local[:max_num]
+
+    keep = original_indices[keep_local]
+    dets = torch.cat((filtered_boxes[keep_local], filtered_scores[keep_local, None]), dim=1)
+    return dets, keep
+
+
+def _batched_nms(
+    boxes: Any,
+    scores: Any,
+    idxs: Any,
+    nms_cfg: dict[str, Any] | None,
+    class_agnostic: bool = False,
+) -> tuple[Any, Any]:
+    import torch
+
+    if boxes.numel() == 0 or scores.numel() == 0:
+        keep = torch.empty((0,), dtype=torch.long, device=scores.device)
+        dets = torch.cat((boxes.reshape(0, boxes.shape[-1]), scores.reshape(0, 1)), dim=1)
+        return dets, keep
+
+    if nms_cfg is None:
+        order = scores.argsort(descending=True)
+        return torch.cat((boxes[order], scores[order, None]), dim=1), order
+
+    nms_cfg = dict(nms_cfg)
+    iou_threshold = nms_cfg.pop("iou_threshold", nms_cfg.pop("iou_thr", 0.5))
+    score_threshold = nms_cfg.pop("score_threshold", 0)
+    max_num = nms_cfg.pop("max_num", -1)
+
+    if class_agnostic:
+        boxes_for_nms = boxes
+    else:
+        max_coordinate = boxes.max()
+        offsets = idxs.to(boxes) * (max_coordinate + boxes.new_tensor(1))
+        boxes_for_nms = boxes + offsets[:, None]
+
+    if score_threshold > 0:
+        valid = scores > score_threshold
+        original_indices = torch.nonzero(valid, as_tuple=False).squeeze(1)
+        boxes_for_nms = boxes_for_nms[valid]
+        scores_for_nms = scores[valid]
+    else:
+        original_indices = torch.arange(scores.numel(), device=scores.device)
+        scores_for_nms = scores
+
+    keep_local = _torchvision_nms()(boxes_for_nms, scores_for_nms, float(iou_threshold))
+    if max_num > 0:
+        keep_local = keep_local[:max_num]
+
+    keep = original_indices[keep_local]
+    dets = torch.cat((boxes[keep], scores[keep, None]), dim=1)
+    return dets, keep
+
+
+def apply_runtime_patches() -> None:
+    try:
+        mmcv = importlib.import_module("mmcv")
+    except ImportError:
+        mmcv = None
+
+    ops_module = sys.modules.get("mmcv.ops")
+    if ops_module is None:
+        ops_module = types.ModuleType("mmcv.ops")
+        sys.modules["mmcv.ops"] = ops_module
+
+    nms_module = sys.modules.get("mmcv.ops.nms")
+    if nms_module is None:
+        nms_module = types.ModuleType("mmcv.ops.nms")
+        sys.modules["mmcv.ops.nms"] = nms_module
+
+    ops_module.nms = _nms
+    ops_module.batched_nms = _batched_nms
+    nms_module.nms = _nms
+    nms_module.batched_nms = _batched_nms
+
+    if mmcv is not None:
+        mmcv.ops = ops_module

requirements.txt

@@ -22,6 +22,7 @@ pycocotools==2.0.10
 trimesh==4.11.3
 sortedcontainers==2.4.0
 openmim==0.3.9
+mmcv-lite==2.1.0
 mmengine==0.10.7
 yapf==0.43.0
 lmdb==2.0.0