device: torch.device,

dtype: torch.dtype = torch.float32,

sigma: float = 1.0,

double_size: bool = False) -> torch.Tensor:

"""Generate isotropic Gaussian-blurred uniform noise.

gen_w, gen_h = ((width + 1) // 2, (height + 1) // 2) if double_size else (width, height)

noise = torch.rand(gen_h, gen_w, device=device, dtype=dtype)

if sigma > 0.0:

noise = noise.unsqueeze(0) # [h, w] -> [1, h, w]

ks = _blur_kernel_size(sigma)

noise = torchvision.transforms.GaussianBlur((ks, 1), sigma=sigma)(noise)

noise = torchvision.transforms.GaussianBlur((1, ks), sigma=sigma)(noise)

noise = noise.squeeze(0) # -> [h, w]

if double_size:

noise = noise.repeat_interleave(2, dim=-1).repeat_interleave(2, dim=-2)[:height, :width]

device: torch.device,

dtype: torch.dtype = torch.float32,

mode: str = "PAL",

ntsc_chroma: str = "nearest",

double_size: bool = False) -> torch.Tensor:

"""Generate VHS noise.

gen_w, gen_h = ((width + 1) // 2, (height + 1) // 2) if double_size else (width, height)

# Y channel (luma): fine-grained horizontal runs.

noise_y = torch.rand(gen_h, gen_w, device=device, dtype=dtype).unsqueeze(0) # [1, H, W]

noise_y = torchvision.transforms.GaussianBlur((5, 1), sigma=2.0)(noise_y)

if mode == "PAL":

result = noise_y

elif mode == "NTSC":

# NTSC: independent U/V planes with coarser horizontal blur

# (lower chroma bandwidth → wider, smoother runs).

#

# 4:2:0 chroma subsampling: generate U/V at half luma resolution and

# upscale. Human vision has low chroma spatial acuity, and the coarser

# chroma blocks dramatically improve delta-based compression (QOI) —

# random per-pixel chroma is the main entropy source.

chroma_h = (gen_h + 1) // 2

chroma_w = (gen_w + 1) // 2

noise_u = torch.rand(chroma_h, chroma_w, device=device, dtype=dtype).unsqueeze(0)

noise_u = torchvision.transforms.GaussianBlur((11, 1), sigma=4.0)(noise_u)

noise_u -= 0.5

noise_v = torch.rand(chroma_h, chroma_w, device=device, dtype=dtype).unsqueeze(0)

noise_v = torchvision.transforms.GaussianBlur((11, 1), sigma=4.0)(noise_v)

noise_v -= 0.5

# Upscale chroma to luma resolution.

chroma = torch.cat([noise_u, noise_v], dim=0) # [2, chroma_h, chroma_w]

if ntsc_chroma == "bilinear": # analog aesthetic

chroma = torch.nn.functional.interpolate(

chroma.unsqueeze(0), size=(gen_h, gen_w), mode="bilinear", align_corners=False,

).squeeze(0)

else: # ntsc_chroma == "nearest": # retro digital aesthetic

chroma = (chroma.repeat_interleave(2, dim=-1)

.repeat_interleave(2, dim=-2)[:, :gen_h, :gen_w])

noise_u = chroma[0:1] # [1, H, W]

noise_v = chroma[1:2] # [1, H, W]

result = torch.cat([noise_y, noise_u, noise_v], dim=0) # [3, H, W]

else:

raise ValueError(f"vhs_noise: unknown mode {mode!r}; expected 'PAL' or 'NTSC'")

if double_size:

result = result.repeat_interleave(2, dim=-1).repeat_interleave(2, dim=-2)[..., :height, :width]

device: torch.device,

dtype: torch.dtype = torch.float32,

tint: Tuple[float, float, float] = (0.92, 0.92, 1.0),

brightness: Tuple[float, float] = (0.04, 0.40),

mode: str = "PAL") -> List[torch.Tensor]:

