Enable CUDA llama-cpp in ZeroGPU vision workers

Vision was still running on CPU because requirements.txt pinned the CPU-only
llama-cpp-python wheel. ZeroGPU workers had n_gpu_layers=-1 but no CUDA
backend to honor it, which explains ~62s image decode times.

Switch the Space runtime to the cu124 prebuilt wheel so @spaces.GPU workers
can offload the GGUF. Keep the main Gradio process on n_gpu_layers=0 for
text replies; lazy import still defers model load until first use. Route
single-image extract() through the GPU worker as well, and reload the
backend when switching between CPU and GPU load modes in the same instance.

Adds unit tests for _n_gpu_layers and GPU routing of extract().

requirements.txt

@@ -1,10 +1,12 @@
 # Hugging Face Spaces installs this file with pip (Spaces do not use uv).
 # Keep it in sync with pyproject.toml: `make lock-requirements`.
 # Runtime dependencies only; dev tools (ruff, mypy, pytest) live in pyproject.
-# CPU wheel: CUDA wheels need libcudart at import time, which ZeroGPU lacks.
+# CUDA wheel: vision runs in @spaces.GPU workers with n_gpu_layers=-1. The main
+# Gradio process still uses n_gpu_layers=0 for text; lazy import avoids loading
+# the model until first use. HF ZeroGPU images ship libcudart for both paths.
 --prefer-binary
 --only-binary llama-cpp-python
---extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cpu
+--extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cu124
 llama-cpp-python>=0.3.26
 spaces
 huggingface_hub

src/vivamais/adapters/models/llama_cpp.py

@@ -108,7 +108,7 @@ def _image_mime(data: bytes) -> str:
 
 
 def _import_llama() -> tuple[Any, Any]:
-    """Import llama_cpp lazily so Space startup does not load CUDA libs."""
+    """Import llama_cpp lazily so Space startup does not load the GGUF."""
     from llama_cpp import Llama
     from llama_cpp.llama_chat_format import MiniCPMv26ChatHandler
 
@@ -116,7 +116,7 @@ def _import_llama() -> tuple[Any, Any]:
 
 
 def _n_gpu_layers(*, use_gpu: bool = False) -> int:
-    """CPU in the Gradio process; offload all layers inside ZeroGPU workers."""
+    """CPU in the Gradio process; full offload inside ZeroGPU workers."""
     env_val = os.environ.get("VIVAMAIS_N_GPU_LAYERS")
     if env_val is not None:
         return int(env_val)
@@ -139,6 +139,7 @@ class LlamaCppBackend:
         self._n_ctx = n_ctx or (int(env_ctx) if env_ctx else 4096)
         self._verbose = verbose
         self._llama: Any | None = None
+        self._loaded_use_gpu: bool | None = None
         self._last_raw_output: str = ""
 
     def __del__(self) -> None:
@@ -165,8 +166,11 @@ class LlamaCppBackend:
         )
 
     def _ensure_loaded(self, use_gpu: bool = False) -> None:
-        if self._llama is not None:
+        if self._llama is not None and self._loaded_use_gpu == use_gpu:
             return
+        if self._llama is not None:
+            del self._llama
+            self._llama = None
         model_path = _resolve_path(
             self._model_path,
             "VIVAMAIS_MODEL_PATH",
@@ -181,13 +185,17 @@ class LlamaCppBackend:
         )
         Llama, MiniCPMv26ChatHandler = _import_llama()
         handler = MiniCPMv26ChatHandler(clip_model_path=mmproj_path, verbose=self._verbose)
+        gpu_layers = _n_gpu_layers(use_gpu=use_gpu)
         self._llama = Llama(
             model_path=model_path,
             n_ctx=self._n_ctx,
-            n_gpu_layers=_n_gpu_layers(use_gpu=use_gpu),
+            n_gpu_layers=gpu_layers,
             verbose=self._verbose,
             chat_handler=handler,
         )
+        self._loaded_use_gpu = use_gpu
+        if self._verbose:
+            print(f"llamacpp loaded n_gpu_layers={gpu_layers} use_gpu={use_gpu}", flush=True)
 
     def _build_messages(self, image: bytes, prompt: str) -> list[dict[str, Any]]:
         image_b64 = base64.b64encode(image).decode("utf-8")
@@ -241,7 +249,13 @@ class LlamaCppBackend:
         self._last_raw_output = text
         return text
 
