Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
124 lines
4.8 KiB
Python
124 lines
4.8 KiB
Python
"""
|
|
The skeleton the scalar-field generators share.
|
|
|
|
domain_warp draws a field, standardises it, masks it, hands channel 0 to
|
|
fill_channels and draws the rest through the same closures; five golden fixtures
|
|
pin that generator to the byte, so its fbm_noise stays where it is. Everything
|
|
here is for the generators added after it, which have no fixtures to keep: the
|
|
same skeleton, and an FBM whose 3D path uses the four-corner simplex rather than
|
|
the corner-0 variant those fixtures froze.
|
|
|
|
`seed` may be an int or an int64 tensor of seeds, one per channel. The
|
|
coordinates then grow a leading axis at the first hash and everything after it
|
|
broadcasts, which is how fill_channels draws the whole channel axis in one call;
|
|
see shaders/simplex.py.
|
|
"""
|
|
import torch
|
|
|
|
from .base import BaseNoiseGenerator
|
|
from .simplex import simplex_2d, simplex_3d_full
|
|
from ..utils.shape_masks import apply_shape_mask
|
|
|
|
MAX_OCTAVES = 8
|
|
|
|
|
|
def time_is_an_axis(time_axis, time):
|
|
"""
|
|
Whether to evaluate the field in 3D with time as the third axis.
|
|
|
|
core.shader_noise sets `time_axis` once per clip from the frame count. A
|
|
direct caller that leaves it unset gets the older test, whether time is
|
|
nonzero, which is what domain_warp's single-image fixtures pin.
|
|
"""
|
|
return (time != 0) if time_axis is None else time_axis
|
|
|
|
|
|
def with_time(p, z):
|
|
"""`p[..., :2]` with a constant third coordinate."""
|
|
return torch.cat([p[..., :2], torch.full_like(p[..., :1], float(z))], dim=-1)
|
|
|
|
|
|
def simplex_warp(p, seed, strength, frequency):
|
|
"""Displace `p` by two low-frequency simplex fields, at most `strength` apart."""
|
|
if strength <= 0.0:
|
|
return p
|
|
q = p * frequency
|
|
dx = simplex_2d(q, seed + 100)
|
|
dy = simplex_2d(q, seed + 200)
|
|
return torch.cat([p[..., 0:1] + dx * strength, p[..., 1:2] + dy * strength], dim=-1)
|
|
|
|
|
|
def fbm(p, octaves, seed, persistence=0.5, lacunarity=2.0, time=None, octave_offset=None):
|
|
"""
|
|
Layered simplex noise over `p[..., :2]`.
|
|
|
|
With `time` the layers are evaluated in 3D, each a little further along the
|
|
time axis than the last so they do not move in lockstep. `octave_offset` is a
|
|
translation of `[dx, dy]` per layer index, so the layers are drawn from
|
|
different parts of the same field.
|
|
"""
|
|
total = None
|
|
amp, freq, norm = 1.0, 1.0, 0.0
|
|
for i in range(min(max(int(octaves), 1), MAX_OCTAVES)):
|
|
q = p[..., :2] * freq
|
|
if octave_offset is not None:
|
|
q = q + octave_offset * i
|
|
if time is None:
|
|
layer = simplex_2d(q, seed + i)
|
|
else:
|
|
layer = simplex_3d_full(with_time(q, time * (0.2 + 0.05 * i)), seed + i)
|
|
total = amp * layer if total is None else total + amp * layer
|
|
norm += amp
|
|
amp *= persistence
|
|
freq *= lacunarity
|
|
return total / norm
|
|
|
|
|
|
def standardise(field, seed):
|
|
"""
|
|
Zero mean and unit deviation: over the whole draw for one seed, per leading
|
|
slice for a tensor of seeds, so a batched channel matches the draw it replaces.
|
|
"""
|
|
if torch.is_tensor(seed):
|
|
dims = tuple(range(1, field.dim()))
|
|
return (field - field.mean(dim=dims, keepdim=True)) / (field.std(dim=dims, keepdim=True) + 1e-8)
|
|
return (field - field.mean()) / (field.std() + 1e-8)
|
|
|
|
|
|
def draw_channels(field, coords, params, seed, target_channels, contrast=1.0, palette=True):
|
|
"""
|
|
Fill the channel axis from `field(seed) -> [B, H, W, 1]` (or `[N, B, H, W, 1]`
|
|
for a tensor of seeds), finishing each draw the way every generator does:
|
|
standardised, halved, masked, clamped.
|
|
|
|
The shape mask is drawn once and lerped into every channel: it is a spatial
|
|
shape, and apply_shape_mask reseeds the global RNG. Channel 0 stays on the
|
|
scalar path, which is what keeps a one-channel draw and the travel-mode bases
|
|
built from it unchanged; the palette, if any, maps that one field to three.
|
|
"""
|
|
shape_type = params.shape_type
|
|
shape_strength = params.shape_strength
|
|
mask = None
|
|
if shape_type not in ["none", "0"] and shape_strength > 0:
|
|
mask = apply_shape_mask(coords, shape_type, params.time, seed, shape_strength)
|
|
|
|
def finish(raw, channel_seed):
|
|
out = standardise(raw, channel_seed) * (0.5 * contrast)
|
|
if mask is not None:
|
|
out = torch.lerp(out, out * mask, shape_strength)
|
|
return torch.clamp(out, -1.0, 1.0)
|
|
|
|
base = finish(field(seed), seed).permute(0, 3, 1, 2)
|
|
if palette:
|
|
base = BaseNoiseGenerator.palette_channels(base, params)
|
|
|
|
def draw(channel_seed):
|
|
return finish(field(channel_seed), channel_seed).permute(0, 3, 1, 2)
|
|
|
|
def draw_many(channel_seeds):
|
|
seeds = channel_seeds.reshape(-1, 1, 1, 1, 1)
|
|
# [N, B, H, W, 1] -> [B, N, H, W]
|
|
return finish(field(seeds), seeds).squeeze(-1).permute(1, 0, 2, 3)
|
|
|
|
return BaseNoiseGenerator.fill_channels(draw, base, target_channels, seed, render_many=draw_many)
|