"""Generate *n* unique tinted VHS noise tiles.

lo, hi = brightness

tr, tg, tb = tint

tiles: List[torch.Tensor] = []

for _ in range(n):

noise = vhs_noise(width, height, device=device, dtype=dtype, mode=mode)

if mode == "PAL":

luma = lo + (hi - lo) * noise.squeeze(0) # [H, W]

r = luma * tr

g = luma * tg

b = luma * tb

else:

# NTSC: Y plane → brightness-mapped luma, U/V → small chroma offsets.

# Each tile gets its own random color cast — some warm, some cool.

noise_y = noise[0] # [H, W]

noise_u = noise[1] # [H, W], pre-scaled by 0.5

noise_v = noise[2]

y = lo + (hi - lo) * noise_y # [H, W]

u = 0.5 * noise_u # centered, -0.25 ... +0.25

v = 0.5 * noise_v

yuv = torch.stack([y, u, v], dim=0) # [3, H, W]

rgb = yuv_to_rgb(yuv, clamp=True) # [3, H, W]

r = rgb[0] * tr

g = rgb[1] * tg

b = rgb[2] * tb

a = torch.ones(height, width, device=device, dtype=dtype)

rgba = torch.stack([r, g, b, a], dim=0).clamp(0.0, 1.0) # [4, H, W]

tiles.append(rgba)

def __init__(self,

device: torch.device,

dtype: torch.dtype,

nbins: int = 256):

self.device = device

self.dtype = dtype

self.nbins = nbins

self.display_range = np.log(255.0 / 1.0) # max/min brightness, ~45 dB for 8 bits per channel

def _compute_cdf(self,

x: torch.Tensor,

bin_edges: torch.Tensor,

discard_first_bin: bool = False,

discard_last_bin: bool = False):

"""Compute the cumulative distribution function of the data, given `bin_edges`.

level = torch.bucketize(x, bin_edges, out_int32=True) # sum(x > v for v in bin_edges)

count = x.new_zeros(len(bin_edges) + 1, dtype=torch.float32)

count.index_put_((level,), count.new_ones(1), True)

if discard_first_bin:

count = count[1:]

if discard_last_bin:

count = count[:-1]

norm = x.numel()

cdf = torch.cumsum(count, dim=0) / norm

pdf = count / norm

return pdf, cdf

def _loghist(self, x: torch.Tensor, min_nonzero_x, max_x):

# Histogramming linearly in logarithmic space gives us a logarithmically spaced histogram.

log_bin_edges = torch.linspace(torch.log(min_nonzero_x), torch.log(max_x), self.nbins, dtype=self.dtype, device=self.device)

bin_edges = torch.exp(log_bin_edges)

# - Ignore the less-than-first-bin-edge slot to ignore black pixels (empty background). (TODO: should do this for *blank*, not *black*, but that would need some kind of mask support.)

# - Ignore the greater-than-last-bin-edge slot, since we always put the last bin edge at max(x).

pdf, cdf = self._compute_cdf(x, bin_edges, discard_first_bin=True, discard_last_bin=True)

zero = torch.tensor([0.0], device=self.device, dtype=self.dtype)

pdf = torch.cat([zero, pdf], dim=0) # pretend there's no data below the first bin edge

cdf = torch.cat([zero, cdf], dim=0) # a CDF should start from zero

return bin_edges, pdf, cdf

def loghist(self, x: torch.Tensor):

"""Logarithmically spaced histogram.

max_x = torch.max(x)

# min_nonzero_x = torch.min(torch.where(x > 0.0, x, max_x)) # for general data

min_nonzero_x = torch.tensor(0.01, device=self.device, dtype=self.dtype) # for images

return self._loghist(x, min_nonzero_x, max_x)

def larson(self, x: torch.Tensor):

"""Data-adaptive histogram corrector of Ward-Larson, Rushmeier and Piatko (1997).

max_x = torch.max(x)

# min_nonzero_x = torch.min(torch.where(x > 0.0, x, max_x)) # for general data

min_nonzero_x = torch.tensor(0.01, device=self.device, dtype=self.dtype) # for images

bin_edges, pdf, cdf = self._loghist(x, min_nonzero_x, max_x)

data_range = torch.log(max_x / min_nonzero_x).cpu()

if data_range <= self.display_range: # data is LDR, nothing to do

return bin_edges, pdf, cdf

orig_pdf = pdf

orig_cdf = cdf

# Tolerance of the iteration (an arbitrary small number).

#

# Smaller tolerances produce better conformance to the ceiling (i.e. better prevent superlinear growth).

# The original article used the value 0.025.

#

tol = 0.005 * torch.max(pdf)

delta_b = data_range / self.nbins # width of one logarithmic bin

