From e43941754cc0bb75c2274b7221a3bbe5799c37cd Mon Sep 17 00:00:00 2001 From: Rasmus Ros Date: Mon, 24 Aug 2026 06:22:15 +0200 Subject: [PATCH 1/6] feat: add CUDA SAMVG renderer --- pyproject.toml | 15 +- scripts/bench_samvg_renderer.py | 54 ++ setup.py | 29 ++ src/vectrify/refine/_samvg_cuda.cu | 135 +++++ src/vectrify/refine/cuda_renderer.py | 70 +++ src/vectrify/refine/paths.py | 741 ++++++++++++++++++++++++--- src/vectrify/refine/samvg.py | 63 ++- tests/refine/test_filled_paths.py | 139 ++++- tests/refine/test_samvg.py | 10 + uv.lock | 19 +- 10 files changed, 1171 insertions(+), 104 deletions(-) create mode 100644 scripts/bench_samvg_renderer.py create mode 100644 setup.py create mode 100644 src/vectrify/refine/_samvg_cuda.cu create mode 100644 src/vectrify/refine/cuda_renderer.py diff --git a/pyproject.toml b/pyproject.toml index 352eb52f..d0e9d192 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,6 +57,16 @@ 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", ] @@ -64,7 +74,7 @@ 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"] diff --git a/scripts/bench_samvg_renderer.py b/scripts/bench_samvg_renderer.py new file mode 100644 index 00000000..09bb42ae --- /dev/null +++ b/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() diff --git a/setup.py b/setup.py new file mode 100644 index 00000000..3ebd77a6 --- /dev/null +++ b/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) diff --git a/src/vectrify/refine/_samvg_cuda.cu b/src/vectrify/refine/_samvg_cuda.cu new file mode 100644 index 00000000..61e6e3e7 --- /dev/null +++ b/src/vectrify/refine/_samvg_cuda.cu @@ -0,0 +1,135 @@ +#include +#include +#include + +namespace { +constexpr int kCubics = 16; +constexpr int kSamples = 32; + +__device__ inline void point(const float* control, int cubic, int sample, int samples, float& x, float& y) { + const float t = float(sample) / float(samples - 1), u = 1.f - t; + const float* c = control + cubic * 8; + x = u*u*u*c[0] + 3*u*u*t*c[2] + 3*u*t*t*c[4] + t*t*t*c[6]; + y = u*u*u*c[1] + 3*u*u*t*c[3] + 3*u*t*t*c[5] + t*t*t*c[7]; +} + +__device__ inline void sample_path(const float* path, float* points, int samples) { + for (int index = threadIdx.x; index < kCubics * samples; index += blockDim.x) { + point(path, index / samples, index % samples, samples, points[index * 2], points[index * 2 + 1]); + } + __syncthreads(); +} + +__global__ void forward_kernel(const float* controls, float* output, int batches, + int height, int width, int samples, float xo, float yo) { + const int pixels = height * width; + const int batch = blockIdx.x; + if (batch >= batches) return; + const float* path = controls + batch * kCubics * 8; + __shared__ float points[kCubics * kSamples * 2]; + sample_path(path, points, samples); + for (int pixel = threadIdx.x; pixel < pixels; pixel += blockDim.x) { + const float px = xo + float(pixel % width), py = yo + float(pixel / width); + float winding = 0.f; + for (int edge = 0; edge < kCubics * samples; ++edge) { + const int next = (edge + 1) % (kCubics * samples); + const float ax = points[edge * 2] - px, ay = points[edge * 2 + 1] - py; + const float bx = points[next * 2] - px, by = points[next * 2 + 1] - py; + winding += atan2f(ax * by - ay * bx, ax * bx + ay * by); + } + output[batch * pixels + pixel] = winding; + } +} + +__global__ void backward_kernel(const float* controls, const float* upstream, + float* gradients, int batches, int height, int width, + int samples, float xo, float yo) { + const int pixels = height * width; + const int batch = blockIdx.x; + if (batch >= batches) return; + const float* path = controls + batch * kCubics * 8; + float* gradient = gradients + batch * kCubics * 8; + __shared__ float points[kCubics * kSamples * 2]; + __shared__ float reduction[8][256]; + sample_path(path, points, samples); + float accumulated[kCubics * 8] = {0.f}; + for (int pixel = threadIdx.x; pixel < pixels; pixel += blockDim.x) { + const float px = xo + float(pixel % width), py = yo + float(pixel / width); + const float d_winding = upstream[batch * pixels + pixel]; + for (int edge = 0; edge < kCubics * samples; ++edge) { + const int next = (edge + 1) % (kCubics * samples); + const float ax = points[edge * 2] - px, ay = points[edge * 2 + 1] - py; + const float bx = points[next * 2] - px, by = points[next * 2 + 1] - py; + const float cross = ax * by - ay * bx, dot = ax * bx + ay * by; + const float scale = d_winding / (cross * cross + dot * dot + 1e-20f); + const float dcross = scale * dot, ddot = -scale * cross; + const float gx_a = dcross * by + ddot * bx; + const float gy_a = -dcross * bx + ddot * by; + const float gx_b = -dcross * ay + ddot * ax; + const float gy_b = dcross * ax + ddot * ay; + const float ta = float(edge % samples) / float(samples - 1), ua = 1.f - ta; + const float tb = float(next % samples) / float(samples - 1), ub = 1.f - tb; + const float ba[4] = {ua*ua*ua, 3*ua*ua*ta, 3*ua*ta*ta, ta*ta*ta}; + const float bb[4] = {ub*ub*ub, 3*ub*ub*tb, 3*ub*tb*tb, tb*tb*tb}; + for (int control = 0; control < 4; ++control) { + const int a = (edge / samples) * 8 + control * 2; + const int b = (next / samples) * 8 + control * 2; + accumulated[a] += ba[control] * gx_a; + accumulated[a + 1] += ba[control] * gy_a; + accumulated[b] += bb[control] * gx_b; + accumulated[b + 1] += bb[control] * gy_b; + } + } + } + for (int cubic = 0; cubic < kCubics; ++cubic) { + for (int component = 0; component < 8; ++component) { + reduction[component][threadIdx.x] = accumulated[cubic * 8 + component]; + } + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride /= 2) { + if (threadIdx.x < stride) { + for (int component = 0; component < 8; ++component) { + reduction[component][threadIdx.x] += reduction[component][threadIdx.x + stride]; + } + } + __syncthreads(); + } + if (threadIdx.x == 0) { + for (int component = 0; component < 8; ++component) { + gradient[cubic * 8 + component] = reduction[component][0]; + } + } + __syncthreads(); + } +} +} + +torch::Tensor forward(torch::Tensor controls, int64_t height, int64_t width, int64_t samples, + double x_origin, double y_origin) { + TORCH_CHECK(controls.is_cuda() && controls.scalar_type() == torch::kFloat32); + at::cuda::CUDAGuard guard(controls.device()); + auto output = torch::zeros({controls.size(0), height, width}, controls.options()); + constexpr int threads = 256; + forward_kernel<<>>( + controls.data_ptr(), output.data_ptr(), controls.size(0), height, width, samples, + float(x_origin), float(y_origin)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return output; +} + +torch::Tensor backward(torch::Tensor controls, torch::Tensor upstream, int64_t height, + int64_t width, int64_t samples, double x_origin, double y_origin) { + at::cuda::CUDAGuard guard(controls.device()); + auto gradients = torch::zeros_like(controls); + constexpr int threads = 256; + backward_kernel<<>>( + controls.data_ptr(), upstream.data_ptr(), gradients.data_ptr(), + controls.size(0), height, width, samples, float(x_origin), float(y_origin)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return gradients; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("forward", &forward); + m.def("backward", &backward); +} diff --git a/src/vectrify/refine/cuda_renderer.py b/src/vectrify/refine/cuda_renderer.py new file mode 100644 index 00000000..b5179b6e --- /dev/null +++ b/src/vectrify/refine/cuda_renderer.py @@ -0,0 +1,70 @@ +"""Optional CUDA winding operator for SAMVG's fixed 16-cubic contours.""" + +from __future__ import annotations + +import importlib +from functools import lru_cache +from typing import Any + + +@lru_cache(maxsize=1) +def _extension() -> Any | None: + """Load the ahead-of-time extension when the installed wheel contains it.""" + try: + return importlib.import_module("vectrify._samvg_cuda") + except ImportError: + return None + + +def available() -> bool: + """Whether this installation can execute the fixed-contour CUDA path.""" + return _extension() is not None + + +def winding( + controls: Any, + box: tuple[int, int, int, int], + *, + samples: int, + x_offset: float, + y_offset: float, +) -> Any | None: + """Return a differentiable native winding field, otherwise ``None``.""" + import torch + + extension = _extension() + if ( + extension is None + or not controls.is_cuda + or controls.dtype != torch.float32 + or controls.ndim != 4 + or controls.shape[1:] != (16, 4, 2) + or samples not in {8, 16, 32} + ): + return None + left, top, right, bottom = box + height, width = bottom - top, right - left + + class Winding(torch.autograd.Function): + @staticmethod + def forward(ctx, value): + value = value.contiguous() + ctx.save_for_backward(value) + return extension.forward( + value, height, width, samples, left + x_offset, top + y_offset + ) + + @staticmethod + def backward(ctx, upstream): + (value,) = ctx.saved_tensors + return extension.backward( + value, + upstream.contiguous(), + height, + width, + samples, + left + x_offset, + top + y_offset, + ) + + return Winding.apply(controls) diff --git a/src/vectrify/refine/paths.py b/src/vectrify/refine/paths.py index a4a5145d..40608271 100644 --- a/src/vectrify/refine/paths.py +++ b/src/vectrify/refine/paths.py @@ -16,6 +16,7 @@ from collections import defaultdict from collections.abc import Iterator, Mapping from contextlib import contextmanager +from functools import lru_cache from typing import Any import numpy as np @@ -292,6 +293,7 @@ def coverage( # sampled chords. _UNITS_PER_SAMPLE = 15.0 _MIN_SAMPLES, _MAX_SAMPLES = 8, 48 +_FUSED_CUBICS = 16 def _samples_for(control: Any) -> int: @@ -515,6 +517,182 @@ def _fill_coverage( return torch.stack(coverages).mean(dim=0) +def _fill_winding_chunk(start: Any, end: Any, pixels: Any) -> Any: + """Sum the winding angles of batched sampled contours at ``pixels``. + + Keeping this primitive separate gives ``torch.compile`` one regular, + side-effect-free GPU expression to fuse. In eager mode it is deliberately + the same arithmetic previously in :func:`_fill_coverages`. + """ + import torch + + offset_start = start[:, None] - pixels[None, :, None] + offset_end = end[:, None] - pixels[None, :, None] + cross = ( + offset_start[..., 0] * offset_end[..., 1] + - offset_start[..., 1] * offset_end[..., 0] + ) + dot = (offset_start * offset_end).sum(dim=-1) + return torch.atan2(cross, dot).sum(dim=-1) + + +def _dynamic_fill_winding_chunk(start: Any, end: Any, pixels: Any) -> Any: + """The tiled counterpart with only its pixel dimension left symbolic.""" + return _fill_winding_chunk(start, end, pixels) + + +@lru_cache(maxsize=1) +def _compiled_fill_winding_chunk() -> Any: + """Return the CUDA-fused winding primitive when this torch supports it. + + This is intentionally lazy: installing Vectrify must not require a CUDA + compiler, and the normal CPU renderer remains useful for tests and small + jobs. Inductor generates a kernel for this exact operation rather than + adding an external renderer dependency. + """ + import torch + + compile_fn = getattr(torch, "compile", None) + if compile_fn is None: + return _fill_winding_chunk + try: + return compile_fn( + _fill_winding_chunk, + fullgraph=True, + dynamic=False, + # This primitive is invoked repeatedly while its earlier outputs + # still participate in one optimisation graph. CUDA graph replay + # cannot safely reuse those outputs and adds substantial overhead. + options={"triton.cudagraphs": False}, + ) + except (RuntimeError, TypeError): + log.warning("CUDA winding fusion is unavailable; using eager torch.") + return _fill_winding_chunk + + +@lru_cache(maxsize=1) +def _compiled_tiled_fill_winding_chunk() -> Any: + """Fuse arbitrary-size clipped tiles after path shapes were normalized.""" + import torch + + compile_fn = getattr(torch, "compile", None) + if compile_fn is None: + return _fill_winding_chunk + try: + return compile_fn( + _dynamic_fill_winding_chunk, + fullgraph=True, + dynamic=True, + options={"triton.cudagraphs": False}, + ) + except (RuntimeError, TypeError): + log.warning("CUDA tiled winding fusion is unavailable; using eager torch.") + return _fill_winding_chunk + + +def _pad_fused_cubics(controls: Any) -> Any: + """Pad a short closed contour with zero-length cubics for CUDA fusion.""" + import torch + + count = controls.shape[1] + if count >= _FUSED_CUBICS: + return controls + # All four points are the contour origin, so every added sampled segment + # has zero angle at every pixel and cannot change its winding number. + point = controls[:, :1, :1].expand(-1, _FUSED_CUBICS - count, 4, -1) + return torch.cat((controls, point), dim=1) + + +def _fill_batched_windings( + controls: Any, + box: tuple[int, int, int, int], + *, + samples: int, + x_offset: float, + y_offset: float, + batch_size: int = 4, + pixel_chunk: int = 4_096, + winding_chunk: Any | None = None, +) -> Any: + """Return one winding field per equal-sized contour on CUDA. + + Unlike :func:`_fill_coverages`, this stops before applying a fill rule. + A multi-contour SVG path needs its contour windings summed before that + nonlinearity, so this is the reusable GPU building block for holes. + """ + import torch + + # The packaged native operator uses SAMVG's fixed 16-cubic contour + # representation. It remains deliberately narrow: every + # other contour/layout continues through the proven Torch implementation. + if samples in {8, 16, 32}: + from vectrify.refine.cuda_renderer import winding as cuda_winding + + native = cuda_winding( + controls, + box, + samples=samples, + x_offset=x_offset, + y_offset=y_offset, + ) + if native is not None: + return native + + left, top, right, bottom = box + height, width = bottom - top, right - left + steps = torch.linspace(0, 1, samples, device=controls.device, dtype=controls.dtype) + basis = torch.stack( + [ + (1 - steps) ** 3, + 3 * steps * (1 - steps) ** 2, + 3 * steps**2 * (1 - steps), + steps**3, + ], + dim=-1, + ) + ys, xs = torch.meshgrid( + torch.arange(height, device=controls.device, dtype=controls.dtype) + + top + + y_offset, + torch.arange(width, device=controls.device, dtype=controls.dtype) + + left + + x_offset, + indexing="ij", + ) + pixels = torch.stack((xs, ys), dim=-1).reshape(-1, 2) + if winding_chunk is None: + winding_chunk = ( + _compiled_fill_winding_chunk() + if controls.shape[1] <= _FUSED_CUBICS + else _fill_winding_chunk + ) + output = [] + for control in controls.split(batch_size): + count = len(control) + if count < batch_size: + # Keep the compiled kernel's leading dimension static. The + # padding is sliced away before it reaches the caller, so it has + # no effect on pixels or gradients of the real contours. + control = torch.cat( + (control, control[:1].expand(batch_size - count, -1, -1, -1)) + ) + if control.shape[1] <= _FUSED_CUBICS: + control = _pad_fused_cubics(control) + curve = torch.einsum("sk,nqkc->nqsc", basis, control).flatten(1, 2) + curve = torch.cat((curve, curve[:, :1]), dim=1) + winding = [] + for pixel_start in range(0, len(pixels), pixel_chunk): + winding.append( + winding_chunk( + curve[:, :-1], + curve[:, 1:], + pixels[pixel_start : pixel_start + pixel_chunk], + ) + ) + output.append(torch.cat(winding, dim=1)[:count]) + return torch.cat(output).reshape(-1, height, width) + + def _fill_path_coverage( contours: list[Any], box: tuple[int, int, int, int], @@ -523,21 +701,109 @@ def _fill_path_coverage( samples: int = 32, softness: float = 0.25, subpixels: int = 4, + fuse: bool = True, + dynamic_fuse: bool = False, ) -> Any: """Rasterise every contour according to SVG's fill-rule semantics.""" import torch + if contours and contours[0].is_cuda: + # SAM's detailed masks can have many contours. Sum each contour's + # winding before applying the SVG fill rule, exactly as the eager + # implementation below does, but keep the pixel/segment loops in the + # fused CUDA primitive. + fused = [contour for contour in contours if contour.shape[0] <= _FUSED_CUBICS] + unfused: dict[tuple[int, ...], list[Any]] = defaultdict(list) + for contour in contours: + if contour.shape[0] > _FUSED_CUBICS: + unfused[tuple(contour.shape)].append(contour) + # SAMVG's tracer normally emits at most 16 cubics per contour. Pad + # those contours once and render a whole path in one large GPU batch, + # rather than launching a tiny batch for each contour-shape group. + fused_controls = ( + torch.cat([_pad_fused_cubics(contour[None]) for contour in fused]) + if fused + else None + ) + if fuse: + winding_chunk = _compiled_fill_winding_chunk() + elif dynamic_fuse: + winding_chunk = _compiled_tiled_fill_winding_chunk() + else: + winding_chunk = _fill_winding_chunk + contour_batch_size = 64 if (fuse or dynamic_fuse) else 4 + coverages = [] + for y in range(subpixels): + for x in range(subpixels): + winding = torch.zeros( + (box[3] - box[1], box[2] - box[0]), + dtype=contours[0].dtype, + device=contours[0].device, + ) + if fused_controls is not None: + winding = winding + _fill_batched_windings( + fused_controls, + box, + samples=samples, + x_offset=(x + 0.5) / subpixels, + y_offset=(y + 0.5) / subpixels, + batch_size=contour_batch_size, + winding_chunk=winding_chunk, + ).sum(dim=0) + for group in unfused.values(): + winding = winding + _fill_batched_windings( + torch.stack(group), + box, + samples=samples, + x_offset=(x + 0.5) / subpixels, + y_offset=(y + 0.5) / subpixels, + winding_chunk=winding_chunk, + ).sum(dim=0) + if fill_rule == "evenodd": + coverages.append(0.5 * (1 - torch.cos(winding / 2))) + else: + coverages.append( + torch.sigmoid((winding.abs() - math.pi) / softness) + ) + return torch.stack(coverages).mean(dim=0) + + def contour_winding(contour: Any, x_offset: float, y_offset: float) -> Any: + # A noisy SAM mask can contain dozens of enclosed contours. Keeping + # every pixel-by-segment intermediate alive until its layer loss is + # backpropagated exhausts VRAM even on a small working canvas. + # Checkpointing recomputes the same differentiable winding field during + # backward, preserving renderer semantics and gradients exactly. + if contour.requires_grad: + from torch.utils.checkpoint import checkpoint + + return checkpoint( + lambda value: _fill_winding( + value, + box, + samples=samples, + x_offset=x_offset, + y_offset=y_offset, + ), + contour, + use_reentrant=False, + ) + return _fill_winding( + contour, + box, + samples=samples, + x_offset=x_offset, + y_offset=y_offset, + ) + coverages = [] for y in range(subpixels): for x in range(subpixels): winding = sum( ( - _fill_winding( + contour_winding( contour, - box, - samples=samples, - x_offset=(x + 0.5) / subpixels, - y_offset=(y + 0.5) / subpixels, + (x + 0.5) / subpixels, + (y + 0.5) / subpixels, ) for contour in contours ), @@ -548,9 +814,7 @@ def _fill_path_coverage( # expression is zero for an even count and one for an odd one. coverages.append(0.5 * (1 - torch.cos(winding / 2))) else: - coverages.append( - torch.sigmoid((winding.abs() - math.pi) / softness) - ) + coverages.append(torch.sigmoid((winding.abs() - math.pi) / softness)) return torch.stack(coverages).mean(dim=0) @@ -563,6 +827,8 @@ def _fill_coverages( fill_rule: str = "nonzero", pixel_chunk: int = 1_024, subpixels: int = 4, + fuse: bool = True, + dynamic_fuse: bool = False, ) -> Any: """Rasterise equal-sized closed cubic paths together on the GPU. @@ -585,7 +851,30 @@ def _fill_coverages( dim=-1, ) output = [] + can_fuse = controls.is_cuda and controls.shape[1] <= _FUSED_CUBICS + if fuse and can_fuse: + winding_chunk = _compiled_fill_winding_chunk() + elif dynamic_fuse and can_fuse: + winding_chunk = _compiled_tiled_fill_winding_chunk() + else: + winding_chunk = _fill_winding_chunk + # The fused CUDA kernel consumes far less temporary memory than eager + # broadcasting, so one complete small SAMVG working raster is faster than + # many 1,024-pixel launches. Retain the conservative caller-selected + # chunking on CPU. + active_pixel_chunk = ( + max(pixel_chunk, 4_096) + if controls.is_cuda and (fuse or dynamic_fuse) + else pixel_chunk + ) for control in controls.split(batch_size): + count = len(control) + if controls.is_cuda and count < batch_size: + control = torch.cat( + (control, control[:1].expand(batch_size - count, -1, -1, -1)) + ) + if controls.is_cuda and control.shape[1] <= _FUSED_CUBICS: + control = _pad_fused_cubics(control) curve = torch.einsum("sk,nqkc->nqsc", basis, control).flatten(1, 2) curve = torch.cat((curve, curve[:, :1]), dim=1) start = curve[:, :-1] @@ -604,16 +893,9 @@ def _fill_coverages( ) pixels = torch.stack((xs, ys), dim=-1).reshape(-1, 2) coverages = [] - for pixel_start in range(0, len(pixels), pixel_chunk): - pixel_block = pixels[pixel_start : pixel_start + pixel_chunk] - offset_start = start[:, None] - pixel_block[None, :, None] - offset_end = end[:, None] - pixel_block[None, :, None] - cross = ( - offset_start[..., 0] * offset_end[..., 1] - - offset_start[..., 1] * offset_end[..., 0] - ) - dot = (offset_start * offset_end).sum(dim=-1) - winding = torch.atan2(cross, dot).sum(dim=-1) + for pixel_start in range(0, len(pixels), active_pixel_chunk): + pixel_block = pixels[pixel_start : pixel_start + active_pixel_chunk] + winding = winding_chunk(start, end, pixel_block) if fill_rule == "evenodd": coverages.append(0.5 * (1 - torch.cos(winding / 2))) else: @@ -622,27 +904,154 @@ def _fill_coverages( ) coverage = torch.cat(coverages, dim=1) coverage_sum = ( - coverage - if coverage_sum is None - else coverage_sum + coverage + coverage if coverage_sum is None else coverage_sum + coverage ) assert coverage_sum is not None - output.append(coverage_sum / (subpixels * subpixels)) + output.append((coverage_sum / (subpixels * subpixels))[:count]) return torch.cat(output).reshape(-1, height, width) -def _xing_loss(control: Any) -> Any: - """Return SAMVG's normalized per-cubic Xing regularizer (Eq. 3-6--3-8).""" +def _fill_distance_surrogate_coverages( + controls: Any, + box: tuple[int, int, int, int], + *, + samples: int = 32, + softness: float = 0.25, + subpixels: int = 2, + pixel_chunk: int = 1_024, +) -> Any: + """Exact sampled-fill forward values with closest-cubic surrogate gradients. + + This is the same forward/backward split used by differentiable vector + renderers: topology is a discrete winding decision, while a local signed + distance provides useful boundary gradients. The hard forward samples + retain the existing renderer's fill semantics; only the expensive + pixel-by-512-segment autograd graph is replaced for simple contours. + """ + import torch + + left, top, right, bottom = box + height, width = bottom - top, right - left + result = [] + seeds = torch.linspace(0, 1, 5, dtype=controls.dtype, device=controls.device) + for y in range(subpixels): + for x in range(subpixels): + x_offset, y_offset = (x + 0.5) / subpixels, (y + 0.5) / subpixels + with torch.no_grad(): + winding = _fill_batched_windings( + controls.detach(), + box, + samples=samples, + x_offset=x_offset, + y_offset=y_offset, + batch_size=4, + pixel_chunk=pixel_chunk, + winding_chunk=_compiled_tiled_fill_winding_chunk(), + ) + inside = winding.abs() >= math.pi + ys, xs = torch.meshgrid( + torch.arange(height, device=controls.device, dtype=controls.dtype) + + top + + y_offset, + torch.arange(width, device=controls.device, dtype=controls.dtype) + + left + + x_offset, + indexing="ij", + ) + pixels = torch.stack((xs, ys), dim=-1).reshape(-1, 2) + surrogate_parts = [] + for start in range(0, len(pixels), pixel_chunk): + pixel = pixels[start : start + pixel_chunk] + parameter = seeds[None, None, None].expand( + len(controls), len(pixel), controls.shape[1], -1 + ) + cubic = controls[:, None, :, None] + for _ in range(5): + inverse = 1 - parameter + point = ( + inverse[..., None] ** 3 * cubic[..., 0, :] + + 3 + * inverse[..., None] ** 2 + * parameter[..., None] + * cubic[..., 1, :] + + 3 + * inverse[..., None] + * parameter[..., None] ** 2 + * cubic[..., 2, :] + + parameter[..., None] ** 3 * cubic[..., 3, :] + ) + tangent = ( + 3 + * inverse[..., None] ** 2 + * (cubic[..., 1, :] - cubic[..., 0, :]) + + 6 + * inverse[..., None] + * parameter[..., None] + * (cubic[..., 2, :] - cubic[..., 1, :]) + + 3 + * parameter[..., None] ** 2 + * (cubic[..., 3, :] - cubic[..., 2, :]) + ) + acceleration = 6 * inverse[..., None] * ( + cubic[..., 2, :] - 2 * cubic[..., 1, :] + cubic[..., 0, :] + ) + 6 * parameter[..., None] * ( + cubic[..., 3, :] - 2 * cubic[..., 2, :] + cubic[..., 1, :] + ) + delta = point - pixel[None, :, None, None, :] + parameter = ( + parameter + - (delta * tangent).sum(dim=-1) + / ( + (tangent * tangent).sum(dim=-1) + + (delta * acceleration).sum(dim=-1) + ).clamp_min(1e-8) + ).clamp(0, 1) + inverse = 1 - parameter + point = ( + inverse[..., None] ** 3 * cubic[..., 0, :] + + 3 + * inverse[..., None] ** 2 + * parameter[..., None] + * cubic[..., 1, :] + + 3 + * inverse[..., None] + * parameter[..., None] ** 2 + * cubic[..., 2, :] + + parameter[..., None] ** 3 * cubic[..., 3, :] + ) + distance = ( + (point - pixel[None, :, None, None, :]) + .norm(dim=-1) + .amin(dim=(-1, -2)) + ) + sign = torch.where( + inside.reshape(len(controls), -1)[:, start : start + len(pixel)], + -1.0, + 1.0, + ) + surrogate_parts.append(torch.sigmoid(-sign * distance / softness)) + surrogate = torch.cat(surrogate_parts, dim=1).reshape(-1, height, width) + hard = inside.to(dtype=controls.dtype) + result.append(hard + surrogate - surrogate.detach()) + return torch.stack(result).mean(dim=0) + + +def _xing_penalties(control: Any) -> Any: + """Return SAMVG's normalized Xing penalty for every cubic in ``control``.""" import torch start_handle = control[:, 1] - control[:, 0] end_handle = control[:, 3] - control[:, 2] cross = ( - start_handle[:, 0] * end_handle[:, 1] - - start_handle[:, 1] * end_handle[:, 0] + start_handle[:, 0] * end_handle[:, 1] - start_handle[:, 1] * end_handle[:, 0] ) sine = cross / (start_handle.norm(dim=-1) * end_handle.norm(dim=-1) + 1e-12) - return torch.where(cross < 0, torch.relu(-sine), torch.relu(sine)).mean() + return torch.where(cross < 0, torch.relu(-sine), torch.relu(sine)) + + +def _xing_loss(control: Any) -> Any: + """Return SAMVG's normalized per-cubic Xing regularizer (Eq. 3-6--3-8).""" + return _xing_penalties(control).mean() _HEX_FILL = re.compile(r"^#([0-9a-fA-F]{6})$") @@ -656,6 +1065,44 @@ def _fill_rgb(value: str | None) -> tuple[float, float, float] | None: return tuple(int(digits[index : index + 2], 16) / 255 for index in range(0, 6, 2)) +def _composite_opaque_fills(alphas: Any, colours: Any) -> Any: + """Composite opaque SVG fills in document order without a layer loop. + + Each layer contributes its premultiplied colour through the product of the + transparencies above it. This is algebraically identical to repeatedly + applying ``canvas * (1 - alpha) + colour * alpha`` over a black canvas, + but lets Torch execute the 223-layer cat seed in a few large operations. + """ + import torch + + transparency = 1 - alphas + above_inclusive = torch.cumprod(transparency.flip(0), dim=0).flip(0) + above = torch.cat((above_inclusive[1:], torch.ones_like(alphas[:1])), dim=0) + return ( + colours.clamp(0, 1)[:, None, None, :] * alphas[..., None] * above[..., None] + ).sum(dim=0) + + +@lru_cache(maxsize=1) +def _compiled_opaque_fill_composite() -> Any: + """Return the CUDA-fused painter's-order compositing kernel when possible.""" + import torch + + compile_fn = getattr(torch, "compile", None) + if compile_fn is None: + return _composite_opaque_fills + try: + return compile_fn( + _composite_opaque_fills, + fullgraph=True, + dynamic=False, + options={"triton.cudagraphs": False}, + ) + except (RuntimeError, TypeError): + log.warning("CUDA fill compositing fusion is unavailable; using eager torch.") + return _composite_opaque_fills + + def fit_filled_svg( svg: str, target: Image.Image, @@ -665,6 +1112,9 @@ def fit_filled_svg( color_learning_rate: float = 0.01, xing_weight: float = 0.02, optimisation_long_side: int | None = None, + subpixels: int = 2, + monolithic: bool = False, + curve_samples: int | None = None, ) -> str: """Optimise opaque filled cubic SVG paths against an RGB target. @@ -672,8 +1122,12 @@ def fit_filled_svg( iterations in each of its two passes. This implementation reuses Vectrify's torch renderer instead of requiring DiffVG, while retaining the dissertation's full-resolution Adam defaults: point LR 1, colour LR .01, - and MSE plus .02 Xing loss. ``optimisation_long_side`` is available only - as an explicit caller-selected preview mode. + and MSE plus .02 Xing loss. It uses DiffVG's standard 2x2 optimisation + sampling; the standalone renderer retains its stricter 4x4 default for + Cairo-fidelity checks. Small clipped tiles use fewer cubic samples because + their screen-space deviation is bounded by the tile size; pass + ``curve_samples`` to override that adaptive choice. ``optimisation_long_side`` + is available only as an explicit caller-selected preview mode. """ import xml.etree.ElementTree as ET @@ -740,12 +1194,128 @@ def fit_filled_svg( [control for path in controls for control in path], lr=point_learning_rate ) colour_optimizer = torch.optim.Adam(colours, lr=color_learning_rate) - box = (0, 0, work_width, work_height) - simple_groups: dict[tuple[tuple[int, ...], str], list[int]] = defaultdict(list) - for index, path in enumerate(controls): - if len(path) == 1: - fill_rule = entries[index][3] - simple_groups[(tuple(path[0].shape), fill_rule)].append(index) + # The dissertation averages Xing within each contour then sums contours. + # Keep that weighting while evaluating the 413 cat contours in one CUDA + # expression rather than launching one tiny graph for each. + xing_contour_weights = torch.cat( + [ + torch.full( + (len(control),), + 1 / len(control), + dtype=goal.dtype, + device=device, + ) + for path in controls + for control in path + ] + ) + + def tile_for(path: list[Any]) -> tuple[int, int, int, int]: + """A fixed, antialiased raster tile covering a path's control hull.""" + points = torch.cat([control.detach().reshape(-1, 2) for control in path]) + # Cubic Beziers lie in their control hull. Two pixels retain the + # entire soft edge while avoiding the full-canvas work DiffVG culls. + left = max(0, math.floor(float(points[:, 0].min())) - 2) + top = max(0, math.floor(float(points[:, 1].min())) - 2) + right = min(work_width, math.ceil(float(points[:, 0].max())) + 2) + bottom = min(work_height, math.ceil(float(points[:, 1].max())) + 2) + + # Bucket dimensions keep many unrelated small paths in the same CUDA + # batch. Shift a bucket at the canvas edge rather than clipping its + # protected coverage margin. + tile_width = min(work_width, 8 * math.ceil(max(right - left, 1) / 8)) + tile_height = min(work_height, 8 * math.ceil(max(bottom - top, 1) / 8)) + left = min(left, work_width - tile_width) + top = min(top, work_height - tile_height) + return left, top, tile_width, tile_height + + def samples_for(tile_width: int, tile_height: int) -> int: + """Choose winding tessellation from the cubic's visible pixel extent.""" + if curve_samples is not None: + return curve_samples + longest_side = max(tile_width, tile_height) + if longest_side <= 32: + return 8 + if longest_side <= 64: + return 16 + return 32 + + def restore_tile(alpha: Any, left: int, top: int) -> Any: + return torch.nn.functional.pad( + alpha, + ( + left, + work_width - left - alpha.shape[1], + top, + work_height - top - alpha.shape[0], + ), + ) + + def cropped_simple_groups() -> dict[ + tuple[tuple[int, ...], str, int, int], list[tuple[int, int, int]] + ]: + groups: dict[ + tuple[tuple[int, ...], str, int, int], list[tuple[int, int, int]] + ] = defaultdict(list) + for index, path in enumerate(controls): + if len(path) != 1: + continue + left, top, tile_width, tile_height = tile_for(path) + groups[ + (tuple(path[0].shape), entries[index][3], tile_width, tile_height) + ].append((index, left, top)) + return groups + + def rasterise_simple( + fill_rule: str, + tile_width: int, + tile_height: int, + items: list[tuple[int, int, int]], + ) -> list[tuple[int, Any]]: + translated = torch.stack( + [ + controls[index][0] - controls[index][0].new_tensor((left, top)) + for index, left, top in items + ] + ) + rasterised = _fill_coverages( + translated, + (0, 0, tile_width, tile_height), + fill_rule=fill_rule, + samples=samples_for(tile_width, tile_height), + subpixels=subpixels, + fuse=False, + dynamic_fuse=len(items) >= 4, + ) + return [ + (index, restore_tile(alpha, left, top)) + for (index, left, top), alpha in zip(items, rasterised, strict=True) + ] + + def rasterise_multi(index: int, path: list[Any]) -> Any: + # Very hole-heavy masks already amortise the fused full-frame kernel. + # Tile smaller multi-contour paths, where culling most of their empty + # canvas wins over compiling another specialised large batch. + if len(path) >= 16: + return _fill_path_coverage( + path, + (0, 0, work_width, work_height), + fill_rule=entries[index][3], + samples=samples_for(work_width, work_height), + subpixels=subpixels, + ) + left, top, tile_width, tile_height = tile_for(path) + offset = path[0].new_tensor((left, top)) + alpha = _fill_path_coverage( + [control - offset for control in path], + (0, 0, tile_width, tile_height), + fill_rule=entries[index][3], + samples=samples_for(tile_width, tile_height), + subpixels=subpixels, + fuse=False, + ) + return restore_tile(alpha, left, top) + log.info( "Filled-path optimisation: %d path(s), %dx%d working raster on %s.", len(entries), @@ -756,25 +1326,60 @@ def fit_filled_svg( for _step in range(steps): point_optimizer.zero_grad() colour_optimizer.zero_grad() + simple_groups = cropped_simple_groups() + all_controls = torch.cat([control for path in controls for control in path]) + + if monolithic: + alphas: list[Any | None] = [None] * len(entries) + for ( + _shape, + fill_rule, + tile_width, + tile_height, + ), items in simple_groups.items(): + for index, alpha in rasterise_simple( + fill_rule, tile_width, tile_height, items + ): + alphas[index] = alpha + for index, path in enumerate(controls): + if alphas[index] is None: + alphas[index] = rasterise_multi(index, path) + alpha_stack = torch.stack([alpha for alpha in alphas if alpha is not None]) + composite = ( + _compiled_opaque_fill_composite() + if goal.is_cuda + else _composite_opaque_fills + ) + rendered = composite(alpha_stack, torch.stack(colours)) + loss = ((rendered - goal) ** 2).mean() + loss = ( + loss + + xing_weight + * (_xing_penalties(all_controls) * xing_contour_weights).sum() + ) + loss.backward() + point_optimizer.step() + colour_optimizer.step() + continue # First composite the exact same soft fills without recording an # autograd graph. The saved canvases and suffix transparencies are # enough to derive the MSE gradient of each layer independently. with torch.no_grad(): initial_alphas: list[Any | None] = [None] * len(entries) - for (_shape, fill_rule), indices in simple_groups.items(): - rasterised = _fill_coverages( - torch.stack([controls[index][0] for index in indices]), - box, - fill_rule=fill_rule, - ) - for index, alpha in zip(indices, rasterised, strict=True): + for ( + _shape, + fill_rule, + tile_width, + tile_height, + ), items in simple_groups.items(): + for index, alpha in rasterise_simple( + fill_rule, tile_width, tile_height, items + ): initial_alphas[index] = alpha for index, path in enumerate(controls): if initial_alphas[index] is None: - initial_alphas[index] = _fill_path_coverage( - path, box, fill_rule=entries[index][3] - ) + initial_alphas[index] = rasterise_multi(index, path) before: list[Any] = [] rendered = torch.zeros_like(goal) @@ -810,41 +1415,44 @@ def layer_loss( assert stored_alpha is not None assert suffix is not None colour = colours[index] - path = controls[index] colour_delta = colour.detach().clamp(0, 1) - canvases[index] - alpha_gradient = ( - gradient * suffix[..., None] * colour_delta - ).sum(dim=-1) + alpha_gradient = (gradient * suffix[..., None] * colour_delta).sum(dim=-1) colour_gradient = ( gradient * suffix[..., None] * stored_alpha[..., None] ).sum(dim=(0, 1)) - loss = (alpha * alpha_gradient.detach()).sum() - loss = loss + (colour.clamp(0, 1) * colour_gradient.detach()).sum() - for control in path: - loss = loss + xing_weight * _xing_loss(control) - return loss + return (alpha * alpha_gradient.detach()).sum() + ( + colour.clamp(0, 1) * colour_gradient.detach() + ).sum() # Backpropagate a bounded batch at a time. The compositing derivative # above accounts for all later opaque layers, so this has the same MSE # gradient as one monolithic render without its peak-memory cost. - for (_shape, fill_rule), indices in simple_groups.items(): - for offset in range(0, len(indices), 4): - batch = indices[offset : offset + 4] - rasterised = _fill_coverages( - torch.stack([controls[index][0] for index in batch]), - box, - fill_rule=fill_rule, - ) + for ( + _shape, + fill_rule, + tile_width, + tile_height, + ), items in simple_groups.items(): + for offset in range(0, len(items), 4): + batch = items[offset : offset + 4] loss = torch.zeros((), device=device) - for index, alpha in zip(batch, rasterised, strict=True): + for index, alpha in rasterise_simple( + fill_rule, tile_width, tile_height, batch + ): loss = loss + layer_loss(index, alpha) loss.backward() - simple_indices = {index for group in simple_groups.values() for index in group} + simple_indices = { + index for group in simple_groups.values() for index, _left, _top in group + } for index, path in enumerate(controls): if index not in simple_indices: layer_loss( - index, _fill_path_coverage(path, box, fill_rule=entries[index][3]) + index, + rasterise_multi(index, path), ).backward() + ( + xing_weight * (_xing_penalties(all_controls) * xing_contour_weights).sum() + ).backward() point_optimizer.step() colour_optimizer.step() @@ -853,8 +1461,7 @@ def layer_loss( entries, controls, colours, strict=True ): data = " ".join( - to_path_d((control.detach().cpu() / coordinate_scale_cpu).tolist()) - + " Z" + to_path_d((control.detach().cpu() / coordinate_scale_cpu).tolist()) + " Z" for control in path ) element.set("d", data) diff --git a/src/vectrify/refine/samvg.py b/src/vectrify/refine/samvg.py index d8019aa2..79790ab2 100644 --- a/src/vectrify/refine/samvg.py +++ b/src/vectrify/refine/samvg.py @@ -25,6 +25,11 @@ # ViT-H is the paper-quality default; users who need the smaller checkpoint can # opt down without changing the package through VECTRIFY_SAMVG_MODEL. SAMVG_MODEL = os.environ.get("VECTRIFY_SAMVG_MODEL", "facebook/sam-vit-huge") +# SAMVG's own impact filter selects useful masks against the image. Retaining +# AMG's score gates here discarded the small facial candidates needed by the +# photo seed before that image-aware test could evaluate them. +SAMVG_PRED_IOU_THRESH = 0.0 +SAMVG_STABILITY_SCORE_THRESH = 0.0 @dataclass(frozen=True) @@ -80,6 +85,8 @@ def masks_for(source: Image.Image) -> list[np.ndarray]: points_per_batch=32, points_per_crop=32, crops_n_layers=0, + pred_iou_thresh=SAMVG_PRED_IOU_THRESH, + stability_score_thresh=SAMVG_STABILITY_SCORE_THRESH, )["masks"] ] @@ -110,17 +117,15 @@ def masks_for(source: Image.Image) -> list[np.ndarray]: def _components( - mask: np.ndarray, min_pixels: int, *, fill_holes: bool = False + mask: np.ndarray, min_pixels: int, *, fill_holes: bool = True ) -> list[np.ndarray]: - """Return traceable mask components without inventing filled regions. + """Return traceable AMG components after its required hole cleanup. - AMG already performs its configured small-region postprocessing. Filling - every remaining hole changes an eye, ear, or gap between hairs into a - solid region, and keeping disconnected pieces in one SVG path turns them - into one optimisation unit. SAMVG traces each component, so impact - filtering must receive those components separately. + SAMVG traces each connected component independently. Filling its mask + holes before tracing matches AMG's small-region cleanup and prevents a + noisy mask from becoming hundreds of even-odd SVG contours. """ - from scipy.ndimage import binary_fill_holes, label + from scipy.ndimage import label labels, count = label(mask) components = [] @@ -129,7 +134,21 @@ def _components( if int(component.sum()) < min_pixels: continue if fill_holes: - component = binary_fill_holes(component) + # AMG's postprocessing removes *small* enclosed holes, rather + # than turning meaningful cutouts such as an eye into a solid + # region. The same area cutoff as tiny components keeps those + # two decisions consistent. + background, hole_count = label(~component) + for hole in range(1, hole_count + 1): + points = background == hole + if ( + int(points.sum()) <= min_pixels + and not points[0].any() + and not points[-1].any() + and not points[:, 0].any() + and not points[:, -1].any() + ): + component[points] = True components.append(np.asarray(component, dtype=bool)) return components @@ -193,7 +212,7 @@ def filter_by_impact( min_pixels: int = 32, min_impact: float = 1e-5, max_layers: int = 128, - fill_holes: bool = False, + fill_holes: bool = True, ) -> list[MaskLayer]: """Keep masks that lower blank-canvas reconstruction error. @@ -320,6 +339,7 @@ def retrieve_layers( min_pixels: int = 32, min_impact: float = 1e-5, max_layers: int = 512, + fill_holes: bool = True, ) -> list[MaskLayer]: """Run SAMVG's automatic-mask, coverage-prompt, filter sequence.""" image = image.convert("RGB") @@ -330,6 +350,7 @@ def retrieve_layers( min_pixels=min_pixels, min_impact=min_impact, max_layers=max_layers, + fill_holes=fill_holes, ) layers = recolour_visible_layers(image, layers) points = coverage_prompt_points(layers, (image.height, image.width)) @@ -341,6 +362,7 @@ def retrieve_layers( min_pixels=min_pixels, min_impact=min_impact, max_layers=max_layers, + fill_holes=fill_holes, ) recovered = recolour_visible_layers(image, recovered) log.info( @@ -434,10 +456,7 @@ def solve(parameters: np.ndarray) -> np.ndarray: 3 * (1 - parameters) * parameters**2, ) ) - base = ( - (1 - parameters)[:, None] ** 3 * start - + parameters[:, None] ** 3 * end - ) + base = (1 - parameters)[:, None] ** 3 * start + parameters[:, None] ** 3 * end controls, *_ = np.linalg.lstsq(matrix, points - base, rcond=None) return controls @@ -460,9 +479,8 @@ def solve(parameters: np.ndarray) -> np.ndarray: + 6 * omt[:, None] * t[:, None] * (p1 - p0) + 3 * t[:, None] ** 2 * (end - p1) ) - second = ( - 6 * omt[:, None] * (p1 - 2 * p0 + start) - + 6 * t[:, None] * (end - 2 * p1 + p0) + second = 6 * omt[:, None] * (p1 - 2 * p0 + start) + 6 * t[:, None] * ( + end - 2 * p1 + p0 ) offset = curve - points numerator = (offset * first).sum(axis=1) @@ -532,7 +550,8 @@ def generate_svg( min_pixels: int = 32, min_impact: float = 1e-5, max_layers: int = 512, - segments: int = 8, + segments: int = 16, + fill_holes: bool = True, ) -> str: """Generate SAMVG's traced, pre-optimisation SVG from a target image.""" image = image.convert("RGB") @@ -543,6 +562,7 @@ def generate_svg( min_pixels=min_pixels, min_impact=min_impact, max_layers=max_layers, + fill_holes=fill_holes, ) if masks is not None else retrieve_layers( @@ -550,6 +570,7 @@ def generate_svg( min_pixels=min_pixels, min_impact=min_impact, max_layers=max_layers, + fill_holes=fill_holes, ) ) paths = [] @@ -626,9 +647,9 @@ def _append_layers(svg: str, layers: list[MaskLayer], segments: int) -> str: def _render_svg(svg: str, image: Image.Image, rasterize) -> Image.Image: - return Image.open( - io.BytesIO(rasterize(svg, image.width, image.height)) - ).convert("RGB") + return Image.open(io.BytesIO(rasterize(svg, image.width, image.height))).convert( + "RGB" + ) def _mse(image: Image.Image, rendered: Image.Image) -> float: diff --git a/tests/refine/test_filled_paths.py b/tests/refine/test_filled_paths.py index 70faf13e..9acc496f 100644 --- a/tests/refine/test_filled_paths.py +++ b/tests/refine/test_filled_paths.py @@ -6,6 +6,8 @@ from vectrify.image_utils import rasterize_svg_to_png_bytes from vectrify.refine.paths import ( + _composite_opaque_fills, + _fill_batched_windings, _fill_coverage, _fill_coverages, _fill_path_coverage, @@ -14,6 +16,103 @@ parse_filled_cubics, ) + +def _sixteen_cubic_circle(torch): + angle = torch.arange(17, device="cuda", dtype=torch.float32) * (2 * torch.pi / 16) + points = torch.stack((12 + 7 * torch.cos(angle), 13 + 6 * torch.sin(angle)), -1) + return torch.stack( + ( + points[:-1], + points[:-1] * (2 / 3) + points[1:] * (1 / 3), + points[:-1] * (1 / 3) + points[1:] * (2 / 3), + points[1:], + ), + 1, + )[None] + + +@pytest.mark.parametrize("samples", [8, 16, 32]) +def test_native_winding_matches_torch_forward_and_gradient(samples, monkeypatch): + """Release-wheel CUDA path agrees with the portable sampled renderer.""" + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + from vectrify.refine import cuda_renderer + + if not cuda_renderer.available(): + pytest.skip("optional SAMVG CUDA extension is not installed") + controls = _sixteen_cubic_circle(torch).requires_grad_() + native = _fill_batched_windings( + controls, (0, 0, 24, 24), samples=samples, x_offset=0.25, y_offset=0.25 + ) + extension = cuda_renderer._extension + monkeypatch.setattr(cuda_renderer, "_extension", lambda: None) + portable = _fill_batched_windings( + controls, (0, 0, 24, 24), samples=samples, x_offset=0.25, y_offset=0.25 + ) + upstream = torch.randn_like(native) + (native * upstream).sum().backward() + native_gradient = controls.grad.detach().clone() + controls.grad = None + (portable * upstream).sum().backward() + + assert torch.allclose(native, portable, atol=1e-5, rtol=1e-5) + assert torch.allclose(native_gradient, controls.grad, atol=1e-4, rtol=1e-4) + monkeypatch.setattr(cuda_renderer, "_extension", extension) + + +def test_native_winding_falls_back_without_the_optional_extension(monkeypatch): + torch = pytest.importorskip("torch") + from vectrify.refine import cuda_renderer + + monkeypatch.setattr(cuda_renderer, "_extension", lambda: None) + controls = torch.zeros((1, 16, 4, 2)) + assert ( + cuda_renderer.winding( + controls, (0, 0, 8, 8), samples=16, x_offset=0.5, y_offset=0.5 + ) + is None + ) + + +def test_native_even_odd_coverage_stays_cairo_validated(): + """The native winding path preserves SVG hole coverage, not just tensors.""" + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + from vectrify.refine.cuda_renderer import available + + if not available(): + pytest.skip("optional SAMVG CUDA extension is not installed") + size = 96 + contours = [ + torch.tensor(contour, dtype=torch.float32, device="cuda") + for contour in parse_filled_cubics(DONUT_PATH) + ] + native = _fill_path_coverage( + contours, (0, 0, size, size), fill_rule="evenodd" + ).cpu().numpy() + head = f'' + blank = f"{head}" + drawn = f'{head}' + real = ( + np.asarray( + Image.open( + io.BytesIO(rasterize_svg_to_png_bytes(blank, out_w=size, out_h=size)) + ).convert("L"), + dtype=np.float32, + ) + - np.asarray( + Image.open( + io.BytesIO(rasterize_svg_to_png_bytes(drawn, out_w=size, out_h=size)) + ).convert("L"), + dtype=np.float32, + ) + ) / 255.0 + + assert np.abs(native - real).mean() < 0.002 + assert native[48, 48] < 0.01 + SVG = ( '' '' blank = f"{head}" drawn = f'{head}' + def luminance(svg: str) -> np.ndarray: png = rasterize_svg_to_png_bytes(svg, out_w=size, out_h=size) - return np.asarray( - Image.open(io.BytesIO(png)).convert("L"), dtype=np.float32 - ) + return np.asarray(Image.open(io.BytesIO(png)).convert("L"), dtype=np.float32) real = (luminance(blank) - luminance(drawn)) / 255.0 @@ -99,7 +197,7 @@ def test_filled_fit_preserves_even_odd_contours_at_zero_steps(): fitted = fit_filled_svg(svg, target, steps=0) assert fitted.count("M ") == 2 - assert "fill-rule=\"evenodd\"" in fitted + assert 'fill-rule="evenodd"' in fitted def test_filled_fit_can_optimise_an_even_odd_path_without_losing_its_hole(): @@ -124,10 +222,10 @@ def test_filled_fit_can_optimise_an_even_odd_path_without_losing_its_hole(): rendered = Image.open( io.BytesIO(rasterize_svg_to_png_bytes(fitted, out_w=96, out_h=96)) ).convert("RGB") + def error(candidate: Image.Image) -> float: - difference = ( - np.asarray(candidate, dtype=np.float32) - - np.asarray(target, dtype=np.float32) + difference = np.asarray(candidate, dtype=np.float32) - np.asarray( + target, dtype=np.float32 ) return float((difference**2).mean()) @@ -215,6 +313,25 @@ def test_bounded_compositing_gradient_matches_monolithic_render(): assert torch.allclose(actual, expected, atol=1e-6, rtol=1e-5) +def test_tensorised_opaque_compositing_matches_layer_loop(): + torch = pytest.importorskip("torch") + alphas = torch.tensor( + [ + [[0.2, 0.7], [0.3, 0.5]], + [[0.6, 0.1], [0.8, 0.4]], + [[0.9, 0.2], [0.4, 0.3]], + ] + ) + colours = torch.tensor([[0.1, 0.3, 0.9], [0.9, 0.2, 0.4], [0.2, 0.8, 0.5]]) + expected = torch.zeros(2, 2, 3) + for alpha, colour in zip(alphas, colours, strict=True): + expected = expected * (1 - alpha[..., None]) + colour * alpha[..., None] + + actual = _composite_opaque_fills(alphas, colours) + + assert torch.allclose(actual, expected) + + def test_xing_loss_is_normalized_and_penalizes_either_turn_direction(): torch = pytest.importorskip("torch") controls = torch.tensor( @@ -244,12 +361,8 @@ def test_even_odd_uses_winding_parity_for_a_double_wound_contour(): for contour in parse_filled_cubics(double_loop) ] - even_odd = _fill_path_coverage( - contours, (0, 0, 96, 96), fill_rule="evenodd" - ) - nonzero = _fill_path_coverage( - contours, (0, 0, 96, 96), fill_rule="nonzero" - ) + even_odd = _fill_path_coverage(contours, (0, 0, 96, 96), fill_rule="evenodd") + nonzero = _fill_path_coverage(contours, (0, 0, 96, 96), fill_rule="nonzero") batched_even_odd = _fill_coverages( torch.stack(contours), (0, 0, 96, 96), fill_rule="evenodd" ) diff --git a/tests/refine/test_samvg.py b/tests/refine/test_samvg.py index 889fa76f..2c32ee9c 100644 --- a/tests/refine/test_samvg.py +++ b/tests/refine/test_samvg.py @@ -101,6 +101,16 @@ def test_components_are_separate_and_do_not_fill_meaningful_holes(): assert not components[0][3, 3] +def test_components_fill_only_tiny_enclosed_holes(): + mask = np.ones((8, 8), dtype=bool) + mask[3:5, 3:5] = False + + components = _components(mask, min_pixels=4) + + assert len(components) == 1 + assert components[0].all() + + def test_crop_edge_masks_are_rejected_unless_they_reach_the_image_edge(): cropped = np.ones((20, 30), dtype=bool) at_image_edge = np.zeros((20, 30), dtype=bool) diff --git a/uv.lock b/uv.lock index 6d9f3f6e..6e520606 100644 --- a/uv.lock +++ b/uv.lock @@ -3203,6 +3203,16 @@ dev = [ graphviz = [ { name = "graphviz" }, ] +samvg = [ + { name = "scikit-learn", version = "1.7.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "scikit-learn", version = "1.9.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, + { name = "scipy", version = "1.18.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "torch" }, + { name = "torchvision" }, + { name = "transformers" }, +] typst = [ { name = "typst" }, ] @@ -3233,16 +3243,21 @@ requires-dist = [ { name = "pytest-xdist", marker = "extra == 'dev'" }, { name = "rich", specifier = ">=14.3.3" }, { name = "ruff", marker = "extra == 'dev'" }, + { name = "scikit-learn", marker = "extra == 'samvg'", specifier = ">=1.3.0" }, { name = "scikit-learn", marker = "extra == 'vision'", specifier = ">=1.3.0" }, + { name = "scipy", marker = "extra == 'samvg'", specifier = ">=1.11.0" }, { name = "scipy", marker = "extra == 'vision'", specifier = ">=1.11.0" }, + { name = "torch", marker = "extra == 'samvg'", specifier = ">=2.0.0", index = "https://download.pytorch.org/whl/cu126" }, { name = "torch", marker = "extra == 'vision'", specifier = ">=2.0.0", index = "https://download.pytorch.org/whl/cu126" }, + { name = "torchvision", marker = "extra == 'samvg'", specifier = ">=0.28.0", index = "https://download.pytorch.org/whl/cu126" }, { name = "torchvision", marker = "extra == 'vision'", specifier = ">=0.28.0", index = "https://download.pytorch.org/whl/cu126" }, { name = "tqdm", specifier = ">=4.67.3" }, + { name = "transformers", marker = "extra == 'samvg'", specifier = ">=4.40.0" }, { name = "transformers", marker = "extra == 'vision'", specifier = ">=4.40.0" }, { name = "typst", marker = "extra == 'typst'", specifier = ">=0.11.0" }, - { name = "vectrify", extras = ["graphviz", "typst", "vision"], marker = "extra == 'all'" }, + { name = "vectrify", extras = ["graphviz", "samvg", "typst", "vision"], marker = "extra == 'all'" }, ] -provides-extras = ["vision", "graphviz", "typst", "all", "dev"] +provides-extras = ["vision", "samvg", "graphviz", "typst", "all", "dev"] [[package]] name = "webencodings" From e13cb76830709031dd2d0e885eaeb2d23eb9839b Mon Sep 17 00:00:00 2001 From: Rasmus Ros Date: Mon, 24 Aug 2026 10:29:58 +0200 Subject: [PATCH 2/6] feat: optimize analytic SAMVG tiles --- src/vectrify/refine/_samvg_cuda.cu | 641 +++++++++++++++++++++++++-- src/vectrify/refine/cuda_renderer.py | 262 +++++++++++ src/vectrify/refine/paths.py | 502 ++++++++++++++++++++- tests/refine/test_filled_paths.py | 341 ++++++++++++++ 4 files changed, 1691 insertions(+), 55 deletions(-) diff --git a/src/vectrify/refine/_samvg_cuda.cu b/src/vectrify/refine/_samvg_cuda.cu index 61e6e3e7..f0739b3a 100644 --- a/src/vectrify/refine/_samvg_cuda.cu +++ b/src/vectrify/refine/_samvg_cuda.cu @@ -1,6 +1,7 @@ #include #include #include +#include namespace { constexpr int kCubics = 16; @@ -20,30 +21,453 @@ __device__ inline void sample_path(const float* path, float* points, int samples __syncthreads(); } +__device__ inline void path_bounds(const float* path, float* bounds) { + if (threadIdx.x == 0) { + float min_x = path[0], min_y = path[1], max_x = path[0], max_y = path[1]; + for (int index = 1; index < kCubics * 4; ++index) { + const float x = path[index * 2], y = path[index * 2 + 1]; + min_x = fminf(min_x, x); + min_y = fminf(min_y, y); + max_x = fmaxf(max_x, x); + max_y = fmaxf(max_y, y); + } + bounds[0] = min_x; + bounds[1] = min_y; + bounds[2] = max_x; + bounds[3] = max_y; + } + __syncthreads(); +} + +// Solve a cubic Bezier's y(t) == y as a polynomial, rather than replacing the +// curve with a polyline. Splitting at the (at most two) derivative roots +// leaves monotonic intervals, on which a short bisection finds every crossing. +// This is deliberately a small independent implementation: the renderer only +// needs ray crossings, not a general-purpose path library. +__device__ inline float cubic_component(const float* path, int cubic, float t, int axis) { + const float u = 1.f - t; + const float* c = path + cubic * 8; + return u*u*u*c[axis] + 3.f*u*u*t*c[2 + axis] + + 3.f*u*t*t*c[4 + axis] + t*t*t*c[6 + axis]; +} + +__device__ inline float cubic_derivative(const float* path, int cubic, float t, int axis) { + const float u = 1.f - t; + const float* c = path + cubic * 8; + return 3.f*u*u*(c[2 + axis] - c[axis]) + + 6.f*u*t*(c[4 + axis] - c[2 + axis]) + + 3.f*t*t*(c[6 + axis] - c[4 + axis]); +} + +__device__ inline float cubic_hull_distance_sq(const float* path, int cubic, float px, float py) { + const float* c = path + cubic * 8; + float min_x = c[0], max_x = c[0], min_y = c[1], max_y = c[1]; + for (int point_index = 1; point_index < 4; ++point_index) { + min_x = fminf(min_x, c[point_index * 2]); max_x = fmaxf(max_x, c[point_index * 2]); + min_y = fminf(min_y, c[point_index * 2 + 1]); max_y = fmaxf(max_y, c[point_index * 2 + 1]); + } + const float dx = px < min_x ? min_x - px : (px > max_x ? px - max_x : 0.f); + const float dy = py < min_y ? min_y - py : (py > max_y ? py - max_y : 0.f); + return dx*dx + dy*dy; +} + +__device__ inline int ray_winding(const float* path, float px, float py) { + int winding = 0; + for (int cubic = 0; cubic < kCubics; ++cubic) { + const float* c = path + cubic * 8; + // dy/dt = A t^2 + B t + C. Its roots partition y(t) into monotonic + // intervals, so this finds cubic intersections without tessellation. + const float A = 3.f * (-c[1] + 3.f*c[3] - 3.f*c[5] + c[7]); + const float B = 6.f * (c[1] - 2.f*c[3] + c[5]); + const float C = 3.f * (c[3] - c[1]); + float cuts[4] = {0.f, 1.f, 1.f, 1.f}; + int cut_count = 2; + if (fabsf(A) > 1e-8f) { + const float disc = B*B - 4.f*A*C; + if (disc > 0.f) { + const float root = sqrtf(disc); + const float t0 = (-B - root) / (2.f*A); + const float t1 = (-B + root) / (2.f*A); + if (t0 > 1e-6f && t0 < 1.f-1e-6f) cuts[cut_count++] = t0; + if (t1 > 1e-6f && t1 < 1.f-1e-6f) cuts[cut_count++] = t1; + } + } else if (fabsf(B) > 1e-8f) { + const float t = -C / B; + if (t > 1e-6f && t < 1.f-1e-6f) cuts[cut_count++] = t; + } + // There are at most four cuts; insertion sort keeps the endpoint + // convention stable at extrema and shared cubic endpoints. + for (int i = 1; i < cut_count; ++i) { + float value = cuts[i]; int j = i - 1; + while (j >= 0 && cuts[j] > value) { cuts[j+1] = cuts[j]; --j; } + cuts[j+1] = value; + } + for (int interval = 0; interval + 1 < cut_count; ++interval) { + float lo = cuts[interval], hi = cuts[interval + 1]; + float yl = cubic_component(path, cubic, lo, 1) - py; + float yh = cubic_component(path, cubic, hi, 1) - py; + // Half-open interval: count a root only when the curve crosses + // the horizontal ray, never when it merely touches at an extremum. + if (!((yl <= 0.f && yh > 0.f) || (yh <= 0.f && yl > 0.f))) continue; + // 1/256 in parameter space is already subpixel-accurate on the + // 64px optimisation canvas; filtered coverage absorbs the small + // residual before export is validated by Cairo. + for (int iteration = 0; iteration < 8; ++iteration) { + const float mid = .5f * (lo + hi); + const float ym = cubic_component(path, cubic, mid, 1) - py; + if ((yl <= 0.f && ym <= 0.f) || (yl >= 0.f && ym >= 0.f)) { + lo = mid; yl = ym; + } else { + hi = mid; + } + } + const float t = .5f * (lo + hi); + if (cubic_component(path, cubic, t, 0) > px) { + winding += cubic_derivative(path, cubic, t, 1) > 0.f ? 1 : -1; + } + } + } + return winding; +} + +__global__ void coverage_forward_kernel(const float* controls, float* output, int batches, + int height, int width, int subpixels, + float x_base, float y_base, bool evenodd) { + const int pixels = height * width, batch = blockIdx.x; + if (batch >= batches) return; + const float* path = controls + batch * kCubics * 8; + __shared__ float bounds[4]; + path_bounds(path, bounds); + for (int pixel = threadIdx.x; pixel < pixels; pixel += blockDim.x) { + float coverage = 0.f; + for (int subpixel = 0; subpixel < subpixels * subpixels; ++subpixel) { + const float px = x_base + float(pixel % width) + (float(subpixel % subpixels) + .5f) / subpixels; + const float py = y_base + float(pixel / width) + (float(subpixel / subpixels) + .5f) / subpixels; + if (px >= bounds[0] && px <= bounds[2] && py >= bounds[1] && py <= bounds[3]) { + const int w = ray_winding(path, px, py); + coverage += evenodd ? float(abs(w) & 1) : float(w != 0); + } + } + output[batch * pixels + pixel] = coverage / float(subpixels * subpixels); + } +} + +__global__ void coverage_backward_kernel(const float* controls, const float* upstream, + float* gradients, int batches, int height, int width, + int subpixels, float x_base, float y_base) { + const int pixels = height * width, batch = blockIdx.x; + if (batch >= batches) return; + const float* path = controls + batch * kCubics * 8; + float* gradient = gradients + batch * kCubics * 8; + __shared__ float bounds[4]; + __shared__ float reduction[8][256]; + path_bounds(path, bounds); + float accumulated[kCubics * 8] = {0.f}; + for (int pixel = threadIdx.x; pixel < pixels; pixel += blockDim.x) { + const float d_output = upstream[batch * pixels + pixel] / float(subpixels * subpixels); + // The coverage derivative is local. Interior pixels have constant + // coverage, so only a one-pixel band around the control hull does + // useful work in the gradient pass. + const float base_x = x_base + float(pixel % width); + const float base_y = y_base + float(pixel / width); + if (base_x < bounds[0] - 2.f || base_x > bounds[2] + 2.f || + base_y < bounds[1] - 2.f || base_y > bounds[3] + 2.f) continue; + for (int subpixel = 0; subpixel < subpixels * subpixels; ++subpixel) { + const float px = base_x + (float(subpixel % subpixels) + .5f) / subpixels; + const float py = base_y + (float(subpixel / subpixels) + .5f) / subpixels; + const int winding = ray_winding(path, px, py); + const float sign = winding == 0 ? 1.f : -1.f; + float best_distance = 1e20f, best_t = 0.f; int best_cubic = 0; + for (int cubic = 0; cubic < kCubics; ++cubic) { + if (cubic_hull_distance_sq(path, cubic, px, py) >= best_distance) continue; + // A small fixed set of seeds, followed by Newton projection, + // is a boundary-local closest-point solve. It is independent + // from the forward intersection calculation and only runs in + // the antialias band. + for (int seed = 0; seed < 3; ++seed) { + float t = .5f * seed; + for (int iteration = 0; iteration < 2; ++iteration) { + const float qx = cubic_component(path, cubic, t, 0) - px; + const float qy = cubic_component(path, cubic, t, 1) - py; + const float dx = cubic_derivative(path, cubic, t, 0); + const float dy = cubic_derivative(path, cubic, t, 1); + // Gauss-Newton is stable for the short local update + // steps used by SAMVG and avoids a global curve solve. + t = fminf(1.f, fmaxf(0.f, t - (qx*dx + qy*dy) / (dx*dx + dy*dy + 1e-6f))); + } + const float qx = cubic_component(path, cubic, t, 0) - px; + const float qy = cubic_component(path, cubic, t, 1) - py; + const float distance = qx*qx + qy*qy; + if (distance < best_distance) { best_distance = distance; best_t = t; best_cubic = cubic; } + } + } + const float distance = sqrtf(best_distance + 1e-12f); + const float alpha = 1.f / (1.f + expf(sign * distance / .25f)); + const float factor = d_output * (-sign) * alpha * (1.f-alpha) / .25f / distance; + const float qx = cubic_component(path, best_cubic, best_t, 0) - px; + const float qy = cubic_component(path, best_cubic, best_t, 1) - py; + const float u = 1.f - best_t; + const float basis[4] = {u*u*u, 3.f*u*u*best_t, 3.f*u*best_t*best_t, best_t*best_t*best_t}; + for (int control = 0; control < 4; ++control) { + const int offset = best_cubic * 8 + control * 2; + accumulated[offset] += factor * qx * basis[control]; + accumulated[offset + 1] += factor * qy * basis[control]; + } + } + } + for (int cubic = 0; cubic < kCubics; ++cubic) { + for (int component = 0; component < 8; ++component) + reduction[component][threadIdx.x] = accumulated[cubic * 8 + component]; + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride /= 2) { + if (threadIdx.x < stride) + for (int component = 0; component < 8; ++component) + reduction[component][threadIdx.x] += reduction[component][threadIdx.x + stride]; + __syncthreads(); + } + if (threadIdx.x == 0) + for (int component = 0; component < 8; ++component) + gradient[cubic * 8 + component] = reduction[component][0]; + __syncthreads(); + } +} + +__device__ inline void contours_bounds(const float* controls, int first, int last, float* bounds) { + if (threadIdx.x == 0) { + const float* first_path = controls + first * kCubics * 8; + float min_x = first_path[0], min_y = first_path[1], max_x = first_path[0], max_y = first_path[1]; + for (int contour = first; contour < last; ++contour) { + const float* path = controls + contour * kCubics * 8; + for (int point_index = 0; point_index < kCubics * 4; ++point_index) { + const float x = path[point_index * 2], y = path[point_index * 2 + 1]; + min_x = fminf(min_x, x); min_y = fminf(min_y, y); + max_x = fmaxf(max_x, x); max_y = fmaxf(max_y, y); + } + } + bounds[0] = min_x; bounds[1] = min_y; bounds[2] = max_x; bounds[3] = max_y; + } + __syncthreads(); +} + +__global__ void multi_coverage_forward_kernel(const float* controls, const int64_t* offsets, + float* output, int paths, int height, int width, + int subpixels, float x_base, float y_base, + bool evenodd) { + const int path_index = blockIdx.x, pixels = height * width; + if (path_index >= paths) return; + const int first = int(offsets[path_index]), last = int(offsets[path_index + 1]); + __shared__ float bounds[4]; + contours_bounds(controls, first, last, bounds); + for (int pixel = threadIdx.x; pixel < pixels; pixel += blockDim.x) { + float coverage = 0.f; + for (int subpixel = 0; subpixel < subpixels * subpixels; ++subpixel) { + const float px = x_base + float(pixel % width) + (float(subpixel % subpixels) + .5f) / subpixels; + const float py = y_base + float(pixel / width) + (float(subpixel / subpixels) + .5f) / subpixels; + if (px < bounds[0] || px > bounds[2] || py < bounds[1] || py > bounds[3]) continue; + int winding = 0; + for (int contour = first; contour < last; ++contour) + winding += ray_winding(controls + contour * kCubics * 8, px, py); + coverage += evenodd ? float(abs(winding) & 1) : float(winding != 0); + } + output[path_index * pixels + pixel] = coverage / float(subpixels * subpixels); + } +} + +__global__ void multi_coverage_topology_forward_kernel(const float* controls, const int64_t* offsets, + float* output, uint16_t* topology, int paths, + int height, int width, int subpixels, float x_base, + float y_base, bool evenodd) { + const int path_index = blockIdx.x, pixels = height * width; + if (path_index >= paths) return; + const int first = int(offsets[path_index]), last = int(offsets[path_index + 1]); + __shared__ float bounds[4]; + contours_bounds(controls, first, last, bounds); + for (int pixel = threadIdx.x; pixel < pixels; pixel += blockDim.x) { + float coverage = 0.f; uint16_t mask = 0; + for (int subpixel = 0; subpixel < subpixels * subpixels; ++subpixel) { + const float px = x_base + float(pixel % width) + (float(subpixel % subpixels) + .5f) / subpixels; + const float py = y_base + float(pixel / width) + (float(subpixel / subpixels) + .5f) / subpixels; + int inside = 0; + if (px >= bounds[0] && px <= bounds[2] && py >= bounds[1] && py <= bounds[3]) { + int winding = 0; + for (int contour = first; contour < last; ++contour) + winding += ray_winding(controls + contour * kCubics * 8, px, py); + inside = evenodd ? (abs(winding) & 1) : (winding != 0); + } + mask |= uint16_t(inside) << subpixel; + coverage += float(inside); + } + topology[path_index * pixels + pixel] = mask; + output[path_index * pixels + pixel] = coverage / float(subpixels * subpixels); + } +} + +__global__ void multi_coverage_backward_kernel(const float* controls, const int64_t* offsets, + const float* upstream, float* gradients, int paths, + int height, int width, int subpixels, + float x_base, float y_base) { + const int path_index = blockIdx.x, pixels = height * width; + if (path_index >= paths) return; + const int first = int(offsets[path_index]), last = int(offsets[path_index + 1]); + __shared__ float bounds[4]; + contours_bounds(controls, first, last, bounds); + for (int pixel = threadIdx.x; pixel < pixels; pixel += blockDim.x) { + const float d_output = upstream[path_index * pixels + pixel] / float(subpixels * subpixels); + const float base_x = x_base + float(pixel % width), base_y = y_base + float(pixel / width); + if (base_x < bounds[0] - 2.f || base_x > bounds[2] + 2.f || + base_y < bounds[1] - 2.f || base_y > bounds[3] + 2.f) continue; + for (int subpixel = 0; subpixel < subpixels * subpixels; ++subpixel) { + const float px = base_x + (float(subpixel % subpixels) + .5f) / subpixels; + const float py = base_y + (float(subpixel / subpixels) + .5f) / subpixels; + int winding = 0; + for (int contour = first; contour < last; ++contour) + winding += ray_winding(controls + contour * kCubics * 8, px, py); + const float sign = winding == 0 ? 1.f : -1.f; + float best_distance = 1e20f, best_t = 0.f; int best_contour = first, best_cubic = 0; + for (int contour = first; contour < last; ++contour) { + const float* path = controls + contour * kCubics * 8; + for (int cubic = 0; cubic < kCubics; ++cubic) { + if (cubic_hull_distance_sq(path, cubic, px, py) >= best_distance) continue; + for (int seed = 0; seed < 3; ++seed) { + float t = .5f * seed; + for (int iteration = 0; iteration < 2; ++iteration) { + const float qx = cubic_component(path, cubic, t, 0) - px; + const float qy = cubic_component(path, cubic, t, 1) - py; + const float dx = cubic_derivative(path, cubic, t, 0); + const float dy = cubic_derivative(path, cubic, t, 1); + t = fminf(1.f, fmaxf(0.f, t - (qx*dx + qy*dy) / (dx*dx + dy*dy + 1e-6f))); + } + const float qx = cubic_component(path, cubic, t, 0) - px; + const float qy = cubic_component(path, cubic, t, 1) - py; + const float distance = qx*qx + qy*qy; + if (distance < best_distance) { best_distance = distance; best_t = t; best_contour = contour; best_cubic = cubic; } + } + } + } + const float distance = sqrtf(best_distance + 1e-12f); + const float alpha = 1.f / (1.f + expf(sign * distance / .25f)); + const float factor = d_output * (-sign) * alpha * (1.f-alpha) / .25f / distance; + const float* path = controls + best_contour * kCubics * 8; + const float qx = cubic_component(path, best_cubic, best_t, 0) - px; + const float qy = cubic_component(path, best_cubic, best_t, 1) - py; + const float u = 1.f - best_t; + const float basis[4] = {u*u*u, 3.f*u*u*best_t, 3.f*u*best_t*best_t, best_t*best_t*best_t}; + float* gradient = gradients + best_contour * kCubics * 8 + best_cubic * 8; + for (int control = 0; control < 4; ++control) { + atomicAdd(gradient + control * 2, factor * qx * basis[control]); + atomicAdd(gradient + control * 2 + 1, factor * qy * basis[control]); + } + } + } +} + +// The fill decision is discrete, and its geometry derivative is represented by +// the nearest-boundary surrogate below. Reusing the subpixel decisions from +// forward therefore avoids resolving every cubic/ray intersection again in +// backward, while preserving the same sign convention at the current step. +__global__ void multi_coverage_backward_topology_kernel( + const float* controls, const int64_t* offsets, const int64_t* boundary_offsets, + const int64_t* boundary_indices, const uint16_t* topology, const float* upstream, + float* gradients, int paths, int height, int width, int subpixels, float x_base, + float y_base) { + const int path_index = blockIdx.x, pixels = height * width; + if (path_index >= paths) return; + const int first = int(offsets[path_index]), last = int(offsets[path_index + 1]); + const int boundary_first = int(boundary_offsets[path_index]); + const int boundary_last = int(boundary_offsets[path_index + 1]); + __shared__ float bounds[4]; + contours_bounds(controls, first, last, bounds); + for (int pixel = threadIdx.x; pixel < pixels; pixel += blockDim.x) { + const float d_output = upstream[path_index * pixels + pixel] / float(subpixels * subpixels); + const float base_x = x_base + float(pixel % width), base_y = y_base + float(pixel / width); + if (base_x < bounds[0] - 2.f || base_x > bounds[2] + 2.f || + base_y < bounds[1] - 2.f || base_y > bounds[3] + 2.f) continue; + const uint16_t mask = topology[path_index * pixels + pixel]; + if (boundary_first == boundary_last) continue; + for (int subpixel = 0; subpixel < subpixels * subpixels; ++subpixel) { + const float px = base_x + (float(subpixel % subpixels) + .5f) / subpixels; + const float py = base_y + (float(subpixel / subpixels) + .5f) / subpixels; + const float sign = (mask & (uint16_t(1) << subpixel)) ? -1.f : 1.f; + float best_distance = 1e20f, best_t = 0.f; int best_contour = first, best_cubic = 0; + for (int candidate = boundary_first; candidate < boundary_last; ++candidate) { + const int local_contour = int(boundary_indices[candidate]); + // Boundary indices are local to this packed tile. Keep the + // operator safe for externally supplied candidate tensors; + // internal tile caches always satisfy this condition. + if (local_contour < 0 || first + local_contour >= last) continue; + const int contour = first + local_contour; + const float* path = controls + contour * kCubics * 8; + for (int cubic = 0; cubic < kCubics; ++cubic) { + if (cubic_hull_distance_sq(path, cubic, px, py) >= best_distance) continue; + for (int seed = 0; seed < 3; ++seed) { + float t = .5f * seed; + for (int iteration = 0; iteration < 2; ++iteration) { + const float qx = cubic_component(path, cubic, t, 0) - px; + const float qy = cubic_component(path, cubic, t, 1) - py; + const float dx = cubic_derivative(path, cubic, t, 0); + const float dy = cubic_derivative(path, cubic, t, 1); + t = fminf(1.f, fmaxf(0.f, t - (qx*dx + qy*dy) / (dx*dx + dy*dy + 1e-6f))); + } + const float qx = cubic_component(path, cubic, t, 0) - px; + const float qy = cubic_component(path, cubic, t, 1) - py; + const float distance = qx*qx + qy*qy; + if (distance < best_distance) { best_distance = distance; best_t = t; best_contour = contour; best_cubic = cubic; } + } + } + } + const float distance = sqrtf(best_distance + 1e-12f); + const float alpha = 1.f / (1.f + expf(sign * distance / .25f)); + const float factor = d_output * (-sign) * alpha * (1.f-alpha) / .25f / distance; + const float* path = controls + best_contour * kCubics * 8; + const float qx = cubic_component(path, best_cubic, best_t, 0) - px; + const float qy = cubic_component(path, best_cubic, best_t, 1) - py; + const float u = 1.f - best_t; + const float basis[4] = {u*u*u, 3.f*u*u*best_t, 3.f*u*best_t*best_t, best_t*best_t*best_t}; + float* gradient = gradients + best_contour * kCubics * 8 + best_cubic * 8; + for (int control = 0; control < 4; ++control) { + atomicAdd(gradient + control * 2, factor * qx * basis[control]); + atomicAdd(gradient + control * 2 + 1, factor * qy * basis[control]); + } + } + } +} + __global__ void forward_kernel(const float* controls, float* output, int batches, - int height, int width, int samples, float xo, float yo) { + int height, int width, int samples, int subpixels, + float x_base, float y_base) { const int pixels = height * width; const int batch = blockIdx.x; if (batch >= batches) return; const float* path = controls + batch * kCubics * 8; __shared__ float points[kCubics * kSamples * 2]; + __shared__ float bounds[4]; sample_path(path, points, samples); + path_bounds(path, bounds); for (int pixel = threadIdx.x; pixel < pixels; pixel += blockDim.x) { - const float px = xo + float(pixel % width), py = yo + float(pixel / width); - float winding = 0.f; - for (int edge = 0; edge < kCubics * samples; ++edge) { - const int next = (edge + 1) % (kCubics * samples); - const float ax = points[edge * 2] - px, ay = points[edge * 2 + 1] - py; - const float bx = points[next * 2] - px, by = points[next * 2 + 1] - py; - winding += atan2f(ax * by - ay * bx, ax * bx + ay * by); - } - output[batch * pixels + pixel] = winding; + for (int subpixel = 0; subpixel < subpixels * subpixels; ++subpixel) { + const float px = x_base + float(pixel % width) + + (float(subpixel % subpixels) + .5f) / subpixels; + const float py = y_base + float(pixel / width) + + (float(subpixel / subpixels) + .5f) / subpixels; + if (px < bounds[0] || px > bounds[2] || py < bounds[1] || py > bounds[3]) { + output[(batch * subpixels * subpixels + subpixel) * pixels + pixel] = 0.f; + continue; + } + float winding = 0.f; + for (int edge = 0; edge < kCubics * samples; ++edge) { + const int next = (edge + 1) % (kCubics * samples); + const float ax = points[edge * 2] - px, ay = points[edge * 2 + 1] - py; + const float bx = points[next * 2] - px, by = points[next * 2 + 1] - py; + winding += atan2f(ax * by - ay * bx, ax * bx + ay * by); + } + output[(batch * subpixels * subpixels + subpixel) * pixels + pixel] = winding; + } } } __global__ void backward_kernel(const float* controls, const float* upstream, float* gradients, int batches, int height, int width, - int samples, float xo, float yo) { + int samples, int subpixels, float x_base, float y_base) { const int pixels = height * width; const int batch = blockIdx.x; if (batch >= batches) return; @@ -51,33 +475,43 @@ __global__ void backward_kernel(const float* controls, const float* upstream, float* gradient = gradients + batch * kCubics * 8; __shared__ float points[kCubics * kSamples * 2]; __shared__ float reduction[8][256]; + __shared__ float bounds[4]; sample_path(path, points, samples); + path_bounds(path, bounds); float accumulated[kCubics * 8] = {0.f}; for (int pixel = threadIdx.x; pixel < pixels; pixel += blockDim.x) { - const float px = xo + float(pixel % width), py = yo + float(pixel / width); - const float d_winding = upstream[batch * pixels + pixel]; - for (int edge = 0; edge < kCubics * samples; ++edge) { - const int next = (edge + 1) % (kCubics * samples); - const float ax = points[edge * 2] - px, ay = points[edge * 2 + 1] - py; - const float bx = points[next * 2] - px, by = points[next * 2 + 1] - py; - const float cross = ax * by - ay * bx, dot = ax * bx + ay * by; - const float scale = d_winding / (cross * cross + dot * dot + 1e-20f); - const float dcross = scale * dot, ddot = -scale * cross; - const float gx_a = dcross * by + ddot * bx; - const float gy_a = -dcross * bx + ddot * by; - const float gx_b = -dcross * ay + ddot * ax; - const float gy_b = dcross * ax + ddot * ay; - const float ta = float(edge % samples) / float(samples - 1), ua = 1.f - ta; - const float tb = float(next % samples) / float(samples - 1), ub = 1.f - tb; - const float ba[4] = {ua*ua*ua, 3*ua*ua*ta, 3*ua*ta*ta, ta*ta*ta}; - const float bb[4] = {ub*ub*ub, 3*ub*ub*tb, 3*ub*tb*tb, tb*tb*tb}; - for (int control = 0; control < 4; ++control) { - const int a = (edge / samples) * 8 + control * 2; - const int b = (next / samples) * 8 + control * 2; - accumulated[a] += ba[control] * gx_a; - accumulated[a + 1] += ba[control] * gy_a; - accumulated[b] += bb[control] * gx_b; - accumulated[b + 1] += bb[control] * gy_b; + for (int subpixel = 0; subpixel < subpixels * subpixels; ++subpixel) { + const float px = x_base + float(pixel % width) + + (float(subpixel % subpixels) + .5f) / subpixels; + const float py = y_base + float(pixel / width) + + (float(subpixel / subpixels) + .5f) / subpixels; + if (px < bounds[0] || px > bounds[2] || py < bounds[1] || py > bounds[3]) { + continue; + } + const float d_winding = upstream[(batch * subpixels * subpixels + subpixel) * pixels + pixel]; + for (int edge = 0; edge < kCubics * samples; ++edge) { + const int next = (edge + 1) % (kCubics * samples); + const float ax = points[edge * 2] - px, ay = points[edge * 2 + 1] - py; + const float bx = points[next * 2] - px, by = points[next * 2 + 1] - py; + const float cross = ax * by - ay * bx, dot = ax * bx + ay * by; + const float scale = d_winding / (cross * cross + dot * dot + 1e-20f); + const float dcross = scale * dot, ddot = -scale * cross; + const float gx_a = dcross * by + ddot * bx; + const float gy_a = -dcross * bx + ddot * by; + const float gx_b = -dcross * ay + ddot * ax; + const float gy_b = dcross * ax + ddot * ay; + const float ta = float(edge % samples) / float(samples - 1), ua = 1.f - ta; + const float tb = float(next % samples) / float(samples - 1), ub = 1.f - tb; + const float ba[4] = {ua*ua*ua, 3*ua*ua*ta, 3*ua*ta*ta, ta*ta*ta}; + const float bb[4] = {ub*ub*ub, 3*ub*ub*tb, 3*ub*tb*tb, tb*tb*tb}; + for (int control = 0; control < 4; ++control) { + const int a = (edge / samples) * 8 + control * 2; + const int b = (next / samples) * 8 + control * 2; + accumulated[a] += ba[control] * gx_a; + accumulated[a + 1] += ba[control] * gy_a; + accumulated[b] += bb[control] * gx_b; + accumulated[b + 1] += bb[control] * gy_b; + } } } } @@ -111,8 +545,8 @@ torch::Tensor forward(torch::Tensor controls, int64_t height, int64_t width, int auto output = torch::zeros({controls.size(0), height, width}, controls.options()); constexpr int threads = 256; forward_kernel<<>>( - controls.data_ptr(), output.data_ptr(), controls.size(0), height, width, samples, - float(x_origin), float(y_origin)); + controls.data_ptr(), output.data_ptr(), controls.size(0), height, width, samples, 1, + float(x_origin) - .5f, float(y_origin) - .5f); C10_CUDA_KERNEL_LAUNCH_CHECK(); return output; } @@ -124,7 +558,130 @@ torch::Tensor backward(torch::Tensor controls, torch::Tensor upstream, int64_t h constexpr int threads = 256; backward_kernel<<>>( controls.data_ptr(), upstream.data_ptr(), gradients.data_ptr(), - controls.size(0), height, width, samples, float(x_origin), float(y_origin)); + controls.size(0), height, width, samples, 1, float(x_origin) - .5f, float(y_origin) - .5f); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return gradients; +} + +torch::Tensor forwards(torch::Tensor controls, int64_t height, int64_t width, + int64_t samples, int64_t subpixels, double x_origin, + double y_origin) { + TORCH_CHECK(subpixels >= 1 && subpixels <= 4); + at::cuda::CUDAGuard guard(controls.device()); + auto output = torch::zeros( + {controls.size(0), subpixels * subpixels, height, width}, controls.options()); + constexpr int threads = 256; + forward_kernel<<>>( + controls.data_ptr(), output.data_ptr(), controls.size(0), height, width, + samples, subpixels, float(x_origin), float(y_origin)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return output; +} + +torch::Tensor backwards(torch::Tensor controls, torch::Tensor upstream, int64_t height, + int64_t width, int64_t samples, int64_t subpixels, + double x_origin, double y_origin) { + at::cuda::CUDAGuard guard(controls.device()); + auto gradients = torch::zeros_like(controls); + constexpr int threads = 256; + backward_kernel<<>>( + controls.data_ptr(), upstream.data_ptr(), gradients.data_ptr(), + controls.size(0), height, width, samples, subpixels, float(x_origin), float(y_origin)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return gradients; +} + +torch::Tensor coverage_forward(torch::Tensor controls, int64_t height, int64_t width, + int64_t subpixels, double x_origin, double y_origin, + bool evenodd) { + TORCH_CHECK(controls.is_cuda() && controls.scalar_type() == torch::kFloat32); + TORCH_CHECK(subpixels >= 1 && subpixels <= 4); + at::cuda::CUDAGuard guard(controls.device()); + auto output = torch::zeros({controls.size(0), height, width}, controls.options()); + coverage_forward_kernel<<>>( + controls.data_ptr(), output.data_ptr(), controls.size(0), height, width, + subpixels, float(x_origin), float(y_origin), evenodd); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return output; +} + +torch::Tensor coverage_backward(torch::Tensor controls, torch::Tensor upstream, int64_t height, + int64_t width, int64_t subpixels, double x_origin, + double y_origin) { + at::cuda::CUDAGuard guard(controls.device()); + auto gradients = torch::zeros_like(controls); + coverage_backward_kernel<<>>( + controls.data_ptr(), upstream.data_ptr(), gradients.data_ptr(), + controls.size(0), height, width, subpixels, float(x_origin), float(y_origin)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return gradients; +} + +torch::Tensor multi_coverage_forward(torch::Tensor controls, torch::Tensor offsets, + int64_t height, int64_t width, int64_t subpixels, + double x_origin, double y_origin, bool evenodd) { + TORCH_CHECK(controls.is_cuda() && controls.scalar_type() == torch::kFloat32); + TORCH_CHECK(offsets.is_cuda() && offsets.scalar_type() == torch::kInt64); + TORCH_CHECK(offsets.dim() == 1 && offsets.size(0) >= 2); + at::cuda::CUDAGuard guard(controls.device()); + const auto paths = offsets.size(0) - 1; + auto output = torch::zeros({paths, height, width}, controls.options()); + multi_coverage_forward_kernel<<>>( + controls.data_ptr(), offsets.data_ptr(), output.data_ptr(), paths, + height, width, subpixels, float(x_origin), float(y_origin), evenodd); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return output; +} + +std::vector multi_coverage_forward_topology( + torch::Tensor controls, torch::Tensor offsets, int64_t height, int64_t width, + int64_t subpixels, double x_origin, double y_origin, bool evenodd, + torch::Tensor topology) { + at::cuda::CUDAGuard guard(controls.device()); + const auto paths = offsets.size(0) - 1; + TORCH_CHECK(topology.is_cuda() && topology.scalar_type() == torch::kUInt16); + TORCH_CHECK( + topology.dim() == 3 && topology.size(0) == paths && + topology.size(1) == height && topology.size(2) == width, + "topology workspace has the wrong shape"); + auto output = torch::zeros({paths, height, width}, controls.options()); + multi_coverage_topology_forward_kernel<<>>( + controls.data_ptr(), offsets.data_ptr(), output.data_ptr(), + topology.data_ptr(), paths, height, width, subpixels, float(x_origin), + float(y_origin), evenodd); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {output, topology}; +} + +torch::Tensor multi_coverage_backward(torch::Tensor controls, torch::Tensor offsets, + torch::Tensor upstream, int64_t height, int64_t width, + int64_t subpixels, double x_origin, double y_origin) { + at::cuda::CUDAGuard guard(controls.device()); + auto gradients = torch::zeros_like(controls); + const auto paths = offsets.size(0) - 1; + multi_coverage_backward_kernel<<>>( + controls.data_ptr(), offsets.data_ptr(), upstream.data_ptr(), + gradients.data_ptr(), paths, height, width, subpixels, float(x_origin), float(y_origin)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return gradients; +} + +torch::Tensor multi_coverage_backward_topology( + torch::Tensor controls, torch::Tensor offsets, torch::Tensor boundary_offsets, + torch::Tensor boundary_indices, torch::Tensor topology, torch::Tensor upstream, + int64_t height, int64_t width, int64_t subpixels, double x_origin, double y_origin) { + TORCH_CHECK(topology.is_cuda() && topology.scalar_type() == torch::kUInt16); + TORCH_CHECK(boundary_offsets.is_cuda() && boundary_offsets.scalar_type() == torch::kInt64); + TORCH_CHECK(boundary_indices.is_cuda() && boundary_indices.scalar_type() == torch::kInt64); + TORCH_CHECK(boundary_offsets.dim() == 1 && boundary_offsets.size(0) == offsets.size(0)); + at::cuda::CUDAGuard guard(controls.device()); + auto gradients = torch::zeros_like(controls); + const auto paths = offsets.size(0) - 1; + multi_coverage_backward_topology_kernel<<>>( + controls.data_ptr(), offsets.data_ptr(), + boundary_offsets.data_ptr(), boundary_indices.data_ptr(), + topology.data_ptr(), upstream.data_ptr(), gradients.data_ptr(), + paths, height, width, subpixels, float(x_origin), float(y_origin)); C10_CUDA_KERNEL_LAUNCH_CHECK(); return gradients; } @@ -132,4 +689,12 @@ torch::Tensor backward(torch::Tensor controls, torch::Tensor upstream, int64_t h PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("forward", &forward); m.def("backward", &backward); + m.def("forwards", &forwards); + m.def("backwards", &backwards); + m.def("coverage_forward", &coverage_forward); + m.def("coverage_backward", &coverage_backward); + m.def("multi_coverage_forward", &multi_coverage_forward); + m.def("multi_coverage_forward_topology", &multi_coverage_forward_topology); + m.def("multi_coverage_backward", &multi_coverage_backward); + m.def("multi_coverage_backward_topology", &multi_coverage_backward_topology); } diff --git a/src/vectrify/refine/cuda_renderer.py b/src/vectrify/refine/cuda_renderer.py index b5179b6e..dea6de95 100644 --- a/src/vectrify/refine/cuda_renderer.py +++ b/src/vectrify/refine/cuda_renderer.py @@ -18,6 +18,13 @@ def _extension() -> Any | None: def available() -> bool: """Whether this installation can execute the fixed-contour CUDA path.""" + # Importing Torch first loads libtorch's shared libraries before the + # optional extension is resolved. Without this, a clean wheel process can + # incorrectly report the bundled CUDA operator as unavailable. + try: + import torch # noqa: F401 + except ImportError: + return False return _extension() is not None @@ -68,3 +75,258 @@ def backward(ctx, upstream): ) return Winding.apply(controls) + + +def windings( + controls: Any, + box: tuple[int, int, int, int], + *, + samples: int, + subpixels: int, +) -> Any | None: + """Return all subpixel winding fields with one native backward reduction.""" + import torch + + extension = _extension() + if ( + extension is None + or not controls.is_cuda + or controls.dtype != torch.float32 + or controls.ndim != 4 + or controls.shape[1:] != (16, 4, 2) + or samples not in {8, 16, 32} + or subpixels not in {1, 2, 4} + ): + return None + left, top, right, bottom = box + height, width = bottom - top, right - left + + class Windings(torch.autograd.Function): + @staticmethod + def forward(ctx, value): + value = value.contiguous() + ctx.save_for_backward(value) + return extension.forwards( + value, height, width, samples, subpixels, left, top + ) + + @staticmethod + def backward(ctx, upstream): + (value,) = ctx.saved_tensors + return extension.backwards( + value, + upstream.contiguous(), + height, + width, + samples, + subpixels, + left, + top, + ) + + return Windings.apply(controls) + + +def coverage( + controls: Any, + box: tuple[int, int, int, int], + *, + subpixels: int, + fill_rule: str, +) -> Any | None: + """Analytic cubic coverage with a boundary-local differentiable pass. + + This intentionally accepts one closed contour per batch item. Callers + combine holes through the winding oracle until the native multi-contour + interface can preserve SVG fill-rule composition in one operation. + """ + import torch + + extension = _extension() + if ( + extension is None + or not controls.is_cuda + or controls.dtype != torch.float32 + or controls.ndim != 4 + or controls.shape[1:] != (16, 4, 2) + or subpixels not in {1, 2, 4} + or fill_rule not in {"nonzero", "evenodd"} + ): + return None + left, top, right, bottom = box + height, width = bottom - top, right - left + + class Coverage(torch.autograd.Function): + @staticmethod + def forward(ctx, value): + value = value.contiguous() + ctx.save_for_backward(value) + return extension.coverage_forward( + value, height, width, subpixels, left, top, fill_rule == "evenodd" + ) + + @staticmethod + def backward(ctx, upstream): + (value,) = ctx.saved_tensors + return extension.coverage_backward( + value, + upstream.contiguous(), + height, + width, + subpixels, + left, + top, + ) + + return Coverage.apply(controls) + + +def multi_coverage_forward( + controls: Any, + offsets: list[int], + box: tuple[int, int, int, int], + *, + subpixels: int, + fill_rule: str, +) -> Any | None: + """Return exact filtered multi-contour coverage without an autograd graph. + + This is used for the bounded compositing pass, whose geometry gradients + are replayed separately. A path's contour windings are combined before + the fill-rule decision, retaining holes and self-overlap semantics. + """ + import torch + + extension = _extension() + if ( + extension is None + or not controls.is_cuda + or controls.dtype != torch.float32 + or controls.ndim != 4 + or controls.shape[1:] != (16, 4, 2) + or len(offsets) < 2 + or subpixels not in {1, 2, 4} + or fill_rule not in {"nonzero", "evenodd"} + ): + return None + left, top, right, bottom = box + return extension.multi_coverage_forward( + controls.contiguous(), + torch.tensor(offsets, dtype=torch.int64, device=controls.device), + bottom - top, + right - left, + subpixels, + left, + top, + fill_rule == "evenodd", + ) + + +def multi_coverage( + controls: Any, + offsets: list[int], + box: tuple[int, int, int, int], + *, + subpixels: int, + fill_rule: str, + boundary_indices: Any | None = None, + boundary_offsets: list[int] | None = None, + topology_workspace: Any | None = None, +) -> Any | None: + """Differentiable analytic coverage for paths made of fixed contours.""" + import torch + + extension = _extension() + if ( + extension is None + or not controls.is_cuda + or controls.dtype != torch.float32 + or controls.ndim != 4 + or controls.shape[1:] != (16, 4, 2) + or len(offsets) < 2 + or subpixels not in {1, 2, 4} + or fill_rule not in {"nonzero", "evenodd"} + ): + return None + left, top, right, bottom = box + height, width = bottom - top, right - left + offset_tensor = torch.tensor(offsets, dtype=torch.int64, device=controls.device) + if boundary_indices is None: + boundary_indices = torch.cat( + [ + torch.arange( + offsets[index + 1] - offsets[index], device=controls.device + ) + for index in range(len(offsets) - 1) + ] + ) + if boundary_offsets is None: + boundary_offsets = offsets + if ( + not isinstance(boundary_indices, torch.Tensor) + or boundary_indices.dtype != torch.int64 + or not boundary_indices.is_cuda + or len(boundary_offsets) != len(offsets) + or boundary_offsets[-1] != boundary_indices.numel() + ): + return None + boundary_offset_tensor = torch.tensor( + boundary_offsets, dtype=torch.int64, device=controls.device + ) + topology_shape = (len(offsets) - 1, height, width) + if topology_workspace is None: + topology_workspace = torch.empty( + topology_shape, dtype=torch.uint16, device=controls.device + ) + if ( + not isinstance(topology_workspace, torch.Tensor) + or topology_workspace.dtype != torch.uint16 + or not topology_workspace.is_cuda + or tuple(topology_workspace.shape) != topology_shape + ): + return None + + class MultiCoverage(torch.autograd.Function): + @staticmethod + def forward(ctx, value): + value = value.contiguous() + coverage, topology = extension.multi_coverage_forward_topology( + value, + offset_tensor, + height, + width, + subpixels, + left, + top, + fill_rule == "evenodd", + topology_workspace, + ) + ctx.save_for_backward( + value, offset_tensor, boundary_offset_tensor, boundary_indices, topology + ) + return coverage + + @staticmethod + def backward(ctx, upstream): + ( + value, + saved_offsets, + saved_boundary_offsets, + saved_boundary_indices, + topology, + ) = ctx.saved_tensors + return extension.multi_coverage_backward_topology( + value, + saved_offsets, + saved_boundary_offsets, + saved_boundary_indices, + topology, + upstream.contiguous(), + height, + width, + subpixels, + left, + top, + ) + + return MultiCoverage.apply(controls) diff --git a/src/vectrify/refine/paths.py b/src/vectrify/refine/paths.py index 40608271..bd016881 100644 --- a/src/vectrify/refine/paths.py +++ b/src/vectrify/refine/paths.py @@ -818,6 +818,194 @@ def contour_winding(contour: Any, x_offset: float, y_offset: float) -> Any: return torch.stack(coverages).mean(dim=0) +def _large_path_tile_candidates( + contours: list[Any], + width: int, + height: int, + *, + tile_size: int = 16, + margin: float = 2.0, +) -> list[tuple[int, int, int, int, tuple[int, ...]]]: + """Build conservative ray-crossing candidates for a large filled path. + + A horizontal ray from a tile pixel can only cross a contour whose control + hull overlaps the tile vertically and reaches to the pixel's right. The + latter becomes ``max_x >= tile_left`` for every pixel in a tile. Cubic + curves lie inside their control hulls, making this a conservative spatial + index: it may retain an unnecessary contour but never drops a crossing. + ``margin`` also admits nearby contours to the boundary-gradient pass. + """ + if tile_size <= 0: + raise ValueError("tile_size must be positive") + bounds = [] + for contour in contours: + points = contour.detach().reshape(-1, 2) + bounds.append( + ( + float(points[:, 0].min()), + float(points[:, 1].min()), + float(points[:, 0].max()), + float(points[:, 1].max()), + ) + ) + tiles = [] + for top in range(0, height, tile_size): + for left in range(0, width, tile_size): + right = min(width, left + tile_size) + bottom = min(height, top + tile_size) + candidates = tuple( + index + for index, (_min_x, min_y, max_x, max_y) in enumerate(bounds) + if max_y >= top - margin + and min_y <= bottom + margin + and max_x >= left - margin + ) + if candidates: + tiles.append((left, top, right - left, bottom - top, candidates)) + return tiles + + +def _large_path_tile_boundary_candidates( + contours: list[Any], + tiles: list[tuple[int, int, int, int, tuple[int, ...]]], + *, + margin: float = 2.0, +) -> list[tuple[int, ...]]: + """Return nearby-contour subsets for the boundary-gradient pass. + + Winding rays must retain any contour extending to a tile's right. The + closest-boundary surrogate is local, so it only needs contours whose + conservative control hull overlaps the tile plus its antialias band. + """ + bounds = [] + for contour in contours: + points = contour.detach().reshape(-1, 2) + bounds.append( + ( + float(points[:, 0].min()), + float(points[:, 1].min()), + float(points[:, 0].max()), + float(points[:, 1].max()), + ) + ) + return [ + tuple( + index + for index in ray_candidates + if bounds[index][2] >= left - margin + and bounds[index][0] <= left + tile_width + margin + and bounds[index][3] >= top - margin + and bounds[index][1] <= top + tile_height + margin + ) + for left, top, tile_width, tile_height, ray_candidates in tiles + ] + + +def _tiled_large_path_coverage( + contours: list[Any], + box: tuple[int, int, int, int], + tiles: list[tuple[int, int, int, int, tuple[int, ...]]], + *, + fill_rule: str, + subpixels: int, + packed_contours: Any | None = None, + candidate_indices: list[Any] | None = None, + boundary_candidate_indices: list[Any] | None = None, + topology_workspaces: dict[tuple[int, int], Any] | None = None, +) -> Any | None: + """Analytically rasterise a large path from conservative contour tiles. + + ``packed_contours`` is normally a contiguous fixed-16 slice of the fit + parameter. Reusing it and the device-resident ``candidate_indices`` + avoids rebuilding the same per-tile Python concatenations each Adam step. + """ + import torch + + left, top, right, bottom = box + height, width = bottom - top, right - left + output = None + from vectrify.refine.cuda_renderer import multi_coverage + + # Tile dimensions have only edge variants. Rendering one CUDA batch per + # dimension replaces the old one-launch-per-tile graph without mixing + # candidate sets: each tile remains an independent SVG compound path. + tile_groups: dict[tuple[int, int], list[tuple[int, Any]]] = defaultdict(list) + for tile_number, tile in enumerate(tiles): + tile_groups[(tile[2], tile[3])].append((tile_number, tile)) + for (tile_width, tile_height), group in tile_groups.items(): + packed_tiles = [] + offsets = [0] + boundary_offsets = [0] + boundary_indices = [] + for tile_number, tile in group: + tile_left, tile_top, _tile_width, _tile_height, candidates = tile + offset = contours[0].new_tensor((left + tile_left, top + tile_top)) + if packed_contours is None: + packed = torch.cat( + [ + _pad_fused_cubics((contours[candidate] - offset)[None]) + for candidate in candidates + ] + ) + else: + indices = ( + candidate_indices[tile_number] + if candidate_indices is not None + else torch.tensor( + candidates, dtype=torch.long, device=packed_contours.device + ) + ) + packed = packed_contours.index_select(0, indices) - offset + packed_tiles.append(packed) + offsets.append(offsets[-1] + len(candidates)) + if boundary_candidate_indices is None: + local_boundary = torch.arange( + len(candidates), dtype=torch.long, device=packed.device + ) + else: + local_boundary = boundary_candidate_indices[tile_number] + boundary_indices.append(local_boundary) + boundary_offsets.append(boundary_offsets[-1] + local_boundary.numel()) + topology_workspace = None + if topology_workspaces is not None: + shape = (len(group), tile_height, tile_width) + topology_workspace = topology_workspaces.get((tile_width, tile_height)) + if topology_workspace is None or tuple(topology_workspace.shape) != shape: + topology_workspace = torch.empty( + shape, dtype=torch.uint16, device=packed_tiles[0].device + ) + topology_workspaces[(tile_width, tile_height)] = topology_workspace + coverage = multi_coverage( + torch.cat(packed_tiles), + offsets, + (0, 0, tile_width, tile_height), + subpixels=subpixels, + fill_rule=fill_rule, + boundary_indices=torch.cat(boundary_indices), + boundary_offsets=boundary_offsets, + topology_workspace=topology_workspace, + ) + if coverage is None: + return None + for alpha, (_tile_number, tile) in zip(coverage, group, strict=True): + tile_left, tile_top, _tile_width, _tile_height, _candidates = tile + restored = torch.nn.functional.pad( + alpha, + ( + tile_left, + width - tile_left - tile_width, + tile_top, + height - tile_top - tile_height, + ), + ) + output = restored if output is None else output + restored + if output is None: + return torch.zeros( + (height, width), dtype=contours[0].dtype, device=contours[0].device + ) + return output + + def _fill_coverages( controls: Any, box: tuple[int, int, int, int], @@ -840,6 +1028,20 @@ def _fill_coverages( left, top, right, bottom = box height, width = bottom - top, right - left + # Simple closed contours are the common SAMVG case. Use the native + # cubic-intersection renderer here; the sampled winding implementation + # below remains the portable oracle and handles arbitrary layouts. + if controls.is_cuda and controls.shape[1] <= _FUSED_CUBICS: + from vectrify.refine.cuda_renderer import coverage as cuda_coverage + + native = cuda_coverage( + _pad_fused_cubics(controls), + box, + subpixels=subpixels, + fill_rule=fill_rule, + ) + if native is not None: + return native steps = torch.linspace(0, 1, samples, device=controls.device, dtype=controls.dtype) basis = torch.stack( [ @@ -1113,7 +1315,7 @@ def fit_filled_svg( xing_weight: float = 0.02, optimisation_long_side: int | None = None, subpixels: int = 2, - monolithic: bool = False, + monolithic: bool | None = None, curve_samples: int | None = None, ) -> str: """Optimise opaque filled cubic SVG paths against an RGB target. @@ -1127,7 +1329,9 @@ def fit_filled_svg( Cairo-fidelity checks. Small clipped tiles use fewer cubic samples because their screen-space deviation is bounded by the tile size; pass ``curve_samples`` to override that adaptive choice. ``optimisation_long_side`` - is available only as an explicit caller-selected preview mode. + is available only as an explicit caller-selected preview mode. CUDA uses + one monolithic compositor graph for 64px-or-smaller working canvases by + default; larger canvases retain the memory-bounded replay. """ import xml.etree.ElementTree as ET @@ -1160,6 +1364,13 @@ def fit_filled_svg( else min(1.0, optimisation_long_side / max(width, height)) ) work_width, work_height = round(width * scale), round(height * scale) + if monolithic is None: + # At SAMVG's 64px seed-fitting resolution the complete opaque-layer + # graph is small, avoids renderer replay for each bounded layer batch, + # and has exactly the same painter-order MSE derivative. Preserve the + # bounded path at larger resolutions, where its saved alpha/canvas + # representation is intentionally memory conservative. + monolithic = device == "cuda" and work_width * work_height <= 64 * 64 # The target is resized to the integer working raster. Map coordinates # with those exact axis scales too: applying the single nominal scale to # both axes subtly shifts every horizontal edge when rounding makes the @@ -1169,20 +1380,48 @@ def fit_filled_svg( dtype=torch.float32, device=device, ) - controls = [ + initial_controls = [ [ ( torch.tensor(contour, dtype=torch.float32, device=device) * coordinate_scale - ).requires_grad_() + ) for contour in contours ] for _element, contours, _colour, _fill_rule in entries ] - colours = [ - torch.tensor(colour, dtype=torch.float32, device=device, requires_grad=True) - for _element, _contours, colour, _fill_rule in entries - ] + # A detailed SAMVG seed has hundreds of contours. Keeping each one as a + # separate Adam parameter turns one optimiser update into hundreds of tiny + # CUDA kernels. Store fixed-width contour slots in one parameter and use + # narrow views below, retaining every original contour length in the SVG + # and Xing terms. + flat_controls = [control for path in initial_controls for control in path] + contour_sizes = [len(control) for control in flat_controls] + control_storage = torch.nn.Parameter( + torch.cat([_pad_fused_cubics(control[None]) for control in flat_controls]) + ) + controls = [] + path_storage_spans = [] + storage_offset = 0 + for path in initial_controls: + path_start = storage_offset + views = [] + for _control in path: + size = contour_sizes[storage_offset] + views.append(control_storage[storage_offset, :size]) + storage_offset += 1 + controls.append(views) + path_storage_spans.append((path_start, storage_offset)) + # Match the fixed-width geometry storage above: one colour parameter + # avoids launching Adam's tiny update kernels once per SVG layer. + color_storage = torch.nn.Parameter( + torch.tensor( + [colour for _element, _contours, colour, _fill_rule in entries], + dtype=torch.float32, + device=device, + ) + ) + colours = list(color_storage.unbind(0)) goal = torch.tensor( np.asarray( target.convert("RGB").resize((work_width, work_height)), dtype=np.float32 @@ -1191,9 +1430,11 @@ def fit_filled_svg( device=device, ) point_optimizer = torch.optim.Adam( - [control for path in controls for control in path], lr=point_learning_rate + [control_storage], lr=point_learning_rate, fused=device == "cuda" + ) + colour_optimizer = torch.optim.Adam( + [color_storage], lr=color_learning_rate, fused=device == "cuda" ) - colour_optimizer = torch.optim.Adam(colours, lr=color_learning_rate) # The dissertation averages Xing within each contour then sums contours. # Keep that weighting while evaluating the 413 cat contours in one CUDA # expression rather than launching one tiny graph for each. @@ -1293,10 +1534,65 @@ def rasterise_simple( ] def rasterise_multi(index: int, path: list[Any]) -> Any: - # Very hole-heavy masks already amortise the fused full-frame kernel. - # Tile smaller multi-contour paths, where culling most of their empty - # canvas wins over compiling another specialised large batch. + # Large paths use fixed conservative candidate tiles. Every tile + # sees all contours that can cross one of its horizontal rays, while + # avoiding the old all-contours-at-every-pixel winding fallback. if len(path) >= 16: + # The index has a two-pixel conservative guard band. Rebuild it + # after local optimisation consumes half that allowance, so a + # stale index cannot exclude a valid ray crossing. + reference = large_multi_index_references[index] + movement = max( + float((control.detach() - saved).abs().amax()) + for control, saved in zip(path, reference, strict=True) + ) + if movement > 1.0: + initial_large_multi_tiles[index] = _large_path_tile_candidates( + path, work_width, work_height + ) + large_multi_tile_indices[index] = [ + torch.tensor(candidates, dtype=torch.long, device=device) + for _left, _top, _width, _height, candidates in ( + initial_large_multi_tiles[index] + ) + ] + large_multi_boundary_indices[index] = [ + torch.tensor( + [ray_candidates.index(candidate) for candidate in candidates], + dtype=torch.long, + device=device, + ) + for ( + _left, + _top, + _width, + _height, + ray_candidates, + ), candidates in zip( + initial_large_multi_tiles[index], + _large_path_tile_boundary_candidates( + path, initial_large_multi_tiles[index] + ), + strict=True, + ) + ] + large_multi_index_references[index] = tuple( + control.detach().clone() for control in path + ) + large_multi_topology_workspaces[index].clear() + tiled_alpha = _tiled_large_path_coverage( + path, + (0, 0, work_width, work_height), + initial_large_multi_tiles[index], + fill_rule=entries[index][3], + subpixels=subpixels, + packed_contours=control_storage[slice(*path_storage_spans[index])], + candidate_indices=large_multi_tile_indices[index], + boundary_candidate_indices=large_multi_boundary_indices[index], + topology_workspaces=large_multi_topology_workspaces[index], + ) + if tiled_alpha is not None: + return tiled_alpha return _fill_path_coverage( path, (0, 0, work_width, work_height), @@ -1304,7 +1600,7 @@ def rasterise_multi(index: int, path: list[Any]) -> Any: samples=samples_for(work_width, work_height), subpixels=subpixels, ) - left, top, tile_width, tile_height = tile_for(path) + left, top, tile_width, tile_height = initial_multi_tiles[index] offset = path[0].new_tensor((left, top)) alpha = _fill_path_coverage( [control - offset for control in path], @@ -1316,6 +1612,146 @@ def rasterise_multi(index: int, path: list[Any]) -> Any: ) return restore_tile(alpha, left, top) + def rasterise_multi_group( + fill_rule: str, + tile_width: int, + tile_height: int, + items: list[tuple[int, int, int]], + ) -> list[tuple[int, Any]]: + """Rasterise equal-sized multi-contour paths in one contour batch.""" + translated = [] + spans = [] + for index, left, top in items: + offset = controls[index][0].new_tensor((left, top)) + start = len(translated) + translated.extend(control - offset for control in controls[index]) + spans.append((start, len(translated))) + # Native winding uses one CUDA block per contour. Combining contours + # from otherwise independent paths lets its blocks occupy the GPU at + # once, while summing each recorded span before the fill nonlinearity + # preserves SVG path semantics (including holes). + packed = torch.cat([_pad_fused_cubics(control[None]) for control in translated]) + from vectrify.refine.cuda_renderer import multi_coverage + + analytic = multi_coverage( + packed, + [start for start, _end in spans] + [spans[-1][1]], + (0, 0, tile_width, tile_height), + subpixels=subpixels, + fill_rule=fill_rule, + ) + if analytic is not None: + return [ + (index, restore_tile(alpha, left, top)) + for (index, left, top), alpha in zip(items, analytic, strict=True) + ] + from vectrify.refine.cuda_renderer import windings as cuda_windings + + native_winding = cuda_windings( + packed, + (0, 0, tile_width, tile_height), + samples=samples_for(tile_width, tile_height), + subpixels=subpixels, + ) + if native_winding is not None: + path_winding = torch.stack( + [native_winding[start:end].sum(dim=0) for start, end in spans] + ) + if fill_rule == "evenodd": + coverage = 0.5 * (1 - torch.cos(path_winding / 2)) + else: + coverage = torch.sigmoid((path_winding.abs() - math.pi) / 0.25) + return [ + (index, restore_tile(alpha, left, top)) + for (index, left, top), alpha in zip( + items, coverage.mean(dim=1), strict=True + ) + ] + coverage_sum = None + for y in range(subpixels): + for x in range(subpixels): + winding = _fill_batched_windings( + packed, + (0, 0, tile_width, tile_height), + samples=samples_for(tile_width, tile_height), + x_offset=(x + 0.5) / subpixels, + y_offset=(y + 0.5) / subpixels, + batch_size=64, + ) + path_winding = torch.stack( + [winding[start:end].sum(dim=0) for start, end in spans] + ) + if fill_rule == "evenodd": + coverage = 0.5 * (1 - torch.cos(path_winding / 2)) + else: + coverage = torch.sigmoid((path_winding.abs() - math.pi) / 0.25) + coverage_sum = ( + coverage if coverage_sum is None else coverage_sum + coverage + ) + assert coverage_sum is not None + return [ + (index, restore_tile(alpha, left, top)) + for (index, left, top), alpha in zip( + items, + coverage_sum / (subpixels * subpixels), + strict=True, + ) + ] + + # Tile layout is part of the seed rasterisation setup, not optimisation + # state. Re-reading each CUDA control tensor's extrema every Adam step + # introduces hundreds of device synchronisations on a detailed SAMVG + # seed. The two-pixel antialias margin already makes these fixed tiles + # conservative for the local coordinate updates used by the fit. + initial_simple_groups = cropped_simple_groups() + initial_multi_tiles = { + index: tile_for(path) + for index, path in enumerate(controls) + if len(path) != 1 and len(path) < 16 + } + initial_large_multi_tiles = { + index: _large_path_tile_candidates(path, work_width, work_height) + for index, path in enumerate(controls) + if len(path) >= 16 + } + large_multi_tile_indices = { + index: [ + torch.tensor(candidates, dtype=torch.long, device=device) + for _left, _top, _width, _height, candidates in tiles + ] + for index, tiles in initial_large_multi_tiles.items() + } + large_multi_boundary_indices = { + index: [ + torch.tensor( + [ray_candidates.index(candidate) for candidate in candidates], + dtype=torch.long, + device=device, + ) + for (_left, _top, _width, _height, ray_candidates), candidates in zip( + tiles, + _large_path_tile_boundary_candidates(controls[index], tiles), + strict=True, + ) + ] + for index, tiles in initial_large_multi_tiles.items() + } + large_multi_topology_workspaces: dict[int, dict[tuple[int, int], Any]] = { + index: {} for index in initial_large_multi_tiles + } + large_multi_index_references = { + index: tuple(control.detach().clone() for control in path) + for index, path in enumerate(controls) + if len(path) >= 16 + } + initial_multi_groups: dict[tuple[str, int, int], list[tuple[int, int, int]]] = ( + defaultdict(list) + ) + for index, (left, top, tile_width, tile_height) in initial_multi_tiles.items(): + initial_multi_groups[(entries[index][3], tile_width, tile_height)].append( + (index, left, top) + ) + log.info( "Filled-path optimisation: %d path(s), %dx%d working raster on %s.", len(entries), @@ -1326,7 +1762,7 @@ def rasterise_multi(index: int, path: list[Any]) -> Any: for _step in range(steps): point_optimizer.zero_grad() colour_optimizer.zero_grad() - simple_groups = cropped_simple_groups() + simple_groups = initial_simple_groups all_controls = torch.cat([control for path in controls for control in path]) if monolithic: @@ -1341,6 +1777,15 @@ def rasterise_multi(index: int, path: list[Any]) -> Any: fill_rule, tile_width, tile_height, items ): alphas[index] = alpha + for ( + fill_rule, + tile_width, + tile_height, + ), items in initial_multi_groups.items(): + for index, alpha in rasterise_multi_group( + fill_rule, tile_width, tile_height, items + ): + alphas[index] = alpha for index, path in enumerate(controls): if alphas[index] is None: alphas[index] = rasterise_multi(index, path) @@ -1350,7 +1795,7 @@ def rasterise_multi(index: int, path: list[Any]) -> Any: if goal.is_cuda else _composite_opaque_fills ) - rendered = composite(alpha_stack, torch.stack(colours)) + rendered = composite(alpha_stack, color_storage) loss = ((rendered - goal) ** 2).mean() loss = ( loss @@ -1377,6 +1822,15 @@ def rasterise_multi(index: int, path: list[Any]) -> Any: fill_rule, tile_width, tile_height, items ): initial_alphas[index] = alpha + for ( + fill_rule, + tile_width, + tile_height, + ), items in initial_multi_groups.items(): + for index, alpha in rasterise_multi_group( + fill_rule, tile_width, tile_height, items + ): + initial_alphas[index] = alpha for index, path in enumerate(controls): if initial_alphas[index] is None: initial_alphas[index] = rasterise_multi(index, path) @@ -1444,8 +1898,22 @@ def layer_loss( simple_indices = { index for group in simple_groups.values() for index, _left, _top in group } + multi_indices = { + index + for group in initial_multi_groups.values() + for index, _left, _top in group + } + for (fill_rule, tile_width, tile_height), items in initial_multi_groups.items(): + for offset in range(0, len(items), 64): + batch = items[offset : offset + 64] + loss = torch.zeros((), device=device) + for index, alpha in rasterise_multi_group( + fill_rule, tile_width, tile_height, batch + ): + loss = loss + layer_loss(index, alpha) + loss.backward() for index, path in enumerate(controls): - if index not in simple_indices: + if index not in simple_indices and index not in multi_indices: layer_loss( index, rasterise_multi(index, path), diff --git a/tests/refine/test_filled_paths.py b/tests/refine/test_filled_paths.py index 9acc496f..a9c5889d 100644 --- a/tests/refine/test_filled_paths.py +++ b/tests/refine/test_filled_paths.py @@ -11,6 +11,10 @@ _fill_coverage, _fill_coverages, _fill_path_coverage, + _large_path_tile_boundary_candidates, + _large_path_tile_candidates, + _pad_fused_cubics, + _tiled_large_path_coverage, _xing_loss, fit_filled_svg, parse_filled_cubics, @@ -75,6 +79,40 @@ def test_native_winding_falls_back_without_the_optional_extension(monkeypatch): ) +def test_native_subpixel_windings_share_the_separate_winding_result(): + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + from vectrify.refine.cuda_renderer import available, winding, windings + + if not available(): + pytest.skip("optional SAMVG CUDA extension is not installed") + controls = _sixteen_cubic_circle(torch).requires_grad_() + fused = windings(controls, (0, 0, 24, 24), samples=16, subpixels=2) + separate = torch.stack( + [ + winding( + controls, + (0, 0, 24, 24), + samples=16, + x_offset=(x + 0.5) / 2, + y_offset=(y + 0.5) / 2, + ) + for y in range(2) + for x in range(2) + ], + dim=1, + ) + upstream = torch.randn_like(fused) + (fused * upstream).sum().backward() + fused_gradient = controls.grad.detach().clone() + controls.grad = None + (separate * upstream).sum().backward() + + assert torch.allclose(fused, separate, atol=1e-5, rtol=1e-5) + assert torch.allclose(fused_gradient, controls.grad, atol=1e-4, rtol=1e-4) + + def test_native_even_odd_coverage_stays_cairo_validated(): """The native winding path preserves SVG hole coverage, not just tensors.""" torch = pytest.importorskip("torch") @@ -113,6 +151,77 @@ def test_native_even_odd_coverage_stays_cairo_validated(): assert np.abs(native - real).mean() < 0.002 assert native[48, 48] < 0.01 + +def test_native_analytic_cubic_coverage_stays_cairo_validated(): + """The production single-contour path uses cubic intersections, not samples.""" + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + from vectrify.refine.cuda_renderer import available + + if not available(): + pytest.skip("optional SAMVG CUDA extension is not installed") + size = 96 + contour = torch.tensor( + parse_filled_cubics("M 12 48 C 12 5 84 5 84 48 C 84 91 12 91 12 48 Z")[0], + dtype=torch.float32, + device="cuda", + ) + contour = torch.cat((contour, contour[:1].expand(16 - len(contour), -1, -1)))[ + None + ] + native = _fill_coverages(contour, (0, 0, size, size), subpixels=4).cpu().numpy() + head = f'' + blank = f"{head}" + path = "M 12 48 C 12 5 84 5 84 48 C 84 91 12 91 12 48 Z" + drawn = f'{head}' + real = ( + np.asarray( + Image.open( + io.BytesIO(rasterize_svg_to_png_bytes(blank, out_w=size, out_h=size)) + ).convert("L"), + dtype=np.float32, + ) + - np.asarray( + Image.open( + io.BytesIO(rasterize_svg_to_png_bytes(drawn, out_w=size, out_h=size)) + ).convert("L"), + dtype=np.float32, + ) + ) / 255.0 + + assert np.abs(native - real).mean() < 0.002 + + +def test_native_analytic_multi_contour_coverage_preserves_a_hole(): + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + from vectrify.refine.cuda_renderer import available, multi_coverage + + if not available(): + pytest.skip("optional SAMVG CUDA extension is not installed") + contours = [ + torch.tensor(contour, dtype=torch.float32, device="cuda") + for contour in parse_filled_cubics(DONUT_PATH) + ] + controls = torch.cat( + [ + torch.cat((contour, contour[:1].expand(16 - len(contour), -1, -1)))[ + None + ] + for contour in contours + ] + ).requires_grad_() + coverage = multi_coverage( + controls, [0, 2], (0, 0, 96, 96), subpixels=4, fill_rule="evenodd" + ) + assert coverage is not None + assert coverage[0, 48, 48] < 0.01 + coverage.sum().backward() + assert controls.grad is not None + assert controls.grad.abs().sum() > 0 + SVG = ( '' ' 0 + + +def test_analytic_multi_coverage_reuses_topology_workspace_between_steps(): + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + from vectrify.refine.cuda_renderer import available, multi_coverage + + if not available(): + pytest.skip("optional SAMVG CUDA extension is not installed") + controls = torch.tensor( + [[[[0.0, 0.0], [0.0, 4.0], [4.0, 4.0], [4.0, 0.0]]] * 16], + device="cuda", + requires_grad=True, + ) + workspace = torch.empty((1, 4, 4), dtype=torch.uint16, device="cuda") + pointer = workspace.data_ptr() + for _ in range(2): + controls.grad = None + coverage = multi_coverage( + controls, + [0, 1], + (0, 0, 4, 4), + subpixels=2, + fill_rule="nonzero", + topology_workspace=workspace, + ) + assert coverage is not None + coverage.sum().backward() + assert torch.isfinite(controls.grad).all() + assert workspace.data_ptr() == pointer + + +@pytest.mark.parametrize("fill_rule", ["evenodd", "nonzero"]) +def test_tiled_analytic_large_compound_path_matches_cairo(fill_rule): + """The 16-contour production tile path preserves holes in Cairo terms.""" + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + from vectrify.refine.cuda_renderer import available + + if not available(): + pytest.skip("optional SAMVG CUDA extension is not installed") + size = 64 + pieces = [] + for row in range(2): + for column in range(4): + x, y = 2 + column * 16, 6 + row * 28 + # Reverse the inner contour, so it is a hole for both SVG rules. + pieces.extend( + ( + f"M {x} {y} L {x + 12} {y} L {x + 12} {y + 12} L {x} {y + 12} Z", + ( + f"M {x + 3} {y + 3} L {x + 3} {y + 9} " + f"L {x + 9} {y + 9} L {x + 9} {y + 3} Z" + ), + ) + ) + path = " ".join(pieces) + contours = [ + torch.tensor(contour, dtype=torch.float32, device="cuda") + for contour in parse_filled_cubics(path) + ] + assert len(contours) == 16 + padded = [_pad_fused_cubics(contour[None])[0] for contour in contours] + native = _tiled_large_path_coverage( + padded, + (0, 0, size, size), + _large_path_tile_candidates(padded, size, size), + fill_rule=fill_rule, + subpixels=4, + ) + assert native is not None + head = f'' + blank = f"{head}" + drawn = f'{head}' + cairo = ( + np.asarray( + Image.open( + io.BytesIO(rasterize_svg_to_png_bytes(blank, out_w=size, out_h=size)) + ).convert("L"), + dtype=np.float32, + ) + - np.asarray( + Image.open( + io.BytesIO(rasterize_svg_to_png_bytes(drawn, out_w=size, out_h=size)) + ).convert("L"), + dtype=np.float32, + ) + ) / 255.0 + rendered = native.detach().cpu().numpy() + assert np.abs(rendered - cairo).mean() < 0.002 + assert rendered[12, 8] < 0.01 + + +def test_filled_fit_uses_analytic_tiles_for_large_cuda_paths(monkeypatch): + """A supported large path must not silently enter sampled winding.""" + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + from vectrify.refine import cuda_renderer + from vectrify.refine import paths as filled_paths + + if not cuda_renderer.available(): + pytest.skip("optional SAMVG CUDA extension is not installed") + pieces = [] + for row in range(2): + for column in range(8): + x, y = 2 + column * 7, 8 + row * 24 + pieces.append( + f"M {x} {y} L {x + 5} {y} L {x + 5} {y + 5} L {x} {y + 5} Z" + ) + svg = ( + '' + f'' + "" + ) + + def unexpected_sampled_fallback(*_args, **_kwargs): + raise AssertionError("supported large CUDA path entered sampled winding") + + monkeypatch.setattr( + filled_paths, "_fill_path_coverage", unexpected_sampled_fallback + ) + fitted = fit_filled_svg( + svg, + Image.new("RGB", (64, 64), "white"), + steps=1, + point_learning_rate=0.0, + color_learning_rate=0.0, + optimisation_long_side=64, + ) + assert "path" in fitted + + def test_bounded_compositing_gradient_matches_monolithic_render(): """The memory-bounded fit pass must retain the full painter's-order MSE gradient.""" torch = pytest.importorskip("torch") From e3eb0289ff371623eec26b4951f6b6f8205a67ab Mon Sep 17 00:00:00 2001 From: Rasmus Ros Date: Mon, 24 Aug 2026 10:31:30 +0200 Subject: [PATCH 3/6] style: format Python sources --- src/vectrify/formats/svg/operations.py | 4 +--- src/vectrify/image_utils.py | 4 +--- src/vectrify/score/ensemble.py | 3 +-- src/vectrify/search/engine.py | 6 +----- src/vectrify/vector/runner.py | 4 +--- 5 files changed, 5 insertions(+), 16 deletions(-) diff --git a/src/vectrify/formats/svg/operations.py b/src/vectrify/formats/svg/operations.py index 53cc69e3..b57f924c 100644 --- a/src/vectrify/formats/svg/operations.py +++ b/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 diff --git a/src/vectrify/image_utils.py b/src/vectrify/image_utils.py index 262588ab..7bae4310 100644 --- a/src/vectrify/image_utils.py +++ b/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 diff --git a/src/vectrify/score/ensemble.py b/src/vectrify/score/ensemble.py index d1eafa7b..d2fcd325 100644 --- a/src/vectrify/score/ensemble.py +++ b/src/vectrify/score/ensemble.py @@ -154,8 +154,7 @@ def rank( return [] images = [ - self._decode(png) - or Image.new("RGB", reference.image.size, (255, 255, 255)) + self._decode(png) or Image.new("RGB", reference.image.size, (255, 255, 255)) for png in candidate_pngs ] diff --git a/src/vectrify/search/engine.py b/src/vectrify/search/engine.py index 2fcd5c85..f8ebc150 100644 --- a/src/vectrify/search/engine.py +++ b/src/vectrify/search/engine.py @@ -592,11 +592,7 @@ def _make_node(res: Result, *, new_lineage: bool = False) -> SearchNode[TState]: node_id = run_state.next_node_id # An LLM seed is an independent attempt at the picture, so it opens # a lineage; a local child continues its parent's. - root = ( - node_id - if new_lineage - else node_roots.get(res.parent_id, node_id) - ) + root = node_id if new_lineage else node_roots.get(res.parent_id, node_id) node_roots[node_id] = root origin = node_origins.get(res.parent_id) or node_id node_origins[node_id] = origin diff --git a/src/vectrify/vector/runner.py b/src/vectrify/vector/runner.py index 830ac9d8..fbee3425 100644 --- a/src/vectrify/vector/runner.py +++ b/src/vectrify/vector/runner.py @@ -478,9 +478,7 @@ def run_vector_search( log.info("SAMVG-inspired seed skipped: it is available for SVG only.") else: try: - content = format_plugin.extract_from_llm( - generate_svg(original_img) - ) + content = format_plugin.extract_from_llm(generate_svg(original_img)) valid, error = format_plugin.validate(content) if not valid: raise ValueError(error or "generated SVG failed validation") From e59ac9dfae29f52546a150bd248649fc5afb6479 Mon Sep 17 00:00:00 2001 From: Rasmus Ros Date: Mon, 24 Aug 2026 10:32:31 +0200 Subject: [PATCH 4/6] chore: ignore build staging --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 63a65064..a3b109b4 100644 --- a/.gitignore +++ b/.gitignore @@ -6,6 +6,7 @@ .venv/ .idea/ *.egg-info +/build/ /output /models # Run output: vectrify writes /runs/ next to the output file From e083cade96627096832dfbdb1cca88347477e1ba Mon Sep 17 00:00:00 2001 From: Rasmus Ros Date: Mon, 24 Aug 2026 10:37:06 +0200 Subject: [PATCH 5/6] fix: resolve pyrefly errors --- src/vectrify/refine/cuda_renderer.py | 12 ++++++++---- src/vectrify/refine/paths.py | 15 +++++++++------ src/vectrify/refine/samvg.py | 17 +++++++++++------ 3 files changed, 28 insertions(+), 16 deletions(-) diff --git a/src/vectrify/refine/cuda_renderer.py b/src/vectrify/refine/cuda_renderer.py index dea6de95..5cb25893 100644 --- a/src/vectrify/refine/cuda_renderer.py +++ b/src/vectrify/refine/cuda_renderer.py @@ -62,8 +62,9 @@ def forward(ctx, value): ) @staticmethod - def backward(ctx, upstream): + def backward(ctx: Any, *upstreams: Any) -> Any: (value,) = ctx.saved_tensors + upstream = upstreams[0] return extension.backward( value, upstream.contiguous(), @@ -111,8 +112,9 @@ def forward(ctx, value): ) @staticmethod - def backward(ctx, upstream): + def backward(ctx: Any, *upstreams: Any) -> Any: (value,) = ctx.saved_tensors + upstream = upstreams[0] return extension.backwards( value, upstream.contiguous(), @@ -166,8 +168,9 @@ def forward(ctx, value): ) @staticmethod - def backward(ctx, upstream): + def backward(ctx: Any, *upstreams: Any) -> Any: (value,) = ctx.saved_tensors + upstream = upstreams[0] return extension.coverage_backward( value, upstream.contiguous(), @@ -307,7 +310,7 @@ def forward(ctx, value): return coverage @staticmethod - def backward(ctx, upstream): + def backward(ctx: Any, *upstreams: Any) -> Any: ( value, saved_offsets, @@ -315,6 +318,7 @@ def backward(ctx, upstream): saved_boundary_indices, topology, ) = ctx.saved_tensors + upstream = upstreams[0] return extension.multi_coverage_backward_topology( value, saved_offsets, diff --git a/src/vectrify/refine/paths.py b/src/vectrify/refine/paths.py index bd016881..6760e34f 100644 --- a/src/vectrify/refine/paths.py +++ b/src/vectrify/refine/paths.py @@ -798,17 +798,16 @@ def contour_winding(contour: Any, x_offset: float, y_offset: float) -> Any: coverages = [] for y in range(subpixels): for x in range(subpixels): - winding = sum( - ( + winding = torch.stack( + [ contour_winding( contour, (x + 0.5) / subpixels, (y + 0.5) / subpixels, ) for contour in contours - ), - start=0, - ) + ] + ).sum(dim=0) if fill_rule == "evenodd": # Winding changes by 2π for every crossing. This periodic # expression is zero for an even count and one for an odd one. @@ -1264,7 +1263,11 @@ def _fill_rgb(value: str | None) -> tuple[float, float, float] | None: if not match: return None digits = match.group(1) - return tuple(int(digits[index : index + 2], 16) / 255 for index in range(0, 6, 2)) + return ( + int(digits[0:2], 16) / 255, + int(digits[2:4], 16) / 255, + int(digits[4:6], 16) / 255, + ) def _composite_opaque_fills(alphas: Any, colours: Any) -> Any: diff --git a/src/vectrify/refine/samvg.py b/src/vectrify/refine/samvg.py index 79790ab2..115ed43b 100644 --- a/src/vectrify/refine/samvg.py +++ b/src/vectrify/refine/samvg.py @@ -15,6 +15,7 @@ import xml.etree.ElementTree as ET from collections import defaultdict from dataclasses import dataclass +from typing import cast import numpy as np from PIL import Image @@ -183,8 +184,9 @@ def recolour_visible_layers( visible = layer.mask & ~covered_above colour = layer.colour if visible.any(): - colour = tuple( - int(value) for value in np.rint(target[visible].mean(axis=0)) + colour = cast( + tuple[int, int, int], + tuple(int(value) for value in np.rint(target[visible].mean(axis=0))), ) revised.append( MaskLayer(layer.mask, colour, layer.impact, layer.overlap_pixels) @@ -246,7 +248,10 @@ def filter_by_impact( for mask in candidates: if int(mask.sum()) < min_pixels: continue - colour = tuple(int(v) for v in np.rint(target[mask].mean(axis=0))) + colour = cast( + tuple[int, int, int], + tuple(int(value) for value in np.rint(target[mask].mean(axis=0))), + ) next_canvas = canvas.copy() next_coverage = coverage | mask next_canvas[mask] = colour @@ -277,7 +282,7 @@ def coverage_prompt_points( _canvas, coverage = _render_layers(shape, layers) radius = max(2, round(min(shape) * radius_fraction)) - distance = distance_transform_edt(~coverage) + distance = np.asarray(distance_transform_edt(~coverage)) ys, xs = np.nonzero(distance >= radius) if len(xs) == 0: return [] @@ -393,7 +398,7 @@ def _loops(mask: np.ndarray) -> list[list[tuple[float, float]]]: loops: list[list[tuple[float, float]]] = [] while edges: start = next(iter(edges)) - current, loop = start, [tuple(map(float, start))] + current, loop = start, [cast(tuple[float, float], tuple(map(float, start)))] while current in edges: following = edges[current].pop() if not edges[current]: @@ -401,7 +406,7 @@ def _loops(mask: np.ndarray) -> list[list[tuple[float, float]]]: current = following if current == start: break - loop.append(tuple(map(float, current))) + loop.append(cast(tuple[float, float], tuple(map(float, current)))) if current == start and len(loop) >= 3: loops.append(loop) return loops From b9ab4dbe3041662a8829c3ceb5beb6ea95592f61 Mon Sep 17 00:00:00 2001 From: Rasmus Ros Date: Mon, 24 Aug 2026 10:52:10 +0200 Subject: [PATCH 6/6] fix: avoid persistent color parameter views --- src/vectrify/refine/paths.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/vectrify/refine/paths.py b/src/vectrify/refine/paths.py index 6760e34f..10b85c73 100644 --- a/src/vectrify/refine/paths.py +++ b/src/vectrify/refine/paths.py @@ -1424,7 +1424,6 @@ def fit_filled_svg( device=device, ) ) - colours = list(color_storage.unbind(0)) goal = torch.tensor( np.asarray( target.convert("RGB").resize((work_width, work_height)), dtype=np.float32 @@ -1840,8 +1839,9 @@ def rasterise_multi_group( before: list[Any] = [] rendered = torch.zeros_like(goal) - for alpha, colour in zip(initial_alphas, colours, strict=True): + for index, alpha in enumerate(initial_alphas): assert alpha is not None + colour = color_storage[index] before.append(rendered) rendered = ( rendered * (1 - alpha[..., None]) @@ -1871,7 +1871,7 @@ def layer_loss( suffix = suffixes[index] assert stored_alpha is not None assert suffix is not None - colour = colours[index] + colour = color_storage[index] colour_delta = colour.detach().clamp(0, 1) - canvases[index] alpha_gradient = (gradient * suffix[..., None] * colour_delta).sum(dim=-1) colour_gradient = ( @@ -1928,9 +1928,10 @@ def layer_loss( colour_optimizer.step() coordinate_scale_cpu = coordinate_scale.cpu() - for (element, _contours, _colour, _fill_rule), path, colour in zip( - entries, controls, colours, strict=True + for index, ((element, _contours, _colour, _fill_rule), path) in enumerate( + zip(entries, controls, strict=True) ): + colour = color_storage[index] data = " ".join( to_path_d((control.detach().cpu() / coordinate_scale_cpu).tolist()) + " Z" for control in path