bug fix

app.py

@@ -92,6 +92,7 @@ def ensure_egoforce_repo() -> Path:
     demo_entrypoint = EGOFORCE_ROOT / "demo" / "run_app.py"
     if demo_entrypoint.exists():
         patch_upstream_gradio_for_zerogpu(demo_entrypoint)
+        patch_upstream_tensorrt_fallback(EGOFORCE_ROOT)
         return EGOFORCE_ROOT
 
     if EGOFORCE_ROOT.exists() and any(EGOFORCE_ROOT.iterdir()):
@@ -114,6 +115,7 @@ def ensure_egoforce_repo() -> Path:
         raise RuntimeError(f"EgoForce demo entrypoint not found at {demo_entrypoint}")
 
     patch_upstream_gradio_for_zerogpu(demo_entrypoint)
+    patch_upstream_tensorrt_fallback(EGOFORCE_ROOT)
     return EGOFORCE_ROOT
 
 
@@ -149,9 +151,141 @@ def patch_upstream_gradio_for_zerogpu(demo_entrypoint: Path) -> None:
             1,
         )
 
+    if "def load_gradio_hero_css():\n" not in source:
+        marker = "GRADIO_HERO_CSS_PATH = ASSETS_CSS_DIR / \"gradio_hero.css\"\n"
+        if marker not in source:
+            raise RuntimeError(f"Could not locate CSS path constant in {demo_entrypoint}")
+        source = source.replace(
+            marker,
+            (
+                marker +
+                "\n"
+                "@lru_cache(maxsize=1)\n"
+                "def load_gradio_hero_css():\n"
+                "    if not GRADIO_HERO_CSS_PATH.exists():\n"
+                "        return None\n"
+                "    return GRADIO_HERO_CSS_PATH.read_text(encoding=\"utf-8\")\n"
+            ),
+            1,
+        )
+
+    source = source.replace("        css=load_gradio_hero_css(),\n    ) as app:\n", "    ) as app:\n")
+
+    launch_css_marker = "        server_port=args.server_port,\n"
+    launch_css_line = "        css=load_gradio_hero_css(),\n"
+    if launch_css_line not in source:
+        if launch_css_marker not in source:
+            raise RuntimeError(f"Could not locate Gradio launch arguments in {demo_entrypoint}")
+        source = source.replace(launch_css_marker, launch_css_marker + launch_css_line, 1)
+
     demo_entrypoint.write_text(source, encoding="utf-8")
 
 