trimmings = torch.tensor(np.inf, device=self.device, dtype=self.dtype) # loop at least once

while trimmings > tol:

pdf_sum = torch.sum(pdf)

if pdf_sum < tol: # degenerate case: histogram has been zeroed out -> return original image

logger.warning("HistogramEqualizer.tonemap: histogram zeroed out, returning original histogram")

return bin_edges, orig_pdf, orig_cdf

# Cut any peaks from PDF that are above the current ceiling; calculate how much we cut in total (to decide whether to iterate again).

ceiling = pdf_sum * (delta_b / self.display_range)

excess = pdf - ceiling

trimmings = torch.sum(torch.where(excess > 0.0, excess, 0.0))

pdf = torch.where(pdf > ceiling, ceiling, pdf)

cdf = torch.cumsum(pdf, dim=0)

cdfmax = torch.max(cdf)

pdf /= cdfmax # normalize sum of adjusted histogram to 1

cdf /= cdfmax # normalize maximum of adjusted cumulative distribution function to 1

return bin_edges, pdf, cdf

def equalize_by_cdf(self,

x: torch.Tensor,

bin_edges: torch.Tensor,

cdf: torch.Tensor):

"""Histogram-equalize data by a given CDF.

discretized = torch.bucketize(x, bin_edges)

"""Gaussian blur kernel size for a given sigma.

def decorator(func):

@wraps(func)

def func_with_metadata(*args, **kwargs):

return func(*args, **kwargs)

func_with_metadata.metadata = metadata # stash it on the function object

return func_with_metadata

"""

def __init__(self,

device: torch.device,

dtype: torch.dtype,

chain: List[Tuple[str, Dict[str, MaybeContained[Atom]]]] = None):

# We intentionally keep very little state in this class, for a more FP/REST approach with less bugs.

# Filters for static effects are stateless.

#

# We deviate from FP in that:

# - The filters MUTATE, i.e. they overwrite the image being processed.

# This is to allow optimizing their implementations for memory usage and speed.

# - The filter for a dynamic effect may store state, if needed for performing FPS correction.

self.device = device

self.dtype = dtype

self.chain = chain

self.histeq = HistogramEqualizer(self.device, self.dtype)

# Meshgrid cache for geometric position of each pixel

self._yy = None

self._xx = None

self._meshy = None

self._meshx = None

self._meshgrid_prev_h = None

self._meshgrid_prev_w = None

# FPS correction

self.CALIBRATION_FPS = 25 # design FPS for dynamic effects (for automatic FPS correction)

self.stream_start_timestamp = time.monotonic_ns() # for updating frame counter reliably (no accumulation)

self.frame_no = -1 # float, frame counter for *normalized* frame number *at CALIBRATION_FPS*

self.last_frame_no = -1

# Caches for individual dynamic effects

self.zoom_data = defaultdict(lambda: None)

self.ca_grid_cache = defaultdict(lambda: None) # name -> {"scale": float, "grid_R": Tensor, "grid_B": Tensor}

self.noise_last_image = defaultdict(lambda: None)

self.noise_last_strength = defaultdict(lambda: -1.0)

self.vhs_glitch_interval = defaultdict(lambda: 0.0)

self.vhs_glitch_last_frame_no = defaultdict(lambda: 0.0)

self.vhs_glitch_last_image = defaultdict(lambda: None)

self.vhs_glitch_last_mask = defaultdict(lambda: None)

self.vhs_headswitching_noise = defaultdict(lambda: None)

self.vhs_tracking_noise = defaultdict(lambda: None)

self.digital_glitches_interval = defaultdict(lambda: 0.0)

self.digital_glitches_last_frame_no = defaultdict(lambda: 0.0)

self.digital_glitches_grid = defaultdict(lambda: None)

def _setup_meshgrid(self, h: int, w: int) -> None:

"""Compute base meshgrid for the geometric position of each pixel.

logger.info(f"_setup_meshgrid: Computing pixel position tensors for image size {w}x{h}")

with torch.inference_mode():

self._yy = torch.linspace(-1.0, 1.0, h, dtype=self.dtype, device=self.device)

self._xx = torch.linspace(-1.0, 1.0, w, dtype=self.dtype, device=self.device)

self._meshy, self._meshx = torch.meshgrid((self._yy, self._xx), indexing="ij")

