51 lines
2.6 KiB
Python
51 lines
2.6 KiB
Python
"""Float Lanczos-3, antialiased and separable, without PIL/uint8 round trips."""
|
|
import math
|
|
import torch
|
|
|
|
# Target block size for temporary gathered samples. A single column can exceed
|
|
# this for extremely large inputs; no image or coefficient cache survives a run.
|
|
_GATHER_ELEMENTS = 4 * 1024 * 1024
|
|
|
|
|
|
def _axis(samples, length, axis):
|
|
source_length = samples.shape[axis]
|
|
if source_length == length:
|
|
return samples
|
|
step = source_length / length
|
|
filter_scale = max(1.0, step) # Widen the filter when reducing resolution.
|
|
radius = 3 * filter_scale
|
|
taps = math.ceil(2 * radius) + 1
|
|
positions = (torch.arange(length, device=samples.device, dtype=samples.dtype) + .5) * step - .5
|
|
first = torch.floor(positions - radius).to(dtype=torch.int64)
|
|
indices = first[:, None] + torch.arange(taps, device=samples.device)[None, :]
|
|
distance = (positions[:, None] - indices.to(dtype=samples.dtype)) / filter_scale
|
|
weights = torch.sinc(distance) * torch.sinc(distance / 3)
|
|
valid = (distance.abs() < 3) & (indices >= 0) & (indices < source_length)
|
|
weights = torch.where(valid, weights, 0)
|
|
weights = weights / weights.sum(-1, keepdim=True)
|
|
indices = indices.clamp(0, source_length - 1)
|
|
moved = samples.movedim(axis, -1)
|
|
output = torch.empty(*moved.shape[:-1], length, device=samples.device, dtype=samples.dtype)
|
|
leading = math.prod(moved.shape[:-1])
|
|
chunk = max(1, min(length, _GATHER_ELEMENTS // max(1, leading * taps)))
|
|
for start in range(0, length, chunk):
|
|
stop = min(length, start + chunk)
|
|
selected = moved.index_select(-1, indices[start:stop].reshape(-1))
|
|
selected = selected.reshape(*moved.shape[:-1], stop - start, taps)
|
|
output[..., start:stop] = (selected * weights[start:stop]).sum(-1)
|
|
return output.movedim(-1, axis)
|
|
|
|
|
|
@torch.no_grad()
|
|
def lanczos_resize(samples, width, height):
|
|
"""Resize NCHW float32/float64 tensors on their current device."""
|
|
if not isinstance(samples, torch.Tensor) or samples.ndim != 4 or samples.dtype not in (torch.float32, torch.float64):
|
|
raise ValueError("Lanczos: expected an NCHW float32/float64 working tensor.")
|
|
if min(samples.shape) < 1 or isinstance(width, bool) or isinstance(height, bool) or \
|
|
not isinstance(width, int) or not isinstance(height, int) or min(width, height) < 1:
|
|
raise ValueError("Lanczos: source and output dimensions must be positive integers.")
|
|
# Reduce intermediate canvas size when one dimension is being reduced.
|
|
if height / samples.shape[-2] < width / samples.shape[-1]:
|
|
return _axis(_axis(samples, height, -2), width, -1)
|
|
return _axis(_axis(samples, width, -1), height, -2)
|