Fix HF Spaces build timeout: limit mamba_ssm to sm_89 only

Previous fix (MAX_JOBS=1) prevented OOMKill but turned an already long
build into one exceeding HF Spaces' 60-minute build cap, because
mamba_ssm's setup.py hardcodes 8+ GPU architectures (sm_53..sm_90) and
serial compilation across all of them is too slow.

Real fix: patch mamba_ssm's setup.py at build time to compile only for
sm_89 (L40S / Ada / RTX 40-series, the actual target). With only one
architecture to build, MAX_JOBS=2 + NVCC_THREADS=2 is safe again.
Estimated build time drops from >60 min (timeout) to ~10-12 min.

- scripts/patch_mamba_ssm.py: line-based, idempotent setup.py patcher
with sentinel marker. Replaces cc_flag.append() calls with `pass` so
enclosing `if:` blocks remain syntactically valid.
- Dockerfile: pip download mamba_ssm sdist, run patcher, then install
from the patched source tree with --no-build-isolation.
Bump MAX_JOBS=2, NVCC_THREADS=2 (safe with single-arch build).

Also includes drift cleanup:
- app_gradio_baseline.py: add HfFolder shim + Gradio 4.44.x schema-bug
fix (mirrors app.py); bind to 127.0.0.1 by default on Windows so the
V1 screenshot can be reproduced locally without share=True.
- assets/gradio_v1_{reject,accept_png}.png: paper figure source images
showing gr.Image rejecting .nii.gz vs. accepting a normal PNG.
- .gitignore: ignore local venv directories (.v1/, .venv/, venv/).

.gitignore

@@ -12,6 +12,9 @@ v3_pred.nii.gz
 __pycache__/
 *.pyc
 *.pyo
+.v1/
+.venv/
+venv/
 
 # Research paper / LaTeX build artifacts (kept local only)
 Paper.tex

Dockerfile

@@ -12,12 +12,13 @@ ENV TORCH_CUDA_ARCH_LIST=8.9
 ENV CUDA_HOME=/usr/local/cuda
 ENV MPLCONFIGDIR=/tmp/matplotlib
 
-# Cap CUDA-extension build parallelism so the HF Spaces build VM (~16 GB
-# RAM) does not OOMKill during mamba_ssm / selective_scan_cuda_oflex
-# compilation. Each nvcc worker can use 3-4 GB; serial build is slower
-# (~25 min total) but reliable.
-ENV MAX_JOBS=1
-ENV NVCC_THREADS=1
+# Cap CUDA-extension build parallelism. mamba_ssm hardcodes 8+ GPU arches
+# in setup.py (sm_53..sm_90) which makes the build both memory-heavy AND
+# wall-clock heavy on HF Spaces (~16 GB RAM, ~60 min build cap).
+# We further patch mamba_ssm below to compile for only sm_89 (L40S target),
+# so MAX_JOBS=2 is now safe and ~6x faster than the all-arch serial build.
+ENV MAX_JOBS=2
+ENV NVCC_THREADS=2
 
 RUN apt-get update && apt-get install -y \
     python3.10 \
@@ -48,7 +49,19 @@ RUN pip install --no-cache-dir -r requirements.txt
 
 COPY . .
 
-RUN pip install "mamba-ssm>=2.2.2" --no-build-isolation -v
+# Download mamba-ssm source, patch its hardcoded cc_flag list to only build
+# for sm_89 (L40S / Ada / RTX 40-series), then install. This avoids HF
+# Spaces 60-minute build timeouts when compiling for all 8+ archs serially.
+COPY scripts/patch_mamba_ssm.py /tmp/patch_mamba_ssm.py
+RUN set -eux; \
+    pip download "mamba-ssm>=2.2.2" --no-deps --no-binary=:all: -d /tmp/mamba-src; \
+    cd /tmp/mamba-src; \
+    SDIST=$(ls mamba_ssm*.tar.gz mamba-ssm*.tar.gz 2>/dev/null | head -1); \
+    tar -xzf "$SDIST"; \
+    SRCDIR=$(tar -tzf "$SDIST" | head -1 | sed 's:/.*::'); \
+    cd "$SRCDIR"; \
+    python /tmp/patch_mamba_ssm.py setup.py; \
+    pip install . --no-build-isolation -v
 
 RUN python - <<'PY'
 import torch, sys

app_gradio_baseline.py

@@ -47,9 +47,64 @@ from pathlib import Path
 import numpy as np
 import torch
 import nibabel as nib
+
+
+# ---------------------------------------------------------------------------
+# Compatibility shims (mirror app.py).
+# Must run BEFORE `import gradio as gr` for the HfFolder shim, and AFTER for
+# the schema bug fix.
+# ---------------------------------------------------------------------------
+def _patch_hf_folder():
+    """Gradio 4.44.x imports HfFolder from huggingface_hub, which was removed
+    in huggingface_hub 1.0+. Inject a minimal shim."""
+    import huggingface_hub as hh
+    if getattr(hh, "HfFolder", None) is not None:
+        return
+    try:
+        from huggingface_hub import get_token
+    except ImportError:
+        get_token = lambda: None  # noqa: E731
+
+    class HfFolder:
+        @staticmethod
+        def get_token():
+            return get_token()
+
+    hh.HfFolder = HfFolder
+
+
+_patch_hf_folder()
+
 import gradio as gr
 
 