self._meshgrid_prev_h = h

self._meshgrid_prev_w = w

logger.info("_setup_meshgrid: Pixel position tensors cached")

@classmethod

@memoize

def get_filters(cls):

"""Return a list of available postprocessing filters and their default configurations.

filters = []

for name in dir(cls):

if name.startswith("_"):

continue

if name in ("get_filters", "render_into"):

continue

# The authoritative source for parameter defaults is the source code, so:

meth = getattr(cls, name) # get the method from the class

if not callable(meth): # some other kind of attribute? (e.g. cache data for filters)

continue

# is a method

func, _ = getfunc(meth) # get the underlying raw function, for use with `inspect.signature`

sig = inspect.signature(func)

# All of our filter settings have a default value.

settings = {v.name: v.default for v in sig.parameters.values() if v.default is not inspect.Parameter.empty}

# Parameter ranges are specified at definition site via our internal `with_metadata` mechanism.

ranges = {name: meth.metadata[name] for name in settings}

param_info = {"defaults": settings,

"ranges": ranges}

filters.append((name, param_info))

def rendering_priority(metadata_record):

name, _ = metadata_record

meth = getattr(cls, name)

return meth.metadata["_priority"]

return list(sorted(filters, key=rendering_priority))

def render_into(self, image):

"""Apply current postprocess chain, modifying `image` in-place."""

time_render_start = time.monotonic_ns()

chain = self.chain # read just once; other threads might reassign it while we're rendering

if not chain:

return

c, h, w = image.shape

if h != self._meshgrid_prev_h or w != self._meshgrid_prev_w:

self._setup_meshgrid(h, w)

# Update the frame counter.

#

# We consider the frame number to be a float, so that dynamic filters can decide what

# to do at fractional frame positions. For continuously animated effects (e.g. banding)

# it makes sense to interpolate continuously, whereas other effects (e.g. scanlines)

# can make their decisions based on the integer part.

#

# As always with floats, we must be careful. Note that we operate in a mindset of robust

# engineering. Since doing the Right Thing here does not cost significantly more engineering

# effort than doing the intuitive but Wrong Thing, it is preferable to go for the proper solution,

# regardless of whether it would take a centuries-long session to actually trigger a failure

# in the less robust approach.

#

# So, floating point accuracy considerations? First, we note that accumulation invites

# disaster in two ways:

#

# - Accumulating the result accumulates also representation error and roundoff error.

# - When accumulating small positive numbers to a sum total, the update eventually

# becomes too small to add, causing the counter to get stuck. (For floats, `x + ϵ = x`

# for sufficiently small ϵ dependent on the magnitude of `x`.)

#

# Fortunately, frame number is a linear function of time, and time diffs can be measured

# precisely. Thus, we can freshly compute the current frame number at each frame, completely

# bypassing the need for accumulation:

#

seconds_since_stream_start = (time_render_start - self.stream_start_timestamp) / 10**9

self.last_frame_no = self.frame_no

self.frame_no = self.CALIBRATION_FPS * seconds_since_stream_start # float!

# That leaves just the questions of how accurate the calculation is, and for how long.

# As to the first question:

#

# - Timestamps are an integer number of nanoseconds, so they are exact.

# - Dividing by 10**9, we move the decimal point. But floats are base-2, so 0.1

# is not representable in IEEE-754. So there will be some small representation error,

# which for float64 likely appears in the ~15th significant digit.

# - Basic arithmetic, such as multiplication, is guaranteed by IEEE-754

# to be accurate to the ULP.

#

# Thus, as the result, we obtain the closest number that is representable in IEEE-754,

# and the strategy works for the whole range of float64.

#

# As for the second question, floats are logarithmically spaced. So if this is left running

# "for long enough" during the same session, accuracy will eventually suffer. Instead of the

# counter getting stuck, however, this will manifest as the frame number updating by more

# than `1.0` each time it updates (i.e. whenever the elapsed number of frames reaches the

# next representable float).

#

# This could be fixed by resetting `stream_start_timestamp` once the frame number

# becomes too large. But in practice, how long does it take for this issue to occur?

# The ULP becomes 1.0 at ~5e15. To reach frame number 5e15, at the reference 25 FPS,

# the time required is 2e14 seconds, i.e. 2.31e9 days, or 6.34 million years.

# While I can almost imagine the eventual bug report, I think it's safe to ignore this.

# Apply the current filter chain.

with torch.inference_mode():

for filter_name, settings in chain:

apply_filter = getattr(self, filter_name)

apply_filter(image, **settings)

# --------------------------------------------------------------------------------

# Physical input signal

@with_metadata(center_x=[-1.0, 1.0],

def zoom(self, image: torch.tensor, *,

center_x: float = 0.0,

center_y: float = -0.867,

factor: float = 2.0,

quality: str = "low",

name: str = "zoom0"):

"""[dynamic] Simulated optical zoom for anime video.

