colortransfer-triton

简体中文 · Installation · API · Methods · Benchmark

ColorTransfer Triton implements AdaIN, Wavelet, Sliced OT, Sinkhorn and ColorFM-O color transfer with Triton GPU kernels and a PyTorch tensor API. No pretrained weights.

Qualitative examples

Figure 1. Aligned pairs: Astronaut cool-cast recovery and Chelsea warm grading.

Figure 2. Different-content pairs: Astronaut → Coffee and Coffee → Chelsea.

Get started

Linux, Python 3.10+, NVIDIA GPU, CUDA PyTorch 2.6+ and compatible Triton 3.2+.

python -m pip install -e .

Python

import torch
import colortransfer_triton as ct

source = torch.rand(1, 3, 512, 512, device="cuda")  # replace with RGB images
reference = torch.rand_like(source)
result = ct.adain(source, reference)
result = ct.transfer(source, reference, method="sliced_ot")

lut = ct.fit_colorfm(source, reference, steps=700)
result = lut(source, strength=0.8)  # reuse without fitting again

The image API is inference-only; detach inputs that require gradients.

Methods

Wavelet uses an undecimated, dilated binomial pyramid, not a Haar/DWT transform. Its reference must have the same dimensions and spatial alignment. The other methods match color distributions without spatial correspondence or semantic object matching.

ColorFM implements the nonsemantic ColorFM-O variant. It trains a small MLP from scratch for each pair using sampled RGB colors: no pretrained weights or segmentation model, but an optimization stage is required. AdaIN, Wavelet and the OT methods do not train a neural network.

OT/ColorFM use sampled fitting and a 3D LUT approximation. Reuse maps with ct.fit_lut or ct.fit_colorfm. strength controls blending; clamp=True clips to [0,1].

Reuse a color map

lut = ct.fit_lut(
    source,
    reference,
    method="sinkhorn",
    samples=1024,
    lut_size=33,
    epsilon=0.03,
    iterations=100,
    seed=0,
)
result = lut(source, strength=0.8)

Fitting runs on sampled colors; the full image uses Triton trilinear LUT lookup with fused blending and clipping. This avoids a full-resolution pixel-to-pixel transport matrix on large images. A fitted map can be applied at other resolutions, but remains specific to the fitted color distributions; refit when the scene or reference changes.

Key parameters

Sinkhorn convergence is available in lut.diagnostics as marginal_error; a fixed iteration count does not guarantee convergence.

Reproduce figures

python -m pip install -e ".[figures]"
python examples/color_transfer.py --suite --output assets --save-images
python examples/color_transfer.py --source source.png --reference reference.png --save-images

Benchmark

python benchmarks/benchmark.py --height 2160 --width 3840 --repeats 30
python benchmarks/benchmark.py --colorfm --repeats 5

Tests

python -m pip install -e ".[test,figures]"
python -m pytest -q

License

MIT. Example photographs from scikit-image: Astronaut (NASA, public domain), Chelsea (Stefan van der Walt, CC0), Coffee (Rachel Michetti, CC0).

Publishing

Version tags automatically build and publish to PyPI after CI passes. See RELEASING.md for one-time Trusted Publisher setup and release steps.