+def _fix_gradio_schema_bug():
+    """Patch gradio_client.utils.get_type to handle boolean schemas
+    (Gradio 4.44.x crashes on additionalProperties: True)."""
+    try:
+        import gradio_client.utils as gu
+        if not hasattr(gu, "get_type"):
+            return
+        original = gu.get_type
+
+        def patched(schema):
+            if isinstance(schema, bool):
+                return "Any"
+            if isinstance(schema, dict):
+                if schema.get("additionalProperties") is True:
+                    schema["additionalProperties"] = {}
+                elif schema.get("additionalProperties") is False:
+                    schema.pop("additionalProperties", None)
+            return original(schema)
+
+        gu.get_type = patched
+    except Exception:
+        pass
+
+
+_fix_gradio_schema_bug()
+
+
 _model_cache: dict = {}
 
 
@@ -356,8 +411,12 @@ def main() -> None:
     else:
         demo = build_v3_demo()
 
+    server_name = os.environ.get(
+        "GRADIO_SERVER_NAME",
+        "127.0.0.1" if os.name == "nt" else "0.0.0.0",
+    )
     demo.launch(
-        server_name="0.0.0.0",
+        server_name=server_name,
         server_port=args.port,
         share=args.share,
     )

scripts/patch_mamba_ssm.py

@@ -0,0 +1,104 @@
+"""Patch mamba-ssm's setup.py to compile CUDA kernels for only sm_89.
+
+Mamba-ssm's upstream setup.py builds for 8+ GPU architectures (sm_53,
+sm_60, sm_70, sm_75, sm_80, sm_86, sm_89, sm_90). On the HF Spaces
+build VM (~16 GB RAM, ~60 min wall-clock cap) compiling all of them
+either OOMKills or hits the build timeout. This patch reduces the list
+to only sm_89 (the L40S target the production VoxPixel and bench
+Spaces both run on).
+
+Strategy (line-based, robust to upstream formatting changes):
+
+    1. Find the line ``cc_flag = []``.
+    2. Replace it with a list literal pinned to sm_89.
+    3. Comment out every subsequent line containing
+       ``cc_flag.append(`` (these would otherwise re-inject the other
+       architectures, including via conditional ``if`` blocks).
+
+Idempotent: the patch leaves a sentinel marker and is a no-op on
+already-patched files.
+
+Usage:
+    python scripts/patch_mamba_ssm.py path/to/setup.py
+"""
+
+import re
+import sys
+from pathlib import Path
+
+
+SENTINEL = "# PATCHED: cc_flag limited to sm_89 by patch_mamba_ssm.py"
+TARGET_ARCH = "89"
+
+
+def _leading_indent(line: str) -> str:
+    return line[: len(line) - len(line.lstrip())]
+
+
+def patch(path: Path) -> bool:
+    src = path.read_text()
+    if SENTINEL in src:
+        print(f"{path}: already patched (sentinel present), skipping.")
+        return False
+    if "cc_flag" not in src:
+        raise SystemExit(
+            f"ERROR: 'cc_flag' not found in {path}. "
+            "Upstream mamba-ssm setup.py may have changed."
+        )
+
+    init_pattern = re.compile(r"^cc_flag\s*=\s*\[\]\s*$")
+    out_lines: list[str] = []
+    init_patched = False
+    appends_neutralized = 0
+
+    for line in src.splitlines(keepends=True):
+        stripped = line.strip()
+
+        if not init_patched and init_pattern.match(stripped):
+            indent = _leading_indent(line)
+            out_lines.append(f"{indent}{SENTINEL}\n")
+            out_lines.append(
+                f'{indent}cc_flag = ["-gencode", '
+                f'"arch=compute_{TARGET_ARCH},code=sm_{TARGET_ARCH}"]\n'
+            )
+            init_patched = True
+            continue
+
+        if init_patched and "cc_flag.append(" in stripped:
+            # Replace with `pass` (not a comment) so any enclosing
+            # `if:` / `else:` block keeps a syntactically valid body.
+            indent = _leading_indent(line)
+            out_lines.append(
+                f"{indent}pass  # (was: {stripped})\n"
+            )
+            appends_neutralized += 1
+            continue
+
+        out_lines.append(line)
+
+    if not init_patched:
+        raise SystemExit(
+            f"ERROR: line 'cc_flag = []' not found in {path}. "
+            "Upstream mamba-ssm setup.py format may have changed."
+        )
+
+    path.write_text("".join(out_lines))
+    print(
+        f"{path}: patched. cc_flag pinned to sm_{TARGET_ARCH}; "
+        f"{appends_neutralized} subsequent cc_flag.append calls neutralized."
+    )
+    return True
+
+
+def main() -> None:
+    if len(sys.argv) != 2:
+        print("usage: patch_mamba_ssm.py <path/to/setup.py>", file=sys.stderr)
+        sys.exit(2)
+    path = Path(sys.argv[1])
+    if not path.is_file():
+        raise SystemExit(f"ERROR: {path} not found.")
+    patch(path)
+
+
+if __name__ == "__main__":
+    main()