if factor == 1.0:

return

# Recompute mesh when the filter settings change, or the video stream size changes.

do_setup = True

c, h, w = image.shape

cached_grid = self.zoom_data[name]["grid"] if name in self.zoom_data else None

size_changed = cached_grid is None or cached_grid.shape[-3] != h or cached_grid.shape[-2] != w

if not size_changed and name in self.zoom_data:

if (factor == self.zoom_data[name]["factor"] and

center_x == self.zoom_data[name]["center_x"] and

center_y == self.zoom_data[name]["center_y"] and

quality == self.zoom_data[name]["quality"]):

do_setup = False

grid = self.zoom_data[name]["grid"]

upscaler = self.zoom_data[name]["upscaler"]

if do_setup:

meshx = center_x + (self._meshx - center_x) / factor # x coordinate, [h, w]

meshy = center_y + (self._meshy - center_y) / factor # y coordinate, [h, w]

grid = torch.stack((meshx, meshy), 2) # [h, w, x/y]

grid = grid.unsqueeze(0) # batch of one

if quality in ("high", "ultra"): # need an Anime4K?

old_upscaler = self.zoom_data[name]["upscaler"] if name in self.zoom_data else None

old_quality = self.zoom_data[name]["quality"] if name in self.zoom_data else None

if size_changed or old_upscaler is None or quality != old_quality:

upscaler_quality = "high" if quality == "ultra" else "low"

upscaler = Upscaler(device=self.device, dtype=self.dtype,

upscaled_width=w, upscaled_height=h,

preset="C", quality=upscaler_quality)

else:

upscaler = old_upscaler # only factor/center changed - save some compute by recycling the existing upscaler.

else:

upscaler = None

self.zoom_data[name] = {"center_x": center_x,

"center_y": center_y,

"factor": factor,

"quality": quality,

"grid": grid,

"upscaler": upscaler}

if quality == "low": # geometric distortion

image_batch = image.unsqueeze(0) # batch of one -> [1, c, h, w]

warped = torch.nn.functional.grid_sample(image_batch, grid, mode="bilinear", padding_mode="border", align_corners=False)

warped = warped.squeeze(0) # [1, c, h, w] -> [c, h, w]

image[:, :, :] = warped

else: # "high" or "ultra" - crop, then Anime4K upscale

g = grid.squeeze(0)

top_left_xy = g[0, 0]

bottom_right_xy = g[-1, -1]

x0 = int((top_left_xy[0] + 1.0) / 2 * w)

y0 = int((top_left_xy[1] + 1.0) / 2 * h)

x1 = int((bottom_right_xy[0] + 1.0) / 2 * w)

y1 = int((bottom_right_xy[1] + 1.0) / 2 * h)

# print(x0, y0, x1, y1) # DEBUG

cropped = torchvision.transforms.functional.crop(image, top=y0, left=x0, height=(y1 - y0), width=(x1 - x0))

image[:, :, :] = upscaler.upscale(cropped)

# Defaults chosen so that they look good for a handful of characters rendered in SD Forge, with the Wai 14.0 Illustrious-SDXL checkpoint.

