diff --git a/__init__.py b/__init__.py index d97ba83..5d760b1 100644 --- a/__init__.py +++ b/__init__.py @@ -25,8 +25,8 @@ from .shaders.registry import ( ) # Importing the generator modules runs their @shader_generator decorators, -# which register "domain_warp", "tensor_field", "curl_noise" and -# "temporal_coherent". Only the aliases below are registered here. +# which is what registers every shader type. Only the aliases below are +# registered here. from .shaders.domain_warp import ( DomainWarpGenerator, generate_domain_warp_tensor, @@ -47,6 +47,10 @@ from .shaders.spectral import ( SpectralNoiseGenerator, generate_spectral_tensor, ) +from .shaders.gaussian import ( + GaussianNoiseGenerator, + generate_gaussian_tensor, +) register_shader("curl", CurlNoiseGenerator, { "description": "Curl/fluid noise patterns (alias)", diff --git a/direct_shader_ksampler.py b/direct_shader_ksampler.py index 3829e31..3652a2c 100644 --- a/direct_shader_ksampler.py +++ b/direct_shader_ksampler.py @@ -28,7 +28,7 @@ class DirectShaderNoiseKSampler(ShaderNoiseKSampler): "use_temporal_coherence": ("BOOLEAN", {"default": False, "tooltip": "Hold one seed across every video frame so the shader pattern evolves only through time, instead of redrawing per frame. Frames share one pattern, so it reinforces rather than averaging out and shows at lower strength than a redrawn pattern: on MiniMax H3 it took over the picture at 0.5 and was hard to see at 0.2. No effect on single images."}), # New direct shader parameters - "shader_type": (["domain_warp", "tensor_field", "curl_noise", "temporal_coherent", "spectral"], {"default": "domain_warp", "tooltip": "Which noise pattern to blend in; each has its own character when it shows. domain_warp: flowing, intricate distortions, the even-handed default. tensor_field: structured and directional. curl_noise: smooth fluid motion, and the one that shows soonest, at about half the strength of the others. temporal_coherent: 4D simplex with time as a real axis, built for smooth animation; on MiniMax H3 it shows sooner than domain_warp, as a dot-grid pattern past about 0.5. spectral: soft cloud-like fields built from their own frequency band rather than pixel by pixel -- no filaments or swirls, and colour schemes do nothing to it, but noise_scale sets the band directly and it costs almost nothing to draw (about 90ms where the others take seconds on a long video). Its character is untested against real prompts, so treat its strengths as unknown rather than calibrated. The live preview only draws the first three; picking another leaves the preview on its last pattern, which does not affect sampling."}), + "shader_type": (["domain_warp", "tensor_field", "curl_noise", "temporal_coherent", "spectral", "gaussian"], {"default": "domain_warp", "tooltip": "Which noise pattern to blend in; each has its own character when it shows. domain_warp: flowing, intricate distortions, the even-handed default. tensor_field: structured and directional. curl_noise: smooth fluid motion, and the one that shows soonest, at about half the strength of the others. temporal_coherent: 4D simplex with time as a real axis, built for smooth animation; on MiniMax H3 it shows sooner than domain_warp, as a dot-grid pattern past about 0.5. spectral: soft cloud-like fields built from their own frequency band rather than pixel by pixel -- no filaments or swirls, and colour schemes do nothing to it, but noise_scale sets the band directly and it costs almost nothing to draw (about 90ms where the others take seconds on a long video). Its character is untested against real prompts, so treat its strengths as unknown rather than calibrated. gaussian: the control -- plain white noise, the same thing the sampler already starts from, so any strength moves you toward another seed's neighbourhood without adding structure; noise_scale, octaves, warp_strength, phase_shift and colour schemes do nothing to it, shape masks still apply, and with temporal coherence a clip drifts smoothly from one white field to a second. The live preview only draws the first three; picking another leaves the preview on its last pattern, which does not affect sampling."}), "shape_type": (["none", "radial", "linear", "spiral", "checkerboard", "spots", "hexgrid", "stripes", "gradient", "vignette", "cross", "stars", "triangles", "concentric", "rays", "zigzag"], {"default": "none", "tooltip": "Mask the shader noise into a shape before it reaches the sampler (not post-processing). A mask concentrates the noise into hard geometry, which survives denoising far more readily than plain shader noise, so the shape itself starts being drawn into the picture at lower strength: on MiniMax H3 it was at 0.6, which is what the stamp preset uses. Around 0.2 the mask shapes the noise without being drawn."}), "color_scheme": (["none", "blue_red", "viridis", "plasma", "inferno", "magma", "turbo", "jet", "rainbow", "cool", "hot", "parula", "hsv", "autumn", "winter", "spring", "summer", "copper", "pink", "bone", "ocean", "terrain", "neon", "fire"], {"default": "none", "tooltip": "Choose a color palette to apply to the shader noise visualization [not post processing - is applied to the shader noise pattern before rendering]"}), "noise_scale": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.001, "tooltip": "Size of the shader's features: lower is larger and zoomed in, higher is smaller and zoomed out. It changes how the shader shows as much as strength does. Large features carry the pattern into the result and pull away from the seed; small ones are mostly absorbed into the picture while still steering it. On SD 1.5 at strength 0.5, 0.5 turned every seed abstract and 2.0 gave clean portraits again; on MiniMax H3 at the same strength both stayed photographic, with 0.5 re-composing the scene most. Small shifts can lead to large variations."}), diff --git a/shaders/gaussian.py b/shaders/gaussian.py new file mode 100644 index 0000000..c1bcf7b --- /dev/null +++ b/shaders/gaussian.py @@ -0,0 +1,85 @@ +""" +White noise: the control. + +Every other type draws a structured field. This one draws exactly what the +sampler already starts from, an independent Gaussian field per channel, so +blending it in moves the run toward another seed's neighbourhood without adding +any structure of its own. That isolates what shader_strength and the blend mode +do by themselves, which is what a control is for. + +Nothing spatial reaches it: noise_scale, octaves, warp_strength, phase_shift and +colour schemes do nothing, and the tooltip says so. Shape masks still apply, +since a mask is a spatial shape imposed on whatever field it is given. + +Time turns one white field into a second along a great circle, so holding the +seed and advancing it evolves the noise smoothly instead of redrawing it: at +time 0 the field is the first draw exactly, at time 1 the second. +""" +import math + +import torch + +from .base import BaseNoiseGenerator +from .registry import shader_generator +from ..core.constants import DEFAULT_CHANNELS +from ..core.params import ShaderParams +from ..utils.noise_utils import create_coordinate_grid +from ..utils.shape_masks import apply_shape_mask + +# Channel seed stride, shared with shaders/base.py so the conventions agree. +_CHANNEL_SEED_STRIDE = 6151 + + +@shader_generator("gaussian", metadata={ + "description": "White noise, an independent Gaussian field per channel: the control", + "supports_temporal": True, +}) +class GaussianNoiseGenerator(BaseNoiseGenerator): + """An independent Gaussian field per channel, with no spatial structure.""" + + @staticmethod + def generate( + batch_size: int, + height: int, + width: int, + params: ShaderParams, + device: torch.device, + seed: int = 0, + target_channels: int = DEFAULT_CHANNELS, + ) -> torch.Tensor: + time = float(params.time) + base_seed = params.get("base_seed", seed) + current_seed = base_seed if params.use_temporal_coherence else seed + + # One CPU generator per channel, seeded the way shaders/base.py seeds its + # extra channels, so channel 0 of a wide draw is the one-channel draw at + # the same seed; shaders/spectral.py records why one wide randn is not. + angle = time * math.pi / 2.0 + fields = [] + for channel in range(target_channels): + channel_generator = torch.Generator(device="cpu").manual_seed( + int(current_seed) + _CHANNEL_SEED_STRIDE * channel) + first = torch.randn(batch_size, 1, height, width, generator=channel_generator) + second = torch.randn(batch_size, 1, height, width, generator=channel_generator) + fields.append(first * math.cos(angle) + second * math.sin(angle)) + field = torch.cat(fields, dim=1).to(device) + mean = field.mean(dim=(-2, -1), keepdim=True) + std = field.std(dim=(-2, -1), keepdim=True) + field = (field - mean) / (std + 1e-8) + + shape_type = params.shape_type + shape_strength = params.shape_strength + if shape_type not in ["none", "0"] and shape_strength > 0: + coords = create_coordinate_grid(batch_size, height, width, device) + mask = apply_shape_mask(coords, shape_type, time, current_seed, shape_strength) + field = torch.lerp(field, field * mask.permute(0, 3, 1, 2), shape_strength) + + # Same convention as the other generators: standardised, halved, clamped. + return torch.clamp(field * 0.5, -1.0, 1.0) + + +def generate_gaussian_tensor(batch_size, height, width, params, device, seed=0, + target_channels=DEFAULT_CHANNELS): + """Function form, matching the other generators' module-level entry points.""" + return GaussianNoiseGenerator.generate( + batch_size, height, width, params, device, seed, target_channels) diff --git a/tests/golden/image_gaussian.pt b/tests/golden/image_gaussian.pt new file mode 100644 index 0000000..925d8e6 Binary files /dev/null and b/tests/golden/image_gaussian.pt differ diff --git a/tests/golden/video_gaussian.pt b/tests/golden/video_gaussian.pt new file mode 100644 index 0000000..0740c82 Binary files /dev/null and b/tests/golden/video_gaussian.pt differ diff --git a/tests/golden_cases.py b/tests/golden_cases.py index 95bd9ed..b542441 100644 --- a/tests/golden_cases.py +++ b/tests/golden_cases.py @@ -54,6 +54,11 @@ CASES = { # the generators added later can share it; this pins its output across the move. "video_temporal_coherent": dict(kind="flow", video=True, shader_type="temporal_coherent", use_temporal_coherence=True, sequential_stages=2), + # One image and one video case per generator added after spectral, the video + # one with temporal coherence, which is where each differs most from a redraw. + "image_gaussian": dict(shader_type="gaussian"), + "video_gaussian": dict(kind="flow", video=True, shader_type="gaussian", + use_temporal_coherence=True, sequential_stages=2), # The spectral generator builds its field from a frequency band rather than # per pixel, so it shares none of the others' code below core.shader_noise. # One image case and one video case, the video one with temporal coherence diff --git a/tests/test_gaussian.py b/tests/test_gaussian.py new file mode 100644 index 0000000..0f7b96f --- /dev/null +++ b/tests/test_gaussian.py @@ -0,0 +1,79 @@ +""" +The gaussian generator: white noise, the control. + +The properties worth pinning are the absences. It must have no spatial +structure, none of the pattern knobs may reach it, and it must still keep the +contract the other generators keep: a field per channel, channel 0 unchanged by +how many were asked for, and one field evolving under temporal coherence. +""" +import pytest +import torch + +from snk.core.constants import CHANNEL_BASIS +from snk.core.shader_noise import effective_channel_rank, generate + +CPU = torch.device("cpu") +PARAMS = { + "scale": 1.0, "octaves": 3.0, "warp_strength": 0.5, "phase_shift": 0.5, + "shape_type": "none", "color_scheme": "none", "time": 0.0, "base_seed": 8888, +} + + +def _neighbour_correlation(noise): + field = noise[0] + pairs = torch.stack([field[:, :-1, :].flatten(), field[:, 1:, :].flatten()]) + return torch.corrcoef(pairs)[0, 1].item() + + +def test_the_field_is_white(): + noise = generate((1, 24, 22, 38), PARAMS, "gaussian", 8888, CPU) + assert abs(_neighbour_correlation(noise)) < 0.05 + assert noise.abs().max() <= 1.0 + + +@pytest.mark.parametrize("knob", [dict(scale=0.3), dict(scale=3.0), dict(octaves=8.0), + dict(warp_strength=5.0), dict(phase_shift=2.0), + dict(color_scheme="viridis")]) +def test_no_pattern_knob_reaches_it(knob): + """The tooltip promises these do nothing, so a stray use of one fails here.""" + stock = generate((1, 4, 16, 16), PARAMS, "gaussian", 8888, CPU) + assert torch.equal(generate((1, 4, 16, 16), dict(PARAMS, **knob), "gaussian", 8888, CPU), stock) + + +def test_a_shape_mask_still_applies(): + stock = generate((1, 4, 16, 16), PARAMS, "gaussian", 8888, CPU) + masked = generate((1, 4, 16, 16), dict(PARAMS, shape_type="radial"), "gaussian", 8888, CPU) + assert not torch.equal(masked, stock) + + +@pytest.mark.parametrize("shape", [(1, 4, 64, 64), (1, 24, 22, 38), (1, 24, 5, 16, 16), + (1, 128, 3, 8, 8), (1, 128, 3, 32, 32)]) +def test_every_channel_is_its_own_field(shape): + rank = effective_channel_rank(generate(shape, PARAMS, "gaussian", 8888, CPU)) + assert rank > min(shape[1], CHANNEL_BASIS) * 0.85 + + +def test_channel_zero_does_not_depend_on_how_many_channels_were_asked_for(): + wide = generate((1, 24, 5, 22, 38), PARAMS, "gaussian", 8888, CPU) + single = generate((1, 1, 5, 22, 38), PARAMS, "gaussian", 8888, CPU) + assert torch.equal(wide[:, :1], single) + + +def test_holding_the_seed_and_advancing_time_turns_one_field_into_another(): + """ + A 5-frame clip steps time by a quarter, so adjacent frames sit cos(pi/8) + apart on the great circle: 0.92. Without coherence each frame is its own seed. + """ + def frame_correlation(coherent): + noise = generate((1, 24, 5, 22, 38), dict(PARAMS, use_temporal_coherence=coherent), + "gaussian", 8888, CPU, temporal_coherence=coherent) + pairs = torch.stack([noise[0, :, 0].flatten(), noise[0, :, 1].flatten()]) + return torch.corrcoef(pairs)[0, 1].item() + + assert frame_correlation(True) > 0.8 + assert abs(frame_correlation(False)) < 0.1 + # Frame 0 is the first draw exactly, whatever time does to the frames after it. + clip = generate((1, 4, 5, 16, 16), dict(PARAMS, use_temporal_coherence=True), + "gaussian", 8888, CPU, temporal_coherence=True) + still = generate((1, 4, 16, 16), dict(PARAMS, use_temporal_coherence=True), "gaussian", 8888, CPU) + assert torch.equal(clip[:, :, 0], still) diff --git a/tests/test_generators_smoke.py b/tests/test_generators_smoke.py index 41db77b..0b2b0f2 100644 --- a/tests/test_generators_smoke.py +++ b/tests/test_generators_smoke.py @@ -12,7 +12,8 @@ import itertools import pytest import torch -SHADER_TYPES = ["domain_warp", "tensor_field", "curl_noise", "temporal_coherent"] +SHADER_TYPES = ["domain_warp", "tensor_field", "curl_noise", "temporal_coherent", "spectral", + "gaussian"] def make_params(octaves=1, channels=4, color_scheme="none", shape_type="none", temporal=False): diff --git a/tests/test_presets.py b/tests/test_presets.py index fdf3445..61f1a14 100644 --- a/tests/test_presets.py +++ b/tests/test_presets.py @@ -33,7 +33,7 @@ def test_the_modes_are_ordered_by_how_wide_they_leave_the_noise(shape): @pytest.mark.parametrize("shader_type", ["domain_warp", "tensor_field", "curl_noise", - "temporal_coherent"]) + "temporal_coherent", "gaussian"]) def test_jump_collapses_every_generator_not_just_the_narrow_ones(shader_type): """ jump has to mean the same thing whatever shader is chosen, otherwise diff --git a/tests/test_shader_noise.py b/tests/test_shader_noise.py index c2041c5..45fbd7f 100644 --- a/tests/test_shader_noise.py +++ b/tests/test_shader_noise.py @@ -103,7 +103,8 @@ def test_temporal_coherence_holds_the_seed(): assert not torch.allclose(without, with_tc) -@pytest.mark.parametrize("shader_type", ["domain_warp", "tensor_field", "curl_noise", "temporal_coherent"]) +@pytest.mark.parametrize("shader_type", ["domain_warp", "tensor_field", "curl_noise", "temporal_coherent", + "gaussian"]) def test_every_registered_generator_is_usable(shader_type): noise = generate((1, 4, 16, 16), {**PARAMS, "shader_type": shader_type}, shader_type, 8888, CPU) assert tuple(noise.shape) == (1, 4, 16, 16) @@ -123,7 +124,8 @@ def test_unknown_shader_type_raises_a_clear_error(): # --- channel width -------------------------------------------------------------------- -ALL_SHADERS = ["domain_warp", "temporal_coherent", "curl_noise", "tensor_field", "spectral"] +ALL_SHADERS = ["domain_warp", "temporal_coherent", "curl_noise", "tensor_field", "spectral", + "gaussian"] @pytest.mark.parametrize("shape", [(1, 4, 48, 48), (1, 16, 32, 32), (1, 24, 5, 16, 16)]) @@ -139,7 +141,7 @@ def test_generators_span_their_channel_axis(shape, shader_type): assert rank > shape[1] * 0.6, f"{shader_type}: rank {rank:.2f} of {shape[1]}" -@pytest.mark.parametrize("shader_type", ["domain_warp", "temporal_coherent", "curl_noise"]) +@pytest.mark.parametrize("shader_type", ["domain_warp", "temporal_coherent", "curl_noise", "gaussian"]) def test_a_latent_wider_than_the_basis_still_spans_the_basis(shader_type): """Past CHANNEL_BASIS the channels are mixtures, which must not re-collapse the draw.""" rank = effective_channel_rank(generate((1, 128, 3, 8, 8), PARAMS, shader_type, 8888, CPU)) @@ -264,7 +266,7 @@ def test_a_collapse_does_not_render_the_draw_it_discards(shape, shader_type, mon assert set(calls) == {1}, f"a collapse rendered a {max(calls)}-channel draw it cannot use" -@pytest.mark.parametrize("shader_type", ["domain_warp", "curl_noise", "temporal_coherent"]) +@pytest.mark.parametrize("shader_type", ["domain_warp", "curl_noise", "temporal_coherent", "gaussian"]) def test_the_hash_generators_do_not_reseed_the_global_rng(shader_type, monkeypatch): """ These three are pure coordinate hashes of their seed argument, so the diff --git a/tests/test_simplex.py b/tests/test_simplex.py index 5d2b4d6..dc18531 100644 --- a/tests/test_simplex.py +++ b/tests/test_simplex.py @@ -103,7 +103,7 @@ def test_the_seed_is_coerced_to_int64(): @pytest.mark.parametrize("shader_type", ["domain_warp", "curl_noise", - "temporal_coherent", "tensor_field"]) + "temporal_coherent", "tensor_field", "gaussian"]) @pytest.mark.parametrize("hw", RAGGED + ALIGNED) def test_channel_zero_survives_batching_at_every_shape(shader_type, hw): """