MLLM-HWSI Inference App

⚠️ Reproducibility status — read first. While validating this app we found that the public Bastech/MLLM-HWSI checkpoint cannot reproduce the paper's image-conditional results: the released VL projector is "dead" (it maps every slide — and random noise — to the same tokens, so the model can't see the image), and the paper's zero-shot classification relies on Stage-I encoders that were never released. See FINDINGS.md for the evidence and scripts/diagnostics/ for the reproduction. The app code below is correct and runs; the released weights are the limitation.

A small, Dockerized web app to run inference with MLLM-HWSI (Alawode et al., CVPR 2026) — a hierarchical, WSI-level multimodal LLM — on a single 24 GB GPU. It wraps the upstream inference code with a config-driven, two-phase design:

• Offline "build" — run the visual encoders (CellViT, CONCH, HIPT) over each whole-slide image one at a time and cache ~1941 multi-scale tokens per slide to disk. Heavy, one-time, time-bound.
• Online "serve" — load only the LoRA-merged Qwen2.5-7B + projectors (~17 GB fp16) and answer questions from the cached tokens via a Gradio web UI, reachable on this host's public IP.

Supports three tasks over a slide: VQA (Morphology), report generation, and zero-shot classification — they share the same feature cache and serving core; only the prompt/eval differ.

Why this fits 24 GB

The encoders never coexist with the LLM. Offline extraction loads each encoder sequentially (a few GB each); serving loads only LLM+projectors (~17 GB peak incl. activations). No quantization needed on a TITAN RTX / RTX 3090 / 4090-class card.

Layout

config.yaml          # single source of truth (paths, dataset, top_k, magnification, generation, proxy)
app/
  config.py          # typed Settings loaded from config.yaml + .env
  serve_app.py       # serve phase: FastAPI (DeepZoom tiles + viewer) with Gradio mounted at /
  evidence.py        # render HPS-selected evidence patches + slide crops
  classify.py        # zero-shot class-name likelihood scoring (Phase 4)
  netutil.py         # proxy fallback on network error
  static/openseadragon/  # vendored deep-zoom viewer assets
scripts/
  download_weights.py     # HF snapshot (LLM+projectors) + gdown encoders
  download_data.py        # WSI-Bench JSONL + TCGA WSI subsample (GDC), task-aware selection
  build_features.py       # OFFLINE orchestrator: WSI -> cached tokens (chains the 3 ext stages)
  eval.py                 # wraps upstream main_test.py per task, adds bootstrap CIs
  eval_classification.py  # Phase 4: closed-set zero-shot classification (accuracy + macro/micro F1)
upstream/            # vendored upstream HWSI-MLLM code (flat, imports preserved; +utils shim, patches)
Dockerfile, docker-compose.yml, Makefile, pyproject.toml   # house-style scaffold (uv-based)

Web UI (serve)

The serve phase is a FastAPI app with Gradio mounted at /, exposing:

• Slide Q&A — pick a slide, choose a task (Morphology / Report / Classification / Caption), ask.
• Deep-zoom preview — pan/zoom gigapixel viewer (OpenSlide DeepZoom tiles + OpenSeadragon), served at /viewer/<slide> with tiles at /dzi/<slide>....
• Evidence patches — the HPS-selected patches the model conditioned on, cropped from the WSI. These are vision-based (query-independent): the same evidence set for any question on that slide.

Prerequisites

• Linux, Docker + nvidia-container-toolkit, one ≥24 GB NVIDIA GPU.

GPU runtime (snap Docker)

This compose file requests the GPU via runtime: nvidia (not gpus: all). On snap-installed Docker whose toolkit defaults to CDI mode without generated specs, set legacy injection once:

sudo sed -i 's/^mode = "cdi"/mode = "legacy"/' \
  /var/snap/docker/current/etc/nvidia-container-runtime/config.toml
sudo snap restart docker

(A snap refresh can reset this; re-apply if GPU access breaks. Alternatively, keep CDI mode and run sudo nvidia-ctk cdi generate --output=/etc/cdi/nvidia.yaml.) Verify: sudo docker compose run --rm mllm-hwsi nvidia-smi should list the GPU.

• cp .env.example .env and (optionally) set HF_TOKEN / PROXY_URL.
• The inference weights are NOT gated (HF Bastech/MLLM-HWSI); encoders come from the authors' Google Drive — so no Hugging Face access requests are required.

Quickstart

cp .env.example .env

make build                     # build the Docker image (uv installs the stack)
make download-weights          # LLM+projectors (HF) + CellViT/CONCH/HIPT (Google Drive)
make download-data             # WSI-Bench JSONL + a TCGA subsample (config: dataset.n_slides)

make build-features            # OFFLINE: extract & cache tokens for all downloaded slides
#   or one slide:  make build-features SLIDE=<slide_id>

make serve                     # ONLINE: web UI at http://<this-host-public-ip>:7860
make eval TASK=Morphology      # text-metric eval (also: Report | Caption); Classification too
make classify                  # Phase 4: zero-shot classification by class-name scoring
make help                      # list all targets

Scaling the eval

dataset.n_slides controls how many slides are downloaded/built. Selection is task-aware (download_data.py --strategy coverage): it picks slides that have all of Morphology/Report/ Classification, so the evals always have records. Bump n_slides (e.g. 25), then re-run download-data → build-features → the evals. Budget ~0.5–2 GB/slide download and ~minutes/slide extraction; make clean reclaims the large intermediate features/_regions.

Configuration

Everything is in config.yaml. Common edits:

• dataset.n_slides — subsample size (disk control; ~0.5–2 GB/slide).
• build.top_k — patches kept per region after HPS filtering (repo default 16; paper ~48).
• build.magnification — 20 or 40.
• serve.generation.* — decoding params.
• network.use_proxy / proxy_url — retry network ops through a proxy on failure.

Networking / remote access

Gradio binds 0.0.0.0:7860 and the container publishes 7860, so the Q&A box is reachable at http://<host-public-ip>:7860 once the host firewall/security-group allows the port. Set serve.share: true for a temporary public Gradio tunnel instead.

Dataset

First target is WSI-Bench (from WSI-LLaVA): a JSONL of 2837 records over 527 TCGA slides spanning many cancer types, each record tagged with a metadata test_type (Morphology, Report, Classification, Caption, …). download_data.py saves it verbatim (upstream main_test.py reads JSONL) and pulls the matching TCGA .svs from the GDC open-access API.

Results

Preliminary, on a tiny coverage subsample (validates the pipeline; not benchmark-grade — scale n_slides up for real numbers):

Note the contrast: free-text + exact-match scores 0 even when answers are close, and BERTScore (~0.88) is fooled by templated phrasing — so classification uses make classify (likelihood scoring over candidate labels), which reports accuracy + macro/micro F1. Numbers above are from n=3 slides and are illustrative only.

Notes & caveats

• Offline extraction is ~minutes/slide and produces large intermediate region .h5 under features/_regions — make clean purges them after the token cache is built.
• upstream/requirements.txt pins bleeding-edge versions (torch cu130, transformers 5.5); if your driver predates CUDA 13, repin torch and adjust the Dockerfile base tag.
• See HANDOFF.md for full context and the phased build status.