+def patch_upstream_tensorrt_fallback(repo_root: Path) -> None:
+    inference_path = repo_root / "demo" / "inference.py"
+    demo_utils_path = repo_root / "demo" / "demo_utils.py"
+
+    inference_source = inference_path.read_text(encoding="utf-8")
+    if "TORCH_TENSORRT_IMPORT_ERROR = None\n" not in inference_source:
+        import_marker = "import torch\nimport torch_tensorrt\n\n"
+        if import_marker not in inference_source:
+            raise RuntimeError(f"Could not locate torch_tensorrt import in {inference_path}")
+        inference_source = inference_source.replace(
+            import_marker,
+            (
+                "import torch\n"
+                "\n"
+                "try:\n"
+                "    import torch_tensorrt\n"
+                "    TORCH_TENSORRT_IMPORT_ERROR = None\n"
+                "except Exception as exc:\n"
+                "    torch_tensorrt = None\n"
+                "    TORCH_TENSORRT_IMPORT_ERROR = exc\n"
+                "    print(f\"Torch-TensorRT unavailable: {exc}. Falling back to PyTorch inference.\", flush=True)\n"
+                "\n"
+            ),
+            1,
+        )
+
+    runtime_marker = (
+        "torch_tensorrt.runtime.set_multi_device_safe_mode(True)\n"
+        "torch_tensorrt.runtime.set_cudagraphs_mode(True)\n"
+    )
+    if runtime_marker in inference_source:
+        inference_source = inference_source.replace(
+            runtime_marker,
+            (
+                "if torch_tensorrt is not None:\n"
+                "    torch_tensorrt.runtime.set_multi_device_safe_mode(True)\n"
+                "    torch_tensorrt.runtime.set_cudagraphs_mode(True)\n"
+            ),
+            1,
+        )
+    inference_path.write_text(inference_source, encoding="utf-8")
+
+    demo_utils_source = demo_utils_path.read_text(encoding="utf-8")
+    if "Torch-TensorRT backend unavailable" not in demo_utils_source:
+        old_compile_function = """def compile_to_tensorrt(model, device):
+    x1, x2, x3, x4 = torch.rand([2, 1, 3, 224, 224]), torch.rand([2, 1, 3, 6, 2]), torch.rand([2, 1, 3, 224, 224]), torch.rand([2, 1, 3, 6, 2])
+    x1, x2, x3, x4 = x1.to(device), x2.to(device), x3.to(device), x4.to(device)
+
+    with torch.inference_mode():
+        model = model.to(device).half()
+        x1, x2, x3, x4 = x1.half(), x2.half(), x3.half(), x4.half()
+        model = torch.jit.trace(model, (x1, x2, x3, x4), strict=False)
+
+    backend_kwargs = {
+        "enabled_precisions": {torch.half},
+        "min_block_size": 2,
+        "torch_executed_ops": {"torch.ops.aten.sub.Tensor"},
+        "optimization_level": 5,
+        "use_python_runtime": False,
+    }
+
+    model = torch.compile(model, backend="torch_tensorrt", options=backend_kwargs, dynamic=False,)
+    with torch.no_grad(): 
+        model(x1, x2, x3, x4) # compiled on first run
+
+    return model
+"""
+        new_compile_function = """def compile_to_tensorrt(model, device):
+    try:
+        import torch_tensorrt  # noqa: F401
+    except Exception as exc:
+        print(f"Torch-TensorRT backend unavailable: {exc}. Using PyTorch model.", flush=True)
+        return model.to(device).half()
+
+    x1, x2, x3, x4 = torch.rand([2, 1, 3, 224, 224]), torch.rand([2, 1, 3, 6, 2]), torch.rand([2, 1, 3, 224, 224]), torch.rand([2, 1, 3, 6, 2])
+    x1, x2, x3, x4 = x1.to(device), x2.to(device), x3.to(device), x4.to(device)
+
+    with torch.inference_mode():
+        fallback_model = model.to(device).half()
+        x1, x2, x3, x4 = x1.half(), x2.half(), x3.half(), x4.half()
+        traced_model = torch.jit.trace(fallback_model, (x1, x2, x3, x4), strict=False)
+
+    backend_kwargs = {
+        "enabled_precisions": {torch.half},
+        "min_block_size": 2,
+        "torch_executed_ops": {"torch.ops.aten.sub.Tensor"},
+        "optimization_level": 5,
+        "use_python_runtime": False,
+    }
+
+    try:
+        compiled_model = torch.compile(traced_model, backend="torch_tensorrt", options=backend_kwargs, dynamic=False,)
+        with torch.no_grad():
+            compiled_model(x1, x2, x3, x4) # compiled on first run
+        return compiled_model
+    except Exception as exc:
+        print(f"Torch-TensorRT compile failed: {exc}. Using PyTorch model.", flush=True)
+        return fallback_model
+"""
+        if old_compile_function not in demo_utils_source:
+            raise RuntimeError(f"Could not locate compile_to_tensorrt in {demo_utils_path}")
+        demo_utils_source = demo_utils_source.replace(old_compile_function, new_compile_function, 1)
+        demo_utils_path.write_text(demo_utils_source, encoding="utf-8")
+
+
 def package_available(module_name: str) -> bool:
     return importlib.util.find_spec(module_name) is not None
 

requirements.txt

@@ -1,8 +1,10 @@
 --extra-index-url https://miropsota.github.io/torch_packages_builder
+--find-links https://download.pytorch.org/whl/torch-tensorrt
 
 numpy==1.26.4
 torch==2.8.0
 torchvision==0.23.0
+torch_tensorrt==2.8.0+cu128
 pytorch3d==0.7.9+pt2.8.0cu128
 opencv-python==4.11.0.86
 pillow==11.3.0