feat: add CUDA SAMVG renderer

.gitignore

@@ -6,6 +6,7 @@
 .venv/
 .idea/
 *.egg-info
+/build/
 /output
 /models
 # Run output: vectrify writes <output-stem>/runs/ next to the output file

pyproject.toml

@@ -57,14 +57,24 @@ vision = [
     "torchvision>=0.28.0",
     "transformers>=4.40.0",
 ]
+# The CUDA extension is shipped in platform-specific wheels.  Installing this
+# extra is deliberately sufficient for the SAMVG seed even on machines that
+# use the portable Torch renderer fallback.
+samvg = [
+    "scipy>=1.11.0",
+    "scikit-learn>=1.3.0",
+    "torch>=2.0.0",
+    "torchvision>=0.28.0",
+    "transformers>=4.40.0",
+]
 graphviz = [
     "graphviz>=0.21",
 ]
 typst = [
     "typst>=0.11.0",
 ]
 all = [
-    "vectrify[vision,graphviz,typst]",
+    "vectrify[vision,samvg,graphviz,typst]",
 ]
 dev = [
     "pytest",
@@ -96,6 +106,9 @@ package-dir = { "" = "src" }
 [tool.setuptools.packages.find]
 where = ["src"]
 
+[tool.setuptools.package-data]
+vectrify = ["refine/*.cu"]
+
 [tool.pytest.ini_options]
 addopts = "-m 'not llm'"
 testpaths = ["tests"]

scripts/bench_samvg_renderer.py

@@ -0,0 +1,54 @@
+"""Measure one steady SAMVG filled-path optimisation step.
+
+Example:
+    uv run python scripts/bench_samvg_renderer.py /tmp/cat.svg /tmp/cat.jpg
+    uv run python scripts/bench_samvg_renderer.py /tmp/cat.svg /tmp/cat.jpg \
+        --torch-fallback
+"""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+from time import perf_counter
+
+from PIL import Image
+
+from vectrify.refine.paths import fit_filled_svg
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("svg", type=Path)
+    parser.add_argument("target", type=Path)
+    parser.add_argument("--steps", type=int, default=3)
+    parser.add_argument("--long-side", type=int, default=64)
+    parser.add_argument("--torch-fallback", action="store_true")
+    args = parser.parse_args()
+    if args.steps < 1:
+        raise ValueError("--steps must be at least one")
+
+    import torch
+
+    if args.torch_fallback:
+        from vectrify.refine import cuda_renderer
+
+        cuda_renderer._extension = lambda: None
+    svg = args.svg.read_text()
+    target = Image.open(args.target)
+    # Setup, CUDA allocator warm-up, and any Torch compilation happen outside
+    # the timed region so the result is a steady optimisation step.
+    fit_filled_svg(svg, target, steps=1, optimisation_long_side=args.long_side)
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+    started = perf_counter()
+    fit_filled_svg(
+        svg, target, steps=args.steps, optimisation_long_side=args.long_side
+    )
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+    print(f"{(perf_counter() - started) / args.steps:.6f} seconds/step")
+
+
+if __name__ == "__main__":
+    main()

setup.py

@@ -0,0 +1,29 @@
+"""Build the optional SAMVG CUDA extension for release wheels.
+
+Normal source installs intentionally remain pure Python.  Release builders
+set ``VECTRIFY_BUILD_SAMVG_CUDA=1`` after installing the matching CUDA Torch
+wheel; the resulting wheel bundles ``vectrify._samvg_cuda``.
+"""
+
+from __future__ import annotations
+
+import os
+
+from setuptools import setup
+
+
+def cuda_extension():
+    if os.environ.get("VECTRIFY_BUILD_SAMVG_CUDA") != "1":
+        return [], {}
+    from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+    extension = CUDAExtension(
+        "vectrify._samvg_cuda",
+        ["src/vectrify/refine/_samvg_cuda.cu"],
+        extra_compile_args={"cxx": ["-O3"], "nvcc": ["-O3"]},
+    )
+    return [extension], {"build_ext": BuildExtension}
+
+
+ext_modules, cmdclass = cuda_extension()
+setup(ext_modules=ext_modules, cmdclass=cmdclass)

src/vectrify/formats/svg/operations.py

@@ -759,9 +759,7 @@ def apply_mutation(
     fn, name = pick_operator(MUTATIONS, operator)
 
     def run() -> str:
-        targeted_fn = cast(
-            Callable[[str, Mapping[int, float] | None], str], fn
-        )
+        targeted_fn = cast(Callable[[str, Mapping[int, float] | None], str], fn)
         return targeted_fn(parent_svg, targets)
 
     return with_retries(run, fallback=parent_svg), name

src/vectrify/image_utils.py

@@ -29,9 +29,7 @@ def crop_single_color_background(
         return image
 
     rgb = np.asarray(image.convert("RGB"), dtype=np.int16)
-    corners = np.array(
-        [rgb[0, 0], rgb[0, -1], rgb[-1, 0], rgb[-1, -1]], dtype=np.int16
-    )
+    corners = np.array([rgb[0, 0], rgb[0, -1], rgb[-1, 0], rgb[-1, -1]], dtype=np.int16)
     background = np.median(corners, axis=0)
     if np.max(np.abs(corners - background)) > tolerance:
         return image