@with_metadata(threshold=[0.0, 1.0],

def bloom(self, image: torch.tensor, *,

threshold: float = 0.560,

exposure: float = 0.842,

sigma: float = 7.0) -> None:

"""[static] Bloom effect (fake HDR). Makes the image look brighter. Popular in early 2000s anime.

# There are online tutorials for how to create this effect, see e.g.:

# https://learnopengl.com/Advanced-Lighting/Bloom

if threshold < 1.0:

# Find the bright parts.

# original_yuv = rgb_to_yuv(image[:3, :, :])

# Y = original_yuv[0, :, :]

Y = luminance(image[:3, :, :])

mask = torch.ge(Y, threshold) # [h, w]

# Make a copy of the image with just the bright parts.

mask = torch.unsqueeze(mask, 0) # -> [1, h, w]

brights = image * mask # [c, h, w]

# Blur the bright parts. Two-pass blur to save compute, since we need a very large blur kernel.

# It seems that in Torch, one large 1D blur is faster than looping with a smaller one.

#

# Although everything else in Torch takes (height, width), kernel size is given as (size_x, size_y);

# see `gaussian_blur_image` in https://pytorch.org/vision/main/_modules/torchvision/transforms/v2/functional/_misc.html

# for a hint (the part where it computes the padding).

ks = _blur_kernel_size(sigma)

brights = torchvision.transforms.GaussianBlur((ks, 1), sigma=sigma)(brights) # blur along x

brights = torchvision.transforms.GaussianBlur((1, ks), sigma=sigma)(brights) # blur along y

# Additively blend the images. Note we are working in linear intensity space, and we will now go over 1.0 intensity.

image.add_(brights)

# We now have a fake HDR image. Tonemap it back to LDR.

image[:3, :, :] = 1.0 - torch.exp(-image[:3, :, :] * exposure) # RGB: tonemap

# # TEST - apply Larson's adaptive histogram remapper to the Y (luminance) channel, and keep the color channels (U, V) as-is.

# # Doesn't look that good, the classical method seems better for this use.

# hdr_Y = luminance(image[:3, :, :]) # Eris have mercy on us, this function isn't designed for HDR data. Seems to work fine, though.

# bin_edges, pdf, cdf = self.histeq.loghist(hdr_Y) # or use `self.histeq.loghist` for classical histogram equalization

# ldr_Y = self.histeq.equalize_by_cdf(hdr_Y, bin_edges, cdf)

# image[:3, :, :] = yuv_to_rgb(torch.cat([ldr_Y.unsqueeze(0), original_yuv[1:]], dim=0))

if threshold < 1.0:

image[3, :, :] = torch.maximum(image[3, :, :], brights[3, :, :]) # alpha: max-combine

torch.clamp_(image, min=0.0, max=1.0)

# --------------------------------------------------------------------------------

# Video camera

@with_metadata(scale=[0.001, 0.05],

def chromatic_aberration(self, image: torch.tensor, *,

scale: float = 0.005,

sigma: float = 1.0,

name: str = "chromatic_aberration0") -> None:

"""[static] Simulate the two types of chromatic aberration in a camera lens.

# Axial CA: Shrink R (deflected less), pass G through (lens reference wavelength), enlarge B (deflected more).

# Cache grids — they depend only on `scale` + meshgrid (meshgrid invalidates on resolution change).

c, h, w = image.shape

cached = self.ca_grid_cache[name]

size_changed = cached is not None and (cached["grid_R"].shape[-3] != h or cached["grid_R"].shape[-2] != w)

if cached is None or size_changed or scale != cached["scale"]:

grid_R = torch.stack((self._meshx * (1.0 + scale), self._meshy * (1.0 + scale)), 2).unsqueeze(0)

grid_B = torch.stack((self._meshx * (1.0 - scale), self._meshy * (1.0 - scale)), 2).unsqueeze(0)

self.ca_grid_cache[name] = {"scale": scale, "grid_R": grid_R, "grid_B": grid_B}

else:

grid_R, grid_B = cached["grid_R"], cached["grid_B"]

# Batch R+A and B+A to halve grid_sample and GaussianBlur calls.

# R and A-via-grid_R use the same warp grid; same for B and A-via-grid_B.

# GaussianBlur processes all channels, so blurring a 2-channel tensor does both at once.

ks = _blur_kernel_size(sigma)

blur_x = torchvision.transforms.GaussianBlur((ks, 1), sigma=sigma)

blur_y = torchvision.transforms.GaussianBlur((1, ks), sigma=sigma)

# Warp R and A together with grid_R (advanced indexing → copy, safe before writes)

warped_RA = torch.nn.functional.grid_sample(image[[0, 3], :, :].unsqueeze(0),

grid_R, mode="bilinear", padding_mode="border", align_corners=False)

warped_RA = blur_y(blur_x(warped_RA)) # transverse CA

# Warp B and A together with grid_B

warped_BA = torch.nn.functional.grid_sample(image[[2, 3], :, :].unsqueeze(0),

grid_B, mode="bilinear", padding_mode="border", align_corners=False)

warped_BA = blur_y(blur_x(warped_BA)) # transverse CA

image[0, :, :] = warped_RA[0, 0, :, :]

# image[1, :, :] passed through as-is (G is the lens reference wavelength)

image[2, :, :] = warped_BA[0, 0, :, :]

# Alpha: average the R-warped, G-original, and B-warped alpha values (in-place, no temporary)

image[3, :, :].add_(warped_RA[0, 1, :, :]).add_(warped_BA[0, 1, :, :]).mul_(1.0 / 3.0)

@with_metadata(strength=[0.1, 0.5],

def vignetting(self, image: torch.tensor, *,

strength: float = 0.42) -> None:

"""[static] Simulate vignetting (less light hitting the corners of a film frame or CCD sensor).

euclidean_distance_from_center = (self._meshy**2 + self._meshx**2)**0.5 / 2**0.5 # [h, w]

brightness = torch.cos(strength * euclidean_distance_from_center * math.pi)**2 # [h, w]

brightness = torch.unsqueeze(brightness, 0) # -> [1, h, w]

image[:3, :, :] *= brightness

# --------------------------------------------------------------------------------

# Retouching / color grading

def _rgb_to_hue(self, rgb: List[float]) -> float:

"""Convert an RGB color to an HSL hue, for use as `bandpass_hue` in `_desaturate_impl`.

R, G, B = rgb

alpha = 0.5 * (2.0 * R - G - B)

beta = 3.0**0.5 / 2.0 * (G - B)

hue = math.atan2(beta, alpha) / (2.0 * math.pi) # note atan2(0, 0) := 0

hue = hue + 0.5 # convert from `[-0.5, 0.5)` to `[0, 1)`

return hue

def _desaturate_impl(self, image: torch.tensor, *,

strength: float = 1.0,

tint_rgb: List[float] = [1.0, 1.0, 1.0],

bandpass_reference_rgb: List[float] = [1.0, 0.0, 0.0],

bandpass_q: float = 0.0) -> None:

"""Shared implementation for `desaturate` and `monochrome_display`."""

R = image[0, :, :]

G = image[1, :, :]

B = image[2, :, :]

if bandpass_q > 0.0: # hue bandpass enabled?

# Calculate hue of each pixel, using a cartesian-to-polar approximation of the HSL representation.

# An approximation is fine here, because we only use this for a hue detector.

# This is faster and requires less branching than the exact hexagonal representation.

desat_alpha = 0.5 * (2.0 * R - G - B)

desat_beta = 3.0**0.5 / 2.0 * (G - B)

desat_hue = torch.atan2(desat_beta, desat_alpha) / (2.0 * math.pi) # note atan2(0, 0) := 0

desat_hue += 0.5 # convert from `[-0.5, 0.5)` to `[0, 1)`

# -> [h, w]

# Determine whether to keep this pixel or desaturate (and by how much).

#

# Calculate distance of each pixel from reference hue, accounting for wrap-around.

bandpass_hue = self._rgb_to_hue(bandpass_reference_rgb)

desat_temp1 = torch.abs(desat_hue - bandpass_hue)

desat_temp2 = torch.abs((desat_hue + 1.0) - bandpass_hue)

desat_temp3 = torch.abs(desat_hue - (bandpass_hue + 1.0))

desat_hue_distance = 2.0 * torch.minimum(torch.minimum(desat_temp1, desat_temp2),

desat_temp3) # [0, 0.5] -> [0, 1]

# -> [h, w]

# How to interpret the following factor:

# 0: far away from reference hue, pixel should be desaturated

# 1: at reference hue, pixel should be kept as-is

#

# - Pixels with their hue at least `bandpass_q` away from `bandpass_hue` are fully desaturated.

# - As distance falls below `bandpass_q`, a blend starts very gradually.

# - As the hue difference approaches zero, the pixel is fully passed through.

# - The 1.0 - ... together with the square makes a sharp spike at the reference hue.

desat_diff2 = (1.0 - torch.clamp(desat_hue_distance / bandpass_q, min=0.0, max=1.0))**2

# Gray pixels should always be desaturated, so that they get the tint applied.

# In HSL computations, gray pixels have an arbitrary hue, usually red, so we must filter them out separately.

# We can do this in YUV space.

YUV = rgb_to_yuv(image[:3, :, :])

Y = YUV[0, :, :] # -> [h, w]

U = YUV[1, :, :] # -> [h, w]

V = YUV[2, :, :] # -> [h, w]

notgray_threshold = 0.05

urel = torch.clamp(torch.abs(U) / notgray_threshold, min=0.0, max=1.0)

vrel = torch.clamp(torch.abs(V) / notgray_threshold, min=0.0, max=1.0)

notgray = (urel**2 + vrel**2)**0.5 # [h, w]; 0: completely gray; 1: has at least some color

desat_diff2 *= notgray

strength_field = strength * (1.0 - desat_diff2) # [h, w]; "field" as in physics, NOT as in CRT TV

else:

Y = luminance(image[:3, :, :]) # save some compute since in this case we don't need U and V

strength_field = strength # just a scalar!

# Desaturate, then apply tint

Y = Y.unsqueeze(0) # -> [1, h, w]

tint_color = torch.tensor(tint_rgb, device=self.device, dtype=image.dtype).unsqueeze(1).unsqueeze(2) # [c, 1, 1]

tinted_desat_image = Y * tint_color # -> [c, h, w]

# Final blend

image[:3, :, :] = (1.0 - strength_field) * image[:3, :, :] + strength_field * tinted_desat_image

# This filter is adapted from an old GLSL code I made for Panda3D 1.8 back in 2014.

@with_metadata(strength=[0.0, 1.0],

def desaturate(self, image: torch.tensor, *,

strength: float = 1.0,

tint_rgb: List[float] = [1.0, 1.0, 1.0],

bandpass_reference_rgb: List[float] = [1.0, 0.0, 0.0],

bandpass_q: float = 0.0) -> None:

"""[static] Desaturate the image, with optional hue bandpass and tint.

self._desaturate_impl(image, strength=strength, tint_rgb=tint_rgb,

bandpass_reference_rgb=bandpass_reference_rgb, bandpass_q=bandpass_q)

# --------------------------------------------------------------------------------

# General use

def _noise_impl(self, image: torch.tensor, *,

strength: float = 0.3,

sigma: float = 1.0,

channel: str = "Y",

double_size: bool = True,

ntsc_chroma: str = "nearest",

name: str = "noise0") -> None:

"""Shared implementation for `noise` and `analog_vhs_noise`."""

# Re-randomize the noise texture whenever the normalized frame number changes, or the video stream size changes.

c, h, w = image.shape

cached = self.noise_last_image[name]

size_changed = cached is not None and (cached.shape[-2] != h or cached.shape[-1] != w)

strength_changed = (strength != self.noise_last_strength[name])

if cached is None or size_changed or strength_changed or int(self.frame_no) > int(self.last_frame_no):

c, h, w = image.shape

if channel.startswith("VHS_"):

vhs_mode = channel.removeprefix("VHS_") # "PAL" or "NTSC"

noise_image = vhs_noise(w, h, device=self.device, dtype=image.dtype,

mode=vhs_mode, ntsc_chroma=ntsc_chroma, double_size=double_size)

# PAL: [1, H, W]; NTSC: [3, H, W] — stored as-is, distinguished at application time.

# NTSC chroma is already 4:2:0 subsampled inside vhs_noise.

else:

noise_image = isotropic_noise(w, h, device=self.device, dtype=image.dtype,

sigma=sigma, double_size=double_size)

noise_image.mul_(strength) # bake in current strength to save a few full-image multiplications during rendering

self.noise_last_image[name] = noise_image

self.noise_last_strength[name] = strength

else:

noise_image = cached

base_multiplier = 1.0 - strength

if channel == "A": # alpha

image[3, :, :].mul_(base_multiplier + noise_image)

return

image_yuv = rgb_to_yuv(image[:3, :, :])

if channel == "VHS_PAL":

image_yuv[0, :, :].mul_(base_multiplier + noise_image.squeeze(0))

elif channel == "VHS_NTSC":

image_yuv = rgb_to_yuv(image[:3, :, :])

# Luma: multiplicative (same as PAL/Y modes).

image_yuv[0, :, :].mul_(base_multiplier + noise_image[0])