-    def extract(self, image: bytes, schema_hint: str) -> dict[str, Any]:
+    def extract(
+        self,
+        image: bytes,
+        schema_hint: str,
+        *,
+        use_gpu: bool = False,
+    ) -> dict[str, Any]:
         if schema_hint == "ticket":
             prompt = TICKET_PROMPT
         elif schema_hint == "document":
@@ -249,7 +263,10 @@ class LlamaCppBackend:
         else:
             prompt = RECEIPT_PROMPT
         messages = self._build_messages(image, prompt)
-        text = self._completion(messages, max_tokens=384, temperature=0.0)
+        if use_gpu:
+            text = self._completion_gpu(messages, max_tokens=384, temperature=0.0)
+        else:
+            text = self._completion(messages, max_tokens=384, temperature=0.0)
         return _parse_json(text, schema_hint)
 
     def extract_all(
@@ -315,7 +332,12 @@ def _gpu_extract_impl(image: bytes) -> tuple[dict[str, Any], dict[str, Any]]:
     return _worker_backend_instance().extract_all(image, use_gpu=True)
 
 
+def _gpu_extract_schema_impl(image: bytes, schema_hint: str) -> dict[str, Any]:
+    return _worker_backend_instance().extract(image, schema_hint, use_gpu=True)
+
+
 _gpu_extract_one: Callable[[bytes], tuple[dict[str, Any], dict[str, Any]]] | None = None
+_gpu_extract_schema: Callable[[bytes, str], dict[str, Any]] | None = None
 
 if spaces is not None:
 
@@ -323,6 +345,10 @@ if spaces is not None:
     def _gpu_extract_one(image: bytes) -> tuple[dict[str, Any], dict[str, Any]]:
         return _gpu_extract_impl(image)
 
+    @spaces.GPU(duration=_single_image_duration)  # type: ignore[misc]
+    def _gpu_extract_schema(image: bytes, schema_hint: str) -> dict[str, Any]:
+        return _gpu_extract_schema_impl(image, schema_hint)
+
 
 def _run_gpu_images(
     images: list[bytes],
@@ -344,6 +370,8 @@ class LlamaCppVisionModel:
         self._backend = backend
 
     def extract(self, image: bytes, schema_hint: str) -> dict[str, Any]:
+        if spaces is not None and _gpu_extract_schema is not None:
+            return _gpu_extract_schema(image, schema_hint)
         return self._backend.extract(image, schema_hint)
 
     def extract_all(self, image: bytes) -> tuple[dict[str, Any], dict[str, Any]]:

tests/unit/adapters/test_llama_cpp.py

@@ -15,6 +15,7 @@ from vivamais.adapters.models.llama_cpp import (
     VARIANT,
     LlamaCppBackend,
     LlamaCppVisionModel,
+    _n_gpu_layers,
     _parse_json,
     _resolve_path,
     _single_image_duration,
@@ -66,6 +67,19 @@ class TestSingleImageDuration:
         )
 
 
+class TestNGpuLayers:
+    def test_cpu_path_uses_zero_layers(self) -> None:
+        assert _n_gpu_layers(use_gpu=False) == 0
+
+    def test_gpu_worker_offloads_all_layers(self) -> None:
+        assert _n_gpu_layers(use_gpu=True) == -1
+
+    def test_env_override(self) -> None:
+        with patch.dict(os.environ, {"VIVAMAIS_N_GPU_LAYERS": "32"}, clear=False):
+            assert _n_gpu_layers(use_gpu=True) == 32
+            assert _n_gpu_layers(use_gpu=False) == 32
+
+
 class TestLlamaCppVisionModelBatchProcess:
     def _make_mock_backend(self) -> MagicMock:
         backend = MagicMock()
@@ -102,6 +116,17 @@ class TestLlamaCppVisionModelBatchProcess:
         gpu_fn.assert_any_call(b"img1")
         gpu_fn.assert_any_call(b"img2")
 
+    def test_extract_routes_through_gpu_when_spaces_available(self) -> None:
+        backend = self._make_mock_backend()
+        model = LlamaCppVisionModel(backend)
+        gpu_fn = MagicMock(return_value={"passenger": "Maria"})
+        with patch("vivamais.adapters.models.llama_cpp.spaces", MagicMock()):
+            with patch("vivamais.adapters.models.llama_cpp._gpu_extract_schema", gpu_fn):
+                result = model.extract(b"img", "document")
+        assert result == {"passenger": "Maria"}
+        gpu_fn.assert_called_once_with(b"img", "document")
+        backend.extract.assert_not_called()
+
     def test_stream_process_yields_per_image_on_gpu(self) -> None:
         backend = self._make_mock_backend()
         model = LlamaCppVisionModel(backend)