Sync from GitHub 0fe9a08

app.py

@@ -122,6 +122,10 @@ def _build_libraries():
 VISUAL_STORE, PARSED_STORE, PIPELINE = _build_libraries()
 # Valid agent-brain keys, for validating the per-request `agent_model`.
 _AGENT_MODEL_KEYS = {m["key"] for m in AGENT_MODELS}
+# Brains too large to keep the ColEmbed visual retriever resident alongside them
+# (constants.AGENT_MODELS `forbid_visual`). While one is active, a "visual"
+# retrieval_mode is forced to "parsed" so ColEmbed never loads next to it (OOM).
+_NO_VISUAL_MODEL_KEYS = {m["key"] for m in AGENT_MODELS if m.get("forbid_visual")}
 # Valid search-index keys, for validating the per-request `retrieval_mode`.
 _RETRIEVAL_MODE_KEYS = {m["key"] for m in RETRIEVAL_MODES}
 # method -> store, for the picker / pdf lookups. In mock both keys map to the
@@ -312,6 +316,10 @@ def api_find(
         agent_model = DEFAULT_AGENT_MODEL
     if retrieval_mode not in _RETRIEVAL_MODE_KEYS:
         retrieval_mode = DEFAULT_RETRIEVAL_MODE
+    # The big brains can't share VRAM with the ColEmbed visual retriever — force
+    # parsed so ColEmbed never loads alongside them (the UI also greys it out).
+    if retrieval_mode == "visual" and agent_model in _NO_VISUAL_MODEL_KEYS:
+        retrieval_mode = DEFAULT_RETRIEVAL_MODE
     log.info(
         "find: manual=%s k=%s think=%s model=%s search=%s viewer=%s hist=%d q=%r",
         manual, k, bool(think), agent_model, retrieval_mode, viewer,
@@ -369,7 +377,11 @@ def index():
         # so the settings dropdown needs no extra round-trip on load.
         .replace(
             "__AGENT_MODELS_JSON__",
-            json.dumps([{"key": m["key"], "label": m["label"]} for m in AGENT_MODELS]),
+            json.dumps([
+                {"key": m["key"], "label": m["label"],
+                 "forbidVisual": bool(m.get("forbid_visual"))}
+                for m in AGENT_MODELS
+            ]),
         )
         .replace("__AGENT_MODEL__", DEFAULT_AGENT_MODEL)
         # Search-index picker: which index the search tool ranks against

core/constants.py

@@ -45,44 +45,45 @@ MINICPM_AGENT_REVISION = os.environ.get("MINICPM_AGENT_REVISION", "") or None
 
 # Selectable agent brains, offered in the UI settings panel. ONE model is meant
 # to be resident in VRAM at a time — switching evicts the previous and loads the
-# next (models/minicpm_agent.use_model). NOTE: the FIRST/default brain is loaded at
-# import (ZeroGPU emulation phase) and materialized into the forked GPU worker,
-# where empty_cache() can't reclaim it — it stays stuck for the process. A brain
-# SWITCHED IN at runtime must fit the headroom ABOVE the resident set
-# (VLM+ColEmbed+embedder) + that stuck default (~15 GiB free on the 48 GiB `large`
-# slice — see core/vram.py).
+# next (models/minicpm_agent.use_model). The FIRST/default brain is loaded at
+# import (ZeroGPU's startup phase) and "packed" into the forked GPU worker, so it
+# stays resident for the whole process and the common (no-switch) turn pays NO
+# per-turn load cost. A brain SWITCHED IN at runtime is built inside the GPU
+# window instead (not packed), so its first turn after a switch is slower.
+#
+# The default is MiniCPM4.1-8B in bf16 (~16 GiB). It fits because the ColEmbed
+# visual retriever (~8 GiB) is NO LONGER packed at import — it lazy-loads only
+# when the "visual" search index is used (models/colembed.py). So the resident set
+# is the MiniCPM-V "eyes" + the Nemotron text embedder + this 8B brain, which
+# leaves room for the grounding spike on the 48 GiB slice. Because the 8B and
+# ColEmbed cannot BOTH be resident, the 8B sets forbid_visual: a turn that asks for
+# visual search while it is active is served by the parsed index instead (enforced
+# in app.py). Smaller brains leave room for ColEmbed to lazy-load, so they keep
+# visual search.
 #
-# The default is MiniCPM4.1-8B at int8 (~8.5 GiB via bitsandbytes). As the LONE
-# stuck brain it REPLACES the old 1B default (it does not stack on it), so the
-# resident set + one ~8.5 GiB brain still leaves room for the grounding spike — the
-# same footprint class as MiniCPM3-4B, which was VRAM-vetted to fit. (bf16 8B at
-# ~16 GiB still does NOT fit; int8 is what makes the 8B deployable as the default.)
-# Caveat: with an ~8.5 GiB default already stuck, switching ANOTHER 4-8 GiB brain
-# in at runtime for a UI A/B is tight and may OOM at grounding — the default path
-# is fine.
 # Each loads as an AutoModelForCausalLM; `trust_remote_code` (default False) flags
 # the ones that ship custom modeling code (MiniCPM3 / MiniCPM4.1). `thinking` flags
 # whether the chat template accepts enable_thinking (Qwen3, MiniCPM5, MiniCPM4.1 do
-# — tool routing passes it False; MiniCPM3 does not). int8/4bit brains need
-# bitsandbytes (requirements.txt). The FIRST entry is the default at boot; the
-# minicpm5-1b entry stays selectable and still tracks the MINICPM_AGENT_MODEL_ID/
-# REVISION env overrides. Only one brain is resident at a time.
+# — tool routing passes it False; MiniCPM3 does not). The FIRST entry is the
+# default at boot; the minicpm5-1b entry stays selectable and still tracks the
+# MINICPM_AGENT_MODEL_ID/REVISION env overrides. Only one brain is resident at a time.
 AGENT_MODELS = [
     {
-        "key": "minicpm4.1-8b-8bit",
-        "label": "MiniCPM4.1 8B (8-bit)",
+        "key": "minicpm4.1-8b",
+        "label": "MiniCPM4.1 8B",
         "model_id": "openbmb/MiniCPM4.1-8B",
+        # DEFAULT brain. Hybrid-reasoning 8B in bf16 (~16 GiB) — the best eval
+        # config (0.90 tool / 0.90 args, and it unlocks the v3 coincidence fix).
+        # Loaded at import so ZeroGPU packs it (no per-turn reload) and it decodes
+        # in bf16 (faster than a bnb-int8 build). trust_remote_code custom modeling
+        # (sparse "InfLLM v2" attention); runs clean on current transformers, unlike
+        # MiniCPM3-4B — pin a reviewed commit (revision) before a real deploy.
         "revision": None,
-        # DEFAULT brain. Hybrid-reasoning 8B at int8 (~8.5 GiB) — the best
-        # deployable eval config (0.85 tool / 0.91 args; with the v3 prompt it also
-        # recovers half the coincidence fix). trust_remote_code custom modeling
-        # (sparse "InfLLM v2" attention); runs clean on current transformers,
-        # unlike MiniCPM3-4B. Quantized weights load straight onto the GPU, so
-        # device_map is set (skips the host→device .to copy).
         "thinking": True,
         "trust_remote_code": True,
-        "quantization": "8bit",
-        "device_map": {"": 0},
+        # Too large to keep the ColEmbed visual retriever resident alongside it, so
+        # visual search is disabled while this brain is active (falls back to parsed).
+        "forbid_visual": True,
     },
     {
         "key": "minicpm5-1b",
@@ -132,35 +133,6 @@ AGENT_MODELS = [
         # switch time with ~5 GiB to spare after the grounding spike. The 8B (16
         # GiB) didn't fit — see core/vram.py / the find-turn VRAM logs.
     },
-    {
-        "key": "minicpm4.1-8b",
-        "label": "MiniCPM4.1 8B",
-        "model_id": "openbmb/MiniCPM4.1-8B",
-        "revision": None,
-        # Hybrid-reasoning 8B (enable_thinking supported; routing passes it False).
-        # trust_remote_code custom modeling (sparse "InfLLM v2" attention) — same
-        # rot risk that broke MiniCPM3-4B on current transformers; pin a reviewed
-        # commit for any real use. NOTE: 8B / ~16 GiB bf16 does NOT fit the
-        # production find turn (the 8B didn't fit — VRAM logs); benchmarkable in
-        # THIS eval (brain-only load) to gauge the quality ceiling, but shipping it
-        # needs quantization (4-bit ~5-6 GiB) + a coexistence VRAM check.
-        "thinking": True,
-        "trust_remote_code": True,
-    },
-    {
-        "key": "minicpm4.1-8b-4bit",
-        "label": "MiniCPM4.1 8B (4-bit)",
-        "model_id": "openbmb/MiniCPM4.1-8B",
-        "revision": None,
-        "thinking": True,
-        "trust_remote_code": True,
-        # bitsandbytes nf4 (~5-6 GiB vs ~16 GiB bf16) to fit the 8B into the
-        # find-turn VRAM budget; quantized weights load straight onto the GPU, so
-        # device_map is set (skips the host→device .to copy). The question this
-        # answers: does 4-bit hold the bf16 8B's quality (0.87/0.86)?
-        "quantization": "4bit",
-        "device_map": {"": 0},
-    },
 ]
 DEFAULT_AGENT_MODEL = AGENT_MODELS[0]["key"]
 # A tool-call decision is short JSON; a rerank reply is a single number. 96 was

frontend/index.html

@@ -355,19 +355,23 @@
         <p class="mt-1 text-xs text-brand-400">The model that finds pages and points. Switching loads it fresh — the first turn after a change is slower.</p>
       </div>
 
-      <!-- search index: which index the search tool ranks the query against -->
-      <div class="mt-5">
+      <!-- search index: which index the search tool ranks the query against. The
+           big brains (forbidVisual) can't keep ColEmbed resident, so Visual is
+           disabled and force-reset to Parsed while one of them is the agent. -->
+      <div class="mt-5"
+           x-effect="if((agentModels.find(x=>x.key===agentModel)||{}).forbidVisual && retrievalMode==='visual') retrievalMode='parsed'">
         <label class="block text-xs font-semibold uppercase tracking-wide text-brand-500/80 mb-1.5">Search index</label>
         <div class="relative">
           <select x-model="retrievalMode"
             class="w-full appearance-none rounded-xl border border-brand-200 bg-brand-50/50 px-3.5 py-2.5 pr-9 text-sm font-medium text-navy focus:border-brand-400 focus:ring-2 focus:ring-brand-100 outline-none transition">
             <template x-for="m in retrievalModes" :key="m.key">
-              <option :value="m.key" x-text="m.label"></option>
+              <option :value="m.key" x-text="m.label"
+                :disabled="m.key==='visual' && (agentModels.find(x=>x.key===agentModel)||{}).forbidVisual"></option>
             </template>
           </select>
           <i data-lucide="chevron-down" class="pointer-events-none absolute right-3 top-1/2 -translate-y-1/2 w-4 h-4 text-brand-400"></i>
         </div>
-        <p class="mt-1 text-xs text-brand-400">What the search tool ranks against. Parsed (text) wins on specs and tables; Visual (ColEmbed) on diagrams.</p>
+        <p class="mt-1 text-xs text-brand-400">What the search tool ranks against. Parsed (text) wins on specs and tables; Visual (ColEmbed) on diagrams.<span x-show="(agentModels.find(x=>x.key===agentModel)||{}).forbidVisual" class="text-amber-600"> Visual is unavailable with the current brain (too large to load ColEmbed alongside it).</span></p>
       </div>
 
       <!-- think: let the VLM reason before committing to a circle box -->

models/colembed.py

@@ -1,12 +1,15 @@
 """Nemotron ColEmbed v2: late-interaction page embeddings + MaxSim retrieval.
 
-The model is a module-level global: ZeroGPU packs module-level CUDA tensors at
-startup and shares them with the GPU worker, whereas function arguments are
-pickled — and the trust_remote_code model class is not picklable.
+The model is a module-level global, but — unlike the other models — it is NOT
+built at import. It lazy-loads on first use (_load) so ZeroGPU does NOT pack it
+at startup, freeing ~8 GiB for the bf16 8B agent brain to stay resident. The
+trade: the FIRST visual-search call after a cold worker builds it inside the GPU
+window (not packed). The default search index is "parsed" (no ColEmbed) and the
+8B brain forbids visual search, so the production default path never loads it.
 
 _embed_pages_on_gpu is a ZeroGPU entry point (used at index time);
 maxsim_search is a plain function so the ask pipeline can run it inside its
-own single GPU call together with answer generation.
+own single GPU call together with answer generation. Both call _load() first.
 
 forward_images/forward_queries return zero-padded [batch, tokens, dim] tensors
 with real tokens L2-normalized, so padding rows are exactly zero. We strip them
@@ -33,45 +36,56 @@ from core.constants import (
 )
 from core.vram import log_vram
 
-_MODEL = (
-    AutoModel.from_pretrained(
-        COLEMBED_MODEL_ID,
-        revision=COLEMBED_REVISION,
-        trust_remote_code=True,
-        dtype=torch.bfloat16,
-        attn_implementation=COLEMBED_ATTN,
+_MODEL = None
+
+
+def _load():
+    """Build ColEmbed and cache it as the module global, on first use. Deliberately
+    NOT called at import: keeping it off the GPU at startup means ZeroGPU does not
+    pack it, freeing ~8 GiB so the bf16 8B agent brain fits the resident set. Must
+    run on GPU (called from within a @spaces.GPU context). A no-op once loaded."""
+    global _MODEL
+    if _MODEL is not None:
+        return _MODEL
+    model = (
+        AutoModel.from_pretrained(
+            COLEMBED_MODEL_ID,
+            revision=COLEMBED_REVISION,
+            trust_remote_code=True,
+            dtype=torch.bfloat16,
+            attn_implementation=COLEMBED_ATTN,
+        )
+        .to("cuda")
+        .eval()
     )
-    .to("cuda")
-    .eval()
-)
-# Pre-build the processor the remote code would otherwise lazily create per
-# GPU worker (it caches on this exact attribute, see _get_processor).
-_MODEL._processor = AutoProcessor.from_pretrained(
-    COLEMBED_MODEL_ID, revision=COLEMBED_REVISION, trust_remote_code=True
-)
-log_vram("load-colembed")
-
-# The remote code's forward_documents hardcodes DataLoader(num_workers=8), but
-# the ZeroGPU worker is a daemonic process and may not spawn children
-# ("daemonic processes are not allowed to have children"). Patch the DataLoader
-# name in the model's own module to force in-process loading. A subclass (not a
-# wrapper function) because the remote code also uses the name in isinstance().
-_remote_module = sys.modules[type(_MODEL).__module__]
-
-
-class _SingleProcessDataLoader(_remote_module.DataLoader):
-    def __init__(self, *args, **kwargs):
-        kwargs["num_workers"] = 0
-        super().__init__(*args, **kwargs)
+    # Pre-build the processor the remote code would otherwise lazily create per
+    # GPU worker (it caches on this exact attribute, see _get_processor).
+    model._processor = AutoProcessor.from_pretrained(
+        COLEMBED_MODEL_ID, revision=COLEMBED_REVISION, trust_remote_code=True
+    )
+    # The remote code's forward_documents hardcodes DataLoader(num_workers=8), but
+    # the ZeroGPU worker is a daemonic process and may not spawn children
+    # ("daemonic processes are not allowed to have children"). Patch the DataLoader
+    # name in the model's own module to force in-process loading. A subclass (not a
+    # wrapper function) because the remote code also uses the name in isinstance().
+    remote_module = sys.modules[type(model).__module__]
 
+    class _SingleProcessDataLoader(remote_module.DataLoader):
+        def __init__(self, *args, **kwargs):
+            kwargs["num_workers"] = 0
+            super().__init__(*args, **kwargs)
 
-_remote_module.DataLoader = _SingleProcessDataLoader
+    remote_module.DataLoader = _SingleProcessDataLoader
+    _MODEL = model
+    log_vram("load-colembed")
+    return _MODEL
 
 
 @spaces.GPU(duration=EMBED_GPU_DURATION)
 def _embed_pages_on_gpu(images: list[Image.Image]) -> list[np.ndarray]:
+    model = _load()
     with torch.no_grad():
-        embs = _MODEL.forward_images(images, batch_size=EMBED_BATCH_SIZE)
+        embs = model.forward_images(images, batch_size=EMBED_BATCH_SIZE)
     out = []
     for emb in embs:  # [tokens, dim]; zero rows are padding
         mask = emb.abs().sum(dim=-1) > 0
@@ -85,8 +99,9 @@ def maxsim_search(
     """Top-K (doc_id, page_num, score) across docs. Must run on GPU (called
     from within a @spaces.GPU context)."""
     results = []
+    model = _load()
     with torch.no_grad():
-        q = _MODEL.forward_queries([question], batch_size=1)[0].to(torch.float16)
+        q = model.forward_queries([question], batch_size=1)[0].to(torch.float16)
         for refs, batch in store.iter_page_batches(doc_ids, SCORE_PAGES_PER_BATCH):
             emb = torch.from_numpy(batch).to(q.device)  # [B, T, D] float16
             sim = torch.einsum("qd,btd->bqt", q, emb).float()

models/minicpm_agent.py

@@ -328,25 +328,6 @@ def use_model(key: str | None = None) -> str:
     device_map = spec.get("device_map")
     if device_map is not None:
         load_kwargs["device_map"] = device_map
-    # A spec can request on-the-fly bitsandbytes quantization (e.g. "4bit" to fit
-    # an 8B brain into the find-turn VRAM budget). Quantized weights are placed on
-    # the GPU at load time, so such a spec must also set device_map (skips the
-    # .to("cuda") below). bitsandbytes is imported lazily so non-quantized brains
-    # — and the production image, which omits it — never touch it.
-    quant = spec.get("quantization")
-    if quant in ("4bit", "8bit"):
-        from transformers import BitsAndBytesConfig
-
-        load_kwargs["quantization_config"] = (
-            BitsAndBytesConfig(
-                load_in_4bit=True,
-                bnb_4bit_quant_type="nf4",
-                bnb_4bit_compute_dtype=torch.bfloat16,
-                bnb_4bit_use_double_quant=True,
-            )
-            if quant == "4bit"
-            else BitsAndBytesConfig(load_in_8bit=True)
-        )
     model = AutoModelForCausalLM.from_pretrained(spec["model_id"], **load_kwargs)
     if device_map is None:
         model = model.to("cuda")
@@ -357,17 +338,11 @@ def use_model(key: str | None = None) -> str:
 
 
 # Load the default brain eagerly at import so ZeroGPU's startup tensor-packing
-# covers it and the common (no-switch) first turn pays no load cost. EXCEPTION: a
-# bitsandbytes-quantized default must NOT be built in the main process. Plain
-# .to("cuda") models are safe at import because the `spaces` library patches torch
-# and "packs" them into the forked GPU worker — but bitsandbytes initializes CUDA
-# directly (bypassing that patch), which hard-errors on ZeroGPU ("CUDA must not be
-# initialized in the main process") and crashes the Space at boot. So for a
-# quantized default, DEFER the load to first GPU use: the pipeline calls use_model()
-# inside its @spaces.GPU window (pipelines/agent_ask.py), and _generate() lazy-loads
-# as a backstop — both on the GPU, the only supported place to build a bnb model.
-if not _spec(DEFAULT_AGENT_MODEL).get("quantization"):
-    use_model(DEFAULT_AGENT_MODEL)
+# covers it and the common (no-switch) turn pays no per-turn load cost. A plain
+# .to("cuda") model is safe at import on ZeroGPU: the `spaces` library patches
+# torch and "packs" it into the forked GPU worker (the bf16 8B default included),
+# so it stays resident without ever being rebuilt per turn.
+use_model(DEFAULT_AGENT_MODEL)
 
 
 def _template_kwargs() -> dict:
@@ -383,10 +358,9 @@ def _generate(
     """Greedy decode the assistant's next message. Traced as one `generation`
     (the resident brain as the model, the messages as input, the reply and the
     in/out token counts attached) when Langfuse is configured."""
-    # Backstop for a deferred (quantized) default brain: it isn't loaded at import
-    # on ZeroGPU, so build it on first use. Always reached inside a @spaces.GPU
-    # window (the pipeline's find turn / the eval's GPU fn) — the supported place to
-    # instantiate a bitsandbytes model. A no-op once a brain is resident.
+    # Defensive: ensure a brain is resident. The default loads at import, so this
+    # is a no-op in normal operation; it only fires if that eager load was skipped
+    # (e.g. a future deferred default). Always reached inside a @spaces.GPU window.
     if _MODEL is None:
         use_model(DEFAULT_AGENT_MODEL)
     with tracing.generation(

requirements.txt

@@ -2,7 +2,6 @@ spaces
 gradio
 transformers>=4.57.2,<5
 accelerate
-bitsandbytes  # int8/4bit agent-brain quantization (default brain is int8 MiniCPM4.1-8B)
 torchvision
 pymupdf
 pillow