From 048a9ae35a4210fc3140fb0acd3bc83fbe139592 Mon Sep 17 00:00:00 2001 From: Rasmus Ros Date: Wed, 26 Aug 2026 08:46:26 +0200 Subject: [PATCH] fix: make samvg runtime portable --- .github/workflows/build.yml | 2 +- .github/workflows/release.yml | 2 +- README.md | 6 +++- src/vectrify/refine/samvg.py | 32 ++++---------------- src/vectrify/refine/samvg_runtime.py | 25 ++++++++++++++++ src/vectrify/refine/samvg_types.py | 30 +++++++++++++++++++ tests/refine/test_samvg.py | 45 ++++++++++++++++++++++++++++ 7 files changed, 112 insertions(+), 30 deletions(-) create mode 100644 src/vectrify/refine/samvg_runtime.py create mode 100644 src/vectrify/refine/samvg_types.py diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index e68792ea..4a0ec8c5 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -50,7 +50,7 @@ jobs: run: uv python install 3.13 - name: Sync dependencies - run: uv sync --extra dev --extra graphviz --extra typst --extra vision + run: uv sync --extra dev --extra graphviz --extra typst --extra vision --extra samvg - name: Lint run: uv run ruff check src/ tests/ scripts/ diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index bcdff849..aca270a1 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -35,7 +35,7 @@ jobs: run: uv python install 3.13 - name: Sync dependencies - run: uv sync --extra dev --extra graphviz --extra typst + run: uv sync --extra dev --extra graphviz --extra typst --extra samvg - name: Run tests run: uv run pytest -q diff --git a/README.md b/README.md index 05169791..d1331492 100644 --- a/README.md +++ b/README.md @@ -132,7 +132,8 @@ it does not run the path optimiser. This keeps seed quality measurable without mixing in local refinement. For the dissertation-style two-phase measurement (initial fit, residual prompts, -and recovery fit), build with the optional native CUDA renderer and run: +and recovery fit), build a local wheel with the optional native CUDA renderer +and run: ```sh VECTRIFY_BUILD_SAMVG_CUDA=1 uv build --wheel --no-build-isolation @@ -143,3 +144,6 @@ uv pip install --force-reinstall --no-deps dist/vectrify-*.whl `--all` also evaluates every benchmark target and the connect-the-dots duck. Each target directory contains the five stage rasters, SVGs, a gallery, pixel-error table, and per-bounded-group CUDA memory/timing data. + +PyPI releases are portable Python wheels and use the Torch renderer fallback. +They do not currently bundle the optional CUDA extension. diff --git a/src/vectrify/refine/samvg.py b/src/vectrify/refine/samvg.py index 7e22b4c3..6c730530 100644 --- a/src/vectrify/refine/samvg.py +++ b/src/vectrify/refine/samvg.py @@ -25,6 +25,9 @@ import numpy as np from PIL import Image +from vectrify.refine.samvg_runtime import device_name, pipeline_options +from vectrify.refine.samvg_types import MaskLayer, TextLayer + log = logging.getLogger(__name__) # SAMVG's quality depends directly on the granularity of its automatic masks. @@ -66,29 +69,6 @@ OCR_TEXT_RMSE_TOLERANCE = 0.02 -@dataclass(frozen=True) -class MaskLayer: - """One painted segmentation mask, in document compositing order.""" - - mask: np.ndarray - colour: tuple[int, int, int] - impact: float - overlap_pixels: int = 0 - - -@dataclass(frozen=True) -class TextLayer: - """A high-confidence OCR word represented as editable SVG text.""" - - text: str - x: float - y: float - width: float - height: float - colour: tuple[int, int, int] - angle: float = 0.0 - - def _text_colour(pixels: np.ndarray) -> tuple[int, int, int]: """Estimate ink colour by contrasting a word crop with its border.""" height, width, _channels = pixels.shape @@ -478,9 +458,7 @@ def _sam_runtime(*, model: str = SAMVG_MODEL) -> _SamRuntime: raise ImportError( "SAMVG requires the samvg extra. Install 'vectrify[samvg]'." ) from exc - options: dict[str, Any] = {"model": model, "device": 0} - if torch.cuda.is_available(): - options["dtype"] = torch.float16 + options = pipeline_options(torch, model) generator = pipeline("mask-generation", **options) log.info( "SAMVG automatic masks: %s on %s (%s).", @@ -1042,7 +1020,7 @@ def prompted_masks( points = [(round(x * scale), round(y * scale)) for x, y in points] own_runtime = _runtime is None runtime = _runtime or _sam_runtime() - device = "cuda" if torch.cuda.is_available() else "cpu" + device = device_name(torch) log.info("SAMVG prompted masks: using %s.", device) if runtime.processor is None: runtime.processor = SamProcessor(runtime.generator.image_processor) diff --git a/src/vectrify/refine/samvg_runtime.py b/src/vectrify/refine/samvg_runtime.py new file mode 100644 index 00000000..d3ce0bb0 --- /dev/null +++ b/src/vectrify/refine/samvg_runtime.py @@ -0,0 +1,25 @@ +"""Small, dependency-light runtime helpers for SAMVG model loading. + +Keeping device policy here lets the segmentation pipeline be tested without +loading a checkpoint and prevents CUDA assumptions leaking into CPU installs. +""" + +from __future__ import annotations + +from typing import Any + + +def pipeline_options(torch: Any, model: str) -> dict[str, Any]: + """Return Transformers pipeline options for the available Torch device.""" + options: dict[str, Any] = {"model": model} + if torch.cuda.is_available(): + options.update(device=0, dtype=torch.float16) + else: + # Transformers uses -1 for CPU; device=0 explicitly selects cuda:0. + options["device"] = -1 + return options + + +def device_name(torch: Any) -> str: + """Return the device name used by direct SAM prompt decoding.""" + return "cuda" if torch.cuda.is_available() else "cpu" diff --git a/src/vectrify/refine/samvg_types.py b/src/vectrify/refine/samvg_types.py new file mode 100644 index 00000000..27018b7c --- /dev/null +++ b/src/vectrify/refine/samvg_types.py @@ -0,0 +1,30 @@ +"""Shared value objects used by SAMVG's segmentation and SVG stages.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + + +@dataclass(frozen=True) +class MaskLayer: + """One painted segmentation mask, in document compositing order.""" + + mask: np.ndarray + colour: tuple[int, int, int] + impact: float + overlap_pixels: int = 0 + + +@dataclass(frozen=True) +class TextLayer: + """A high-confidence OCR word represented as editable SVG text.""" + + text: str + x: float + y: float + width: float + height: float + colour: tuple[int, int, int] + angle: float = 0.0 diff --git a/tests/refine/test_samvg.py b/tests/refine/test_samvg.py index 459ffb94..95bc6f71 100644 --- a/tests/refine/test_samvg.py +++ b/tests/refine/test_samvg.py @@ -32,6 +32,51 @@ recolour_visible_layers, residual_prompt_points, ) +from vectrify.refine.samvg_runtime import device_name, pipeline_options + + +def test_sam_runtime_uses_cpu_pipeline_without_cuda(): + torch = SimpleNamespace( + float16="fp16", cuda=SimpleNamespace(is_available=lambda: False) + ) + + assert pipeline_options(torch, "sam") == {"model": "sam", "device": -1} + assert device_name(torch) == "cpu" + + +def test_sam_runtime_uses_half_precision_cuda_pipeline(): + torch = SimpleNamespace( + float16="fp16", cuda=SimpleNamespace(is_available=lambda: True) + ) + + assert pipeline_options(torch, "sam") == { + "model": "sam", + "device": 0, + "dtype": "fp16", + } + assert device_name(torch) == "cuda" + + +def test_sam_runtime_loads_transformers_on_cpu(monkeypatch): + calls = [] + generator = SimpleNamespace(device="cpu") + monkeypatch.setitem( + sys.modules, + "torch", + SimpleNamespace( + float16="fp16", cuda=SimpleNamespace(is_available=lambda: False) + ), + ) + monkeypatch.setitem( + sys.modules, + "transformers", + SimpleNamespace( + pipeline=lambda name, **kwargs: calls.append((name, kwargs)) or generator + ), + ) + + assert samvg._sam_runtime(model="example/sam").generator is generator + assert calls == [("mask-generation", {"model": "example/sam", "device": -1})] def test_detect_text_retains_high_confidence_editable_words(monkeypatch):