nodes/AEmotionStudio-ComfyUI-ShaderNoiseKSampler
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
0d7c834d33a223245e563df6615fded27cb2600d
AEmotionStudio-ComfyUI-Shad…/shader_params_reader.py
T
Æmotion StudioandClaude Opus 5 0d7c834d33 feat: shader noise as a NOISE object, for custom sampling 
A standing request asked for a SamplerCustomAdvanced version of the node, so a
custom guider and scheduler could be used. The scheduler half already shipped:
custom_sigmas takes a SIGMAS input and the schedule is sampled, not just
counted. The gap was the guider, and it is the half that matters -- a guider is
the only way to reach BasicGuider with no negative and no CFG, DualCFGGuider,
or anything another pack provides.

This does it with a node that outputs NOISE rather than a second sampler.
SamplerCustomAdvanced consumes one, and so does AddNoise.

A second sampler node was the obvious shape and would have worked: the segment
loop already routes through CFGGuider.sample, because comfy.sample.sample
builds a guider and calls it, so accepting one would have been a substitution
rather than a new mechanism. It was not worth it. A NOISE object is forty lines
against a second wall of shader widgets, five frontend files that name nodes
explicitly, and a sampling seam through the pipeline -- and it composes with
everything in the custom sampling family instead of only with itself.

What it covers is the node's default configuration exactly: one shader stage,
painted into the noise a run starts from. That is pinned bit for bit against
the Direct node across eps, flow and AV latents, presets, batch_index and
progressions, so a seed means the same thing on both. What it cannot do is
stages, because a NOISE object is asked for noise once, before any sampling --
the sampler's stages re-enter the shader at segment boundaries partway through.
The node says so, and so do the README and the tooltips.

Two pieces were lifted out rather than copied: build_shader_params, next to the
reader that owns the params dict, and _paint/starting_noise in the pipeline,
where the painting loop was already written twice. The Direct node is sixty-odd
lines shorter for it. Both extractions were checked against the golden suite
before anything was built on them.

The equivalence test earned its keep immediately. Presets set stage_progression,
and stage_progression shapes a single stage -- it sits at the start of the
trajectory and is drawn from there -- so ignoring it made preset: roam mean two
different things on the two nodes. It is now an input here, with its own wording
because a ramp across stages is not what it does with one. That also showed the
Direct node's tooltip has always been wrong to say it needs more than one stage
to do anything; the tooltip is corrected and the sampling is untouched, which is
how roam and video have always behaved.

Also dropped a refusal check copied into starting_noise: unlike run, there is no
sampling to protect there, and the generator raises the same error one call
down. And _STREAM_SEED_STRIDE turned out to have nothing pinning it anywhere --
the nested golden case paints only the spatial stream, so zeroing the offset
changed nothing in the suite. It has a test now, feeding two identical streams
so the inputs cannot explain the difference.

Verified on H3 through BasicGuider, KSamplerSelect and BasicScheduler into
SamplerCustomAdvanced: seven steps, coherent output.
example_workflows/MiniMaxH3_CustomSampling_SNK_Source.json is that run, and
needs core nodes and this pack and nothing else.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-09-22 00:29:09 -07:00

1592 lines
85 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
import json
import os
import torch
import math
class ShaderParamsReader:
"""
Class for reading and applying shader parameters from JSON file
Implements the Lt=Sα(N)∘Kβ(t) pattern where shader transforms are applied to noise before sampling
"""
# Define valid parameter values for security whitelisting
VALID_SHADER_TYPES = {
"tensor_field", "cellular", "domain_warp", "fractal", "perlin",
"waves", "gaussian", "heterogeneous_fbm", "interference_patterns",
"spectral", "projection_3d", "curl_noise"
}
VALID_SHAPE_TYPES = {
"none", "circle", "square", "radial", "star", "linear",
"radial_animated", "spiral", "checkerboard", "spots", "hexgrid",
"stripes", "radial_gradient_static", "gradient", "vignette",
"cross", "triangles", "concentric", "rays", "zigzag",
"gradient_x", "gradient_y", "stars"
}
# Legacy mapping for integer shape types to string identifiers
# 1: circle/radial (old), 2: square, 3: star
LEGACY_SHAPE_MAPPING = {
1: "circle",
2: "square",
3: "star",
"1": "circle",
"2": "square",
"3": "star"
}
VALID_COLOR_SCHEMES = {
"none", "rgb", "complementary", "monochrome", "gradient",
"blue_red", "viridis", "plasma", "inferno", "magma", "turbo",
"jet", "rainbow", "cool", "hot", "parula", "hsv", "autumn",
"winter", "spring", "summer", "copper", "pink", "bone",
"ocean", "terrain", "neon", "fire", "fantasy"
}
@staticmethod
def smoothstep(edge0, edge1, x):
"""
GLSL-style smoothstep.
"""
# Ensure edges are tensors for broadcasting with x
if not isinstance(edge0, torch.Tensor):
edge0 = torch.full_like(x, float(edge0), device=x.device, dtype=x.dtype)
if not isinstance(edge1, torch.Tensor):
edge1 = torch.full_like(x, float(edge1), device=x.device, dtype=x.dtype)
# Calculate t, handling potential edge0 >= edge1 cases by clamping
delta = edge1 - edge0
# Avoid division by zero or near-zero, maintain sign for correct 1-smoothstep
safe_delta = torch.where(torch.abs(delta) < 1e-8, torch.sign(delta) * 1e-8 + 1e-8*(1-torch.abs(torch.sign(delta))), delta)
t = torch.clamp((x - edge0) / safe_delta, 0.0, 1.0)
return t * t * (3.0 - 2.0 * t)
@staticmethod
def random_val(coords, base_seed, seed_offset):
"""
Generate a random-like value based on coordinates and seed.
Matches the random_val helper in CurlNoiseGenerator.
"""
# Use a simple hash-like function based on coordinates
# Ensuring coords are float for calculations
coords_float = coords.float()
hash_val = torch.sin(coords_float[:, :, :, 0] * (12.9898 + seed_offset) + coords_float[:, :, :, 1] * (78.233 + seed_offset)) * 43758.5453
return torch.frac(hash_val)
@staticmethod
def validate_and_sanitize_params(params):
"""
Validates and sanitizes shader parameters to prevent DoS or unexpected behavior.
Clamps values to reasonable ranges and ensures correct types.
"""
sanitized = params.copy()
# 1. Octaves: Clamp to reasonable range (e.g., 1-20) to prevent massive loops
# Check both parameter names
for key in ["octaves", "shaderOctaves"]:
if key in sanitized:
try:
# Convert to float first to handle string representations of floats
val = float(sanitized[key])
# Clamp between 1 and 20, and convert to int
sanitized[key] = int(max(1.0, min(val, 20.0)))
except (ValueError, TypeError):
print(f"Warning: Invalid octaves value '{sanitized[key]}', defaulting to 3")
sanitized[key] = 3
# 2. Scale: Ensure float and clamp to prevent numerical instability
for key in ["scale", "shaderScale"]:
if key in sanitized:
try:
val = float(sanitized[key])
if math.isnan(val) or math.isinf(val):
val = 1.0
# Clamp to avoid extremely large values
sanitized[key] = max(-1000000.0, min(val, 1000000.0))
except (ValueError, TypeError):
sanitized[key] = 1.0
# 3. Intensity/Strength: Ensure float and clamp to reasonable range
# Though some shaders might allow > 1, extremely high values can cause issues
# Check both snake_case (internal) and camelCase (frontend) key names
for key in ["intensity", "shaderColorIntensity", "shapemaskstrength", "shapeMaskStrength",
"shaderShapeStrength", "warp_strength", "shaderWarpStrength",
"phase_shift", "shaderPhaseShift"]:
if key in sanitized:
try:
val = float(sanitized[key])
if math.isnan(val) or math.isinf(val):
val = 0.0 if "strength" in key.lower() or "shift" in key.lower() else 1.0
# Clamp strictly to reasonable limits (e.g. +/- 1M) to prevent numerical instability
# This prevents DoS via numerical overflow or resource exhaustion
sanitized[key] = max(-1000000.0, min(val, 1000000.0))
except (ValueError, TypeError):
sanitized[key] = 0.0 if "strength" in key.lower() or "shift" in key.lower() else 1.0
# 4. Validate Seeds: Ensure they are within safe integer range for PyTorch
# PyTorch manual_seed expects 64-bit signed integer (approx +/- 9e18)
# Using a slightly safer range to avoid boundary issues
MAX_SEED = 9000000000000000000
MIN_SEED = -9000000000000000000
for key in ["seed", "base_seed"]:
if key in sanitized:
try:
# Check for float inputs first to catch Infinity
if isinstance(sanitized[key], float):
if math.isinf(sanitized[key]) or math.isnan(sanitized[key]):
sanitized[key] = 0
continue
val = int(sanitized[key])
# Clamp to safe range to prevent runtime crashes (DoS)
sanitized[key] = max(MIN_SEED, min(val, MAX_SEED))
except (ValueError, TypeError, OverflowError):
sanitized[key] = 0
# 5. Validate String Enums (Shader Type, Shape Type, Color Scheme)
# Prevent arbitrary strings from flowing through the system
# Check both snake_case (internal) and camelCase (frontend) key names
for key in ["shader_type", "shaderType"]:
if key in sanitized:
st = str(sanitized[key]).lower()
# Handle some common aliases before validation
if st == "tensorfield": st = "tensor_field"
if st == "heterogeneousfbm": st = "heterogeneous_fbm"
if st == "projection3d": st = "projection_3d"
if st == "curl": st = "curl_noise"
if st not in ShaderParamsReader.VALID_SHADER_TYPES:
print(f"Warning: Invalid {key} '{st}', defaulting to 'tensor_field'")
sanitized[key] = "tensor_field"
else:
sanitized[key] = st
for shape_key in ["shape_type", "shaderShapeType"]:
if shape_key in sanitized:
shape_val = sanitized[shape_key]
# Handle integer inputs for legacy shape types (1, 2, 3)
# and map them to their string equivalents if valid
is_legacy = False
if isinstance(shape_val, int) or (isinstance(shape_val, str) and shape_val.isdigit()):
# Convert to integer for lookup (handles string "1" and int 1)
try:
lookup_key = int(shape_val)
if lookup_key in ShaderParamsReader.LEGACY_SHAPE_MAPPING:
# Map to valid string name
sanitized[shape_key] = ShaderParamsReader.LEGACY_SHAPE_MAPPING[lookup_key]
is_legacy = True
else:
print(f"Warning: Invalid legacy integer {shape_key} '{shape_val}', defaulting to 'none'")
sanitized[shape_key] = "none"
is_legacy = True
except (ValueError, TypeError):
# Fallthrough to string handling if conversion fails weirdly
pass
# If not a handled legacy integer, treat as string identifier
if not is_legacy:
st = str(shape_val).lower()
if st not in ShaderParamsReader.VALID_SHAPE_TYPES:
print(f"Warning: Invalid {shape_key} '{st}', defaulting to 'none'")
sanitized[shape_key] = "none"
else:
sanitized[shape_key] = st
if "colorScheme" in sanitized:
cs = str(sanitized["colorScheme"]).lower()
if cs not in ShaderParamsReader.VALID_COLOR_SCHEMES:
print(f"Warning: Invalid colorScheme '{cs}', defaulting to 'none'")
sanitized["colorScheme"] = "none"
else:
sanitized["colorScheme"] = cs
return sanitized
@staticmethod
def get_shader_params(custom_path=None):
"""
Utility function to read shader parameters from file.
Returns a dictionary of shader parameters.
Args:
custom_path: Optional path to a custom JSON file
Returns:
Dictionary of shader parameters
"""
# Default values in case file doesn't exist or is invalid
default_params = {
"shader_type": "tensor_field",
"visualization_type": 3, # ellipses
"scale": 1.0,
"phase_shift": 0.0,
"warp_strength": 0.5,
"time": 0.0,
"octaves": 3.0,
"intensity": 0.8, # influence/strength of the shader
"shapemaskstrength": 1.0, # strength of the shape mask
"shape_type": "none" # type of shape mask
}
# Get the extension directory (where this file is located)
EXTENSION_DIR = os.path.dirname(os.path.abspath(__file__))
# Path to the shader_params.json file (default or custom)
if custom_path:
# Security check for path traversal - resolve symlinks
try:
resolved_path = os.path.realpath(custom_path)
extension_real_path = os.path.realpath(EXTENSION_DIR)
# Use realpath to ensure we are comparing canonical paths (handles symlinks and casing)
data_dir_real_path = os.path.realpath(os.path.join(extension_real_path, "data"))
default_config_real_path = os.path.realpath(os.path.join(extension_real_path, "shader_params.json"))
# 1. Strict extension check
if not resolved_path.lower().endswith('.json'):
print(f"SECURITY WARNING: Invalid file extension (must be .json): {custom_path}")
is_safe = False
else:
# 2. Strict location check: Must be in data/ OR be the root shader_params.json
# Use normcase for platform-appropriate case normalization
# (lowercases on Windows, preserves case on Linux)
resolved_norm = os.path.normcase(resolved_path)
extension_norm = os.path.normcase(extension_real_path)
data_dir_norm = os.path.normcase(data_dir_real_path)
default_config_norm = os.path.normcase(default_config_real_path)
# 2. Security Check: File must be physically inside the extension directory
# This blocks symlinks pointing outside the extension folder
is_inside_extension = os.path.commonpath([resolved_norm, extension_norm]) == extension_norm
# 3. Scope Check: File must be in data/ or be the config file
# This blocks reading source code or secrets in the extension root
is_in_data = os.path.commonpath([resolved_norm, data_dir_norm]) == data_dir_norm
is_default_config = resolved_norm == default_config_norm
is_safe = is_inside_extension and (is_in_data or is_default_config)
except (ValueError, OSError):
is_safe = False
if not is_safe:
print(f"SECURITY WARNING: Prevented access to unauthorized file: {custom_path}")
# Fallback to default path instead of opening potentially dangerous file
params_file = os.path.join(EXTENSION_DIR, "shader_params.json")
if not os.path.exists(params_file):
params_file = os.path.join(EXTENSION_DIR, "data", "shader_params.json")
else:
params_file = resolved_path
else:
# Try to find params in root directory first
params_file = os.path.join(EXTENSION_DIR, "shader_params.json")
# If not found, try the data folder
if not os.path.exists(params_file):
params_file = os.path.join(EXTENSION_DIR, "data", "shader_params.json")
# Try to read the shader parameters from the file
try:
if os.path.exists(params_file):
# print(f"Found parameters file: {params_file}")
with open(params_file, 'r') as f:
loaded_params = json.load(f)
# Map between different parameter naming conventions
param_mapping = {
"shaderType": "shader_type",
"shaderScale": "scale",
"shaderOctaves": "octaves",
"shaderWarpStrength": "warp_strength",
"shaderPhaseShift": "phase_shift",
"shapeMaskStrength": "shapemaskstrength",
"shaderShapeStrength": "shapemaskstrength",
"shaderShapeType": "shape_type"
}
# Output raw loaded params for debugging
# print(f"Raw JSON params: {loaded_params}")
# Convert parameter names if needed
params = {}
for key, value in loaded_params.items():
if key in param_mapping:
params[param_mapping[key]] = value
else:
params[key] = value
# Special handling for shaderColorIntensity to maintain both versions
if "shaderColorIntensity" in loaded_params:
# Keep the original key
params["shaderColorIntensity"] = loaded_params["shaderColorIntensity"]
# Also provide as intensity for backward compatibility
params["intensity"] = loaded_params["shaderColorIntensity"]
# Handle specific shader type mapping
if "shader_type" in params:
shader_type = params["shader_type"]
# Convert string values to standardized format
if shader_type.lower() == "tensor_field" or shader_type.lower() == "tensorfield":
params["shader_type"] = "tensor_field"
elif shader_type.lower() == "heterogeneous_fbm" or shader_type.lower() == "heterogeneousfbm":
params["shader_type"] = "heterogeneous_fbm"
elif shader_type.lower() == "projection_3d" or shader_type.lower() == "projection3d":
params["shader_type"] = "projection_3d"
elif shader_type.lower() == "cellular":
params["shader_type"] = "cellular"
elif shader_type.lower() == "fractal":
params["shader_type"] = "fractal"
elif shader_type.lower() == "perlin":
params["shader_type"] = "perlin"
elif shader_type.lower() == "waves":
params["shader_type"] = "waves"
elif shader_type.lower() == "gaussian":
params["shader_type"] = "gaussian"
elif shader_type.lower() == "domain_warp":
params["shader_type"] = "domain_warp"
elif shader_type.lower() == "interference" or shader_type.lower() == "interference_patterns":
params["shader_type"] = "interference_patterns"
print(f"Mapped shader type '{shader_type}' to 'interference_patterns'")
elif shader_type.lower() == "spectral" or shader_type.lower() == "spectral_noise":
params["shader_type"] = "spectral"
print(f"Mapped shader type '{shader_type}' to 'spectral'")
elif shader_type.lower() == "projection" or shader_type.lower() == "projection_3d" or shader_type.lower() == "3d_projection":
params["shader_type"] = "projection_3d"
print(f"Mapped shaderType '{shader_type}' to 'projection_3d'")
elif shader_type.lower() == "curl" or shader_type.lower() == "curl_noise":
params["shader_type"] = "curl_noise"
# print(f"Mapped shaderType '{shader_type}' to 'curl_noise'") # Mapped to curl_noise
# Also check if shader type is in the shaderType field (alternate field name)
if "shaderType" in loaded_params and "shader_type" not in params:
shader_type = loaded_params["shaderType"]
if isinstance(shader_type, str):
if shader_type.lower() == "tensor_field" or shader_type.lower() == "tensorfield":
params["shader_type"] = "tensor_field"
elif shader_type.lower() == "heterogeneous_fbm" or shader_type.lower() == "heterogeneousfbm":
params["shader_type"] = "heterogeneous_fbm"
elif shader_type.lower() == "projection_3d" or shader_type.lower() == "projection3d":
params["shader_type"] = "projection_3d"
elif shader_type.lower() == "cellular":
params["shader_type"] = "cellular"
elif shader_type.lower() == "fractal":
params["shader_type"] = "fractal"
elif shader_type.lower() == "perlin":
params["shader_type"] = "perlin"
elif shader_type.lower() == "waves":
params["shader_type"] = "waves"
elif shader_type.lower() == "gaussian":
params["shader_type"] = "gaussian"
elif shader_type.lower() == "domain_warp":
params["shader_type"] = "domain_warp"
elif shader_type.lower() == "interference" or shader_type.lower() == "interference_patterns":
params["shader_type"] = "interference_patterns"
print(f"Mapped shaderType '{shader_type}' to 'interference_patterns'")
elif shader_type.lower() == "spectral" or shader_type.lower() == "spectral_noise":
params["shader_type"] = "spectral"
print(f"Mapped shaderType '{shader_type}' to 'spectral'")
# Validate and sanitize loaded parameters before merging
params = ShaderParamsReader.validate_and_sanitize_params(params)
# Fill in any missing parameters with defaults
for key, value in default_params.items():
if key not in params:
params[key] = value
# print(f"Successfully loaded shader parameters: {params}")
return params
else:
print(f"Parameters file not found at: {params_file}")
except Exception as e:
print(f"Error loading shader parameters: {e}")
print(f"Using default shader parameters")
return default_params
@staticmethod
def apply_shader_to_noise(noise, shader_params=None, influence=None):
"""
Apply shader effects to the initial noise before sampling
Implements the Sα(N) part of Lt=Sα(N)∘Kβ(t)
Args:
noise: Initial noise tensor [batch, channels, height, width]
shader_params: Dictionary of shader parameters (or None to load from file)
influence: How much to blend shader noise (0.0-1.0, None uses value from params)
Returns:
Modified noise tensor with same shape as input
"""
if shader_params is None:
shader_params = ShaderParamsReader.get_shader_params()
# Extract basic parameters
batch, channels, height, width = noise.shape
device = noise.device
# Use provided influence or get from parameters
if influence is None:
influence = shader_params.get("intensity", 0.8)
# Ensure influence is a float
influence = float(influence)
# Skip if no influence
if influence <= 0.0:
return noise
# Extract shader parameters
shader_type = shader_params.get("shader_type", "tensor_field")
viz_type = shader_params.get("visualization_type", 3) # default to ellipses
scale = shader_params.get("scale", 1.0)
phase_shift = shader_params.get("phase_shift", 0.0)
warp_strength = shader_params.get("warp_strength", 0.5)
time = shader_params.get("time", 0.0)
octaves = shader_params.get("octaves", 3.0)
seed = shader_params.get("seed", 0)
# Extract shape mask parameters
shape_type = shader_params.get("shape_type", "none")
shape_mask_strength = shader_params.get("shapemaskstrength", 1.0)
# Debug print for shape mask parameters
print(f"Shape mask parameters: type={shape_type}, strength={shape_mask_strength}")
# Create coordinate grid (normalized to [-1, 1])
y, x = torch.meshgrid(torch.linspace(-1, 1, height, device=device),
torch.linspace(-1, 1, width, device=device),
indexing='ij')
# Combine into coordinate tensor
p = torch.stack([x, y], dim=-1).unsqueeze(0).repeat(batch, 1, 1, 1)
# Generate different shader patterns
if False: # Placeholder for any future shader types to be handled here
pass # Generate shader_noise for other types if needed
else:
# If the shader type was one of the removed ones or is not handled,
# print a message and return the original noise unchanged.
print(f"Shader type '{shader_type}' is not handled by apply_shader_to_noise or its simple implementation was removed. Returning original noise.")
return noise
# -- REMOVED Unreachable code: permutation, normalization, expansion, blending --
@staticmethod
def _lerp(a, b, t):
"""Helper for linear interpolation."""
return a + (b - a) * t
@staticmethod
def _hsv_to_rgb(h, s, v):
"""
Convert HSV to RGB.
h, s, v are expected in [0,1] range and shape [B, 1, H, W].
Returns R, G, B components, each as [B, 1, H, W] in [0,1] range.
"""
# Ensure inputs are correctly shaped for broadcasting if they are single values
if not isinstance(h, torch.Tensor): h = torch.full_like(s if isinstance(s, torch.Tensor) else v, float(h)) # Fallback for s or v if h is scalar
if not isinstance(s, torch.Tensor): s = torch.full_like(h, float(s))
if not isinstance(v, torch.Tensor): v = torch.full_like(h, float(v))
c = v * s
h_prime = h * 6.0 # h is [0,1]
# Ensure h_prime is a tensor for fmod
if not isinstance(h_prime, torch.Tensor):
h_prime = torch.full_like(c, float(h_prime))
x = c * (1.0 - torch.abs(torch.fmod(h_prime, 2.0) - 1.0))
m = v - c
r, g, b = torch.zeros_like(h), torch.zeros_like(h), torch.zeros_like(h)
# Masks for hue ranges
mask0 = (h_prime < 1.0)
mask1 = (h_prime >= 1.0) & (h_prime < 2.0)
mask2 = (h_prime >= 2.0) & (h_prime < 3.0)
mask3 = (h_prime >= 3.0) & (h_prime < 4.0)
mask4 = (h_prime >= 4.0) & (h_prime < 5.0)
mask5 = (h_prime >= 5.0) # covers up to 6.0
# Assign R, G, B based on hue
r[mask0], g[mask0], b[mask0] = c[mask0], x[mask0], torch.zeros_like(x)[mask0]
r[mask1], g[mask1], b[mask1] = x[mask1], c[mask1], torch.zeros_like(x)[mask1]
r[mask2], g[mask2], b[mask2] = torch.zeros_like(x)[mask2], c[mask2], x[mask2]
r[mask3], g[mask3], b[mask3] = torch.zeros_like(x)[mask3], x[mask3], c[mask3]
r[mask4], g[mask4], b[mask4] = x[mask4], torch.zeros_like(x)[mask4], c[mask4]
r[mask5], g[mask5], b[mask5] = c[mask5], torch.zeros_like(x)[mask5], x[mask5]
r, g, b = r + m, g + m, b + m
return r, g, b
@staticmethod
def _interpolate_colors(stops, t):
"""
Interpolate colors based on stops using vectorized operations.
t is a normalized value tensor [B, 1, H, W] in [0,1] range.
stops: list of [value, color_tuple_or_tensor e.g. (R,G,B) or [1,3,1,1] tensor].
Returns R, G, B components, each as [B, 1, H, W] in [0,1] range.
"""
device = t.device
dtype = t.dtype
# Prepare stops tensors
stop_vals = []
stop_colors = []
for val, color_val in stops:
stop_vals.append(float(val))
if isinstance(color_val, (list, tuple)):
c_tensor = torch.tensor(color_val, device=device, dtype=dtype).view(1, 3, 1, 1)
else: # assume it's already a tensor
c_tensor = color_val.to(device=device, dtype=dtype)
if c_tensor.numel() == 3:
c_tensor = c_tensor.view(1, 3, 1, 1)
stop_colors.append(c_tensor)
# Create tensors for bucketize/gather
stop_vals_tensor = torch.tensor(stop_vals, device=device, dtype=dtype)
# Concatenate colors to [num_stops, 3] for indexing (remove spatial dims for now)
stop_colors_stack = torch.cat([c.view(1, 3) for c in stop_colors], dim=0)
# Find indices where elements should be inserted to maintain order
# bucketize returns indices such that stop_vals[i-1] <= t < stop_vals[i]
indices = torch.bucketize(t, stop_vals_tensor)
# Clamp indices to be within [1, num_stops-1]
# This ensures we always have a valid previous stop (idx-1) and current stop (idx)
# For t < stops[0], indices=0 -> clamped to 1. Uses segment (stops[0], stops[1]).
# For t >= stops[-1], indices=num_stops -> clamped to num_stops-1. Uses segment (stops[-2], stops[-1]).
idxs = torch.clamp(indices, 1, len(stops) - 1)
# Gather start and end values/colors for each pixel's segment
# idxs is [B, 1, H, W], stop_vals_tensor is [N]
# Advanced indexing: we want result [B, 1, H, W]
# Flatten t and idxs for simpler gathering if needed, but PyTorch handles this
t0 = stop_vals_tensor[idxs - 1] # [B, 1, H, W]
t1 = stop_vals_tensor[idxs] # [B, 1, H, W]
# Gather colors
# stop_colors_stack is [N, 3]. idxs is [B, 1, H, W]
# c0 will be [B, 1, H, W, 3]
c0 = stop_colors_stack[idxs - 1]
c1 = stop_colors_stack[idxs]
# Permute to [B, 3, H, W] and squeeze singleton dimension from original idxs indexing
# Note: Indexing with [B, 1, H, W] into [N, 3] creates [B, 1, H, W, 3]
c0 = c0.permute(0, 4, 2, 3, 1).squeeze(-1) # [B, 3, H, W]
c1 = c1.permute(0, 4, 2, 3, 1).squeeze(-1) # [B, 3, H, W]
# Calculate local interpolation factor
denominator = (t1 - t0)
# Avoid division by zero
safe_denominator = torch.where(torch.abs(denominator) < 1e-8, torch.ones_like(denominator), denominator)
local_t = (t - t0) / safe_denominator
local_t = torch.clamp(local_t, 0.0, 1.0) # [B, 1, H, W]
# Interpolate
# c0, c1 are [B, 3, H, W], local_t is [B, 1, H, W] (broadcasts)
final_color = ShaderParamsReader._lerp(c0, c1, local_t)
# Handle strict out of bounds values (below first stop or above last stop)
# If t < stop[0], local_t was computed relative to stop[0] and stop[1].
# It will be negative, clamped to 0. So result = c0 = stop[0]. Correct.
# If t > stop[-1], local_t > 1, clamped to 1. Result = c1 = stop[-1]. Correct.
return final_color[:, 0:1], final_color[:, 1:2], final_color[:, 2:3]
@staticmethod
def apply_color_scheme(noise_tensor, shader_params=None):
"""
Apply color scheme to a shader noise tensor based on shader_params
Args:
noise_tensor: Input noise tensor of shape [batch, channels, height, width]
shader_params: Dictionary of shader parameters (or None to load from file)
Returns:
Modified noise tensor with color scheme applied
"""
if shader_params is None:
shader_params = ShaderParamsReader.get_shader_params()
# Get color scheme and intensity parameters
color_scheme = shader_params.get("colorScheme", "none")
# Try to get the color intensity with priority for shaderColorIntensity
color_intensity = shader_params.get("shaderColorIntensity",
shader_params.get("intensity", 0.8))
# Skip if no color scheme or zero intensity
if color_scheme == "none" or color_intensity <= 0.0:
print(f"Skipping color scheme application: scheme={color_scheme}, intensity={color_intensity}")
return noise_tensor
print(f"APPLYING COLOR SCHEME: {color_scheme} with intensity {color_intensity}")
# Extract dimensions
batch, channels, height, width = noise_tensor.shape
device = noise_tensor.device
# Create empty color tensor that we'll fill based on the scheme
color_tensor = torch.zeros_like(noise_tensor)
# Make sure to preserve the 4th channel if it exists
if channels > 3:
color_tensor[:, 3:] = noise_tensor[:, 3:]
# Helper function to normalize the noise to 0-1 range for colormaps
def normalize_to_01(tensor):
return (tensor - tensor.min()) / (tensor.max() - tensor.min() + 1e-8)
# Map base noise to 0-1 for color mapping
base_noise = normalize_to_01(noise_tensor[:, 0])
# Create a [B, 1, H, W] version of base_noise for helpers
t_color = base_noise.unsqueeze(1)
# Handle different color schemes
if color_scheme == "rgb":
# RGB color scheme: create three distinct channels
if channels >= 3:
# R channel - emphasize details in first latent dimension
color_tensor[:, 0] = noise_tensor[:, 0] * 1.5
# G channel - use second latent dimension with slight phase shift
color_tensor[:, 1] = noise_tensor[:, 1] * 1.3
# B channel - use third latent dimension with different scaling
color_tensor[:, 2] = noise_tensor[:, 2] * 0.8
elif color_scheme == "complementary":
# Complementary colors: create opposing patterns in different channels
if channels >= 3:
# First channel - original
color_tensor[:, 0] = noise_tensor[:, 0] * 1.5
# Second channel - inverted phase from channel 0
color_tensor[:, 1] = -noise_tensor[:, 0] * 0.8
# Third channel - different frequency
color_tensor[:, 2] = noise_tensor[:, 2] * 1.2
elif color_scheme == "monochrome":
# Monochrome: apply the same pattern to all channels with slight variations
if channels > 1:
base_channel = noise_tensor[:, 0:1].clone()
# Expand to all channels with slight variations in scaling
scales = torch.tensor([1.0, 0.95, 0.9, 0.85][:channels], device=device).view(1, -1, 1, 1)
color_tensor = base_channel * scales
elif color_scheme == "gradient":
# Gradient: create a position-based color gradient
if channels >= 3:
# Create coordinate grid for gradient
y_norm = torch.linspace(0, 1, height, device=device).view(1, 1, -1, 1).expand(batch, 1, -1, width)
x_norm = torch.linspace(0, 1, width, device=device).view(1, 1, 1, -1).expand(batch, 1, height, -1)
# R channel - horizontal gradient + noise
color_tensor[:, 0:1] = x_norm + noise_tensor[:, 0:1] * 0.4
# G channel - vertical gradient + noise
color_tensor[:, 1:2] = y_norm + noise_tensor[:, 1:2] * 0.4
# B channel - diagonal gradient + noise
color_tensor[:, 2:3] = (x_norm + y_norm) / 2 + noise_tensor[:, 2:3] * 0.4
elif color_scheme == "blue_red":
if channels >= 3:
# Blue to red gradient (cold to hot) using lerp
c0 = torch.tensor([0.0, 0.0, 1.0], device=device, dtype=t_color.dtype).view(1, 3, 1, 1) # Blue
c1 = torch.tensor([1.0, 0.0, 0.0], device=device, dtype=t_color.dtype).view(1, 3, 1, 1) # Red
# _lerp expects t_color to be broadcastable with c0, c1.
# t_color is [B,1,H,W], c0/c1 are [1,3,1,1]. Result is [B,3,H,W]
interpolated_color = ShaderParamsReader._lerp(c0, c1, t_color)
color_tensor[:, 0:1] = interpolated_color[:, 0:1] # Red
color_tensor[:, 1:2] = interpolated_color[:, 1:2] # Green
color_tensor[:, 2:3] = interpolated_color[:, 2:3] # Blue
elif color_scheme == "viridis":
if channels >= 3:
stops = [
(0.0, (0.267, 0.005, 0.329)), # #440154
(0.33, (0.188, 0.407, 0.553)), # #30678D
(0.66, (0.208, 0.718, 0.471)), # #35B778
(1.0, (0.992, 0.906, 0.143)) # #FDE724
]
r, g, b = ShaderParamsReader._interpolate_colors(stops, t_color)
color_tensor[:, 0:1], color_tensor[:, 1:2], color_tensor[:, 2:3] = r, g, b
elif color_scheme == "plasma":
if channels >= 3:
# Use robust color stops for plasma, matching curl_noise.py
stops = [
(0.0, (0.05, 0.03, 0.53)),
(0.25, (0.40, 0.00, 0.66)),
(0.5, (0.70, 0.18, 0.53)),
(0.75, (0.94, 0.46, 0.25)),
(1.0, (0.98, 0.80, 0.08))
]
r, g, b = ShaderParamsReader._interpolate_colors(stops, t_color)
color_tensor[:, 0:1], color_tensor[:, 1:2], color_tensor[:, 2:3] = r, g, b
elif color_scheme == "inferno":
if channels >= 3:
stops = [
(0.0, (0.001, 0.001, 0.016)),
(0.25, (0.259, 0.039, 0.408)),
(0.5, (0.576, 0.149, 0.404)),
(0.75, (0.867, 0.318, 0.227)),
(0.85, (0.988, 0.647, 0.039)),
(1.0, (0.988, 1.000, 0.643))
]
r, g, b = ShaderParamsReader._interpolate_colors(stops, t_color)
color_tensor[:, 0:1], color_tensor[:, 1:2], color_tensor[:, 2:3] = r, g, b
elif color_scheme == "magma":
if channels >= 3:
stops = [
(0.0, (0.001, 0.001, 0.016)),
(0.25, (0.231, 0.059, 0.439)),
(0.5, (0.549, 0.161, 0.506)),
(0.75, (0.871, 0.288, 0.408)),
(0.85, (0.996, 0.624, 0.427)),
(1.0, (0.988, 0.992, 0.749))
]
r, g, b = ShaderParamsReader._interpolate_colors(stops, t_color)
color_tensor[:, 0:1], color_tensor[:, 1:2], color_tensor[:, 2:3] = r, g, b
elif color_scheme == "turbo":
if channels >= 3:
stops = [
(0.0, (0.188, 0.071, 0.235)),
(0.25, (0.275, 0.408, 0.859)),
(0.5, (0.149, 0.749, 0.549)),
(0.65, (0.831, 1.000, 0.314)),
(0.85, (0.980, 0.718, 0.298)),
(1.0, (0.729, 0.004, 0.000))
]
r, g, b = ShaderParamsReader._interpolate_colors(stops, t_color)
# Turbo often benefits from a slight boost/rescale
r, g, b = r * 1.2 - 0.1, g * 1.2 - 0.1, b * 1.2 - 0.1
color_tensor[:, 0:1], color_tensor[:, 1:2], color_tensor[:, 2:3] = torch.clamp(r, 0, 1), torch.clamp(g, 0, 1), torch.clamp(b, 0, 1)
elif color_scheme == "jet":
if channels >= 3:
stops = [
(0.0, (0.000, 0.000, 0.5)), # Dark Blue
(0.125, (0.000, 0.000, 1.000)),# Blue
(0.375, (0.000, 1.000, 1.000)),# Cyan
(0.625, (1.000, 1.000, 0.000)),# Yellow
(0.875, (1.000, 0.000, 0.000)),# Red
(1.0, (0.500, 0.000, 0.000)) # Dark Red
]
r, g, b = ShaderParamsReader._interpolate_colors(stops, t_color)
color_tensor[:, 0:1], color_tensor[:, 1:2], color_tensor[:, 2:3] = r, g, b
elif color_scheme == "rainbow":
if channels >= 3:
# Use HSV to RGB for rainbow: hue from t_color, constant saturation and value
hue = t_color # base_noise is already [B,1,H,W] and [0,1]
saturation = torch.ones_like(hue) * 0.9 # High saturation
value = torch.ones_like(hue) * 0.9 # Bright value
r, g, b = ShaderParamsReader._hsv_to_rgb(hue, saturation, value)
color_tensor[:, 0:1], color_tensor[:, 1:2], color_tensor[:, 2:3] = r, g, b
elif color_scheme == "cool":
if channels >= 3:
c0 = torch.tensor([0.0, 1.0, 1.0], device=device, dtype=t_color.dtype).view(1, 3, 1, 1) # Cyan
c1 = torch.tensor([1.0, 0.0, 1.0], device=device, dtype=t_color.dtype).view(1, 3, 1, 1) # Magenta
interpolated_color = ShaderParamsReader._lerp(c0, c1, t_color)
color_tensor[:, 0:1], color_tensor[:, 1:2], color_tensor[:, 2:3] = interpolated_color[:,0:1], interpolated_color[:,1:2], interpolated_color[:,2:3]
elif color_scheme == "hot":
if channels >= 3:
stops = [
(0.0, (0.0, 0.0, 0.0)), # Black
(0.375, (1.0, 0.0, 0.0)), # Red
(0.75, (1.0, 1.0, 0.0)), # Yellow
(1.0, (1.0, 1.0, 1.0)) # White
]
r, g, b = ShaderParamsReader._interpolate_colors(stops, t_color)
color_tensor[:, 0:1], color_tensor[:, 1:2], color_tensor[:, 2:3] = r, g, b
elif color_scheme == "parula":
if channels >= 3:
stops = [
(0.0, (0.208, 0.165, 0.529)), # #352a87
(0.25, (0.059, 0.361, 0.867)), # #0f5cdd
(0.5, (0.000, 0.710, 0.651)), # #00b5a6
(0.75, (1.000, 0.765, 0.216)), # #ffc337
(1.0, (0.988, 0.996, 0.643)) # #fcfea4
]
r, g, b = ShaderParamsReader._interpolate_colors(stops, t_color)
color_tensor[:, 0:1], color_tensor[:, 1:2], color_tensor[:, 2:3] = r, g, b
elif color_scheme == "hsv":
if channels >= 3:
hue = t_color
saturation = torch.ones_like(hue) * 0.95 # Full saturation
value = torch.ones_like(hue) * 0.95 # Full value
r, g, b = ShaderParamsReader._hsv_to_rgb(hue, saturation, value)
# Original SPR HSV scaled output to [-1,1]. We keep [0,1] from _hsv_to_rgb for consistency with other interpolated.
color_tensor[:, 0:1], color_tensor[:, 1:2], color_tensor[:, 2:3] = r, g, b
elif color_scheme == "autumn":
if channels >= 3:
c0 = torch.tensor([1.0, 0.0, 0.0], device=device, dtype=t_color.dtype).view(1, 3, 1, 1) # Red
c1 = torch.tensor([1.0, 1.0, 0.0], device=device, dtype=t_color.dtype).view(1, 3, 1, 1) # Yellow
interpolated_color = ShaderParamsReader._lerp(c0, c1, t_color)
color_tensor[:, 0:1], color_tensor[:, 1:2], color_tensor[:, 2:3] = interpolated_color[:,0:1], interpolated_color[:,1:2], interpolated_color[:,2:3]
elif color_scheme == "winter":
if channels >= 3:
c0 = torch.tensor([0.0, 0.0, 1.0], device=device, dtype=t_color.dtype).view(1, 3, 1, 1) # Blue
c1 = torch.tensor([0.0, 1.0, 0.5], device=device, dtype=t_color.dtype).view(1, 3, 1, 1) # Greenish-Cyan
interpolated_color = ShaderParamsReader._lerp(c0, c1, t_color)
color_tensor[:, 0:1], color_tensor[:, 1:2], color_tensor[:, 2:3] = interpolated_color[:,0:1], interpolated_color[:,1:2], interpolated_color[:,2:3]
elif color_scheme == "spring":
if channels >= 3:
c0 = torch.tensor([1.0, 0.0, 1.0], device=device, dtype=t_color.dtype).view(1, 3, 1, 1) # Magenta
c1 = torch.tensor([1.0, 1.0, 0.0], device=device, dtype=t_color.dtype).view(1, 3, 1, 1) # Yellow
interpolated_color = ShaderParamsReader._lerp(c0, c1, t_color)
color_tensor[:, 0:1], color_tensor[:, 1:2], color_tensor[:, 2:3] = interpolated_color[:,0:1], interpolated_color[:,1:2], interpolated_color[:,2:3]
elif color_scheme == "summer":
if channels >= 3:
c0 = torch.tensor([0.0, 0.5, 0.4], device=device, dtype=t_color.dtype).view(1, 3, 1, 1) # Dark Green
c1 = torch.tensor([1.0, 1.0, 0.4], device=device, dtype=t_color.dtype).view(1, 3, 1, 1) # Yellow
interpolated_color = ShaderParamsReader._lerp(c0, c1, t_color)
color_tensor[:, 0:1], color_tensor[:, 1:2], color_tensor[:, 2:3] = interpolated_color[:,0:1], interpolated_color[:,1:2], interpolated_color[:,2:3]
elif color_scheme == "copper":
if channels >= 3:
c0 = torch.tensor([0.0, 0.0, 0.0], device=device, dtype=t_color.dtype).view(1, 3, 1, 1) # Black
c1 = torch.tensor([1.0, 0.6235, 0.3922], device=device, dtype=t_color.dtype).view(1, 3, 1, 1) # Copper color approx (255,159,100)
interpolated_color = ShaderParamsReader._lerp(c0, c1, t_color)
color_tensor[:, 0:1], color_tensor[:, 1:2], color_tensor[:, 2:3] = interpolated_color[:,0:1], interpolated_color[:,1:2], interpolated_color[:,2:3]
elif color_scheme == "pink":
if channels >= 3:
stops = [
(0.0, (0.05, 0.05, 0.05)), # Dark gray
(0.5, (1.0, 0.41, 0.71)), # Hot Pink approx
(1.0, (1.0, 0.75, 0.80)) # Light Pink
]
r, g, b = ShaderParamsReader._interpolate_colors(stops, t_color)
color_tensor[:, 0:1], color_tensor[:, 1:2], color_tensor[:, 2:3] = r, g, b
elif color_scheme == "bone":
if channels >= 3:
stops = [ # Standard bone colormap
(0.0, (0.0, 0.0, 0.0)),
(0.375, (0.3294, 0.3294, 0.4549)), # (84, 84, 116)
(0.75, (0.6275, 0.7569, 0.7569)), # (160, 193, 193)
(1.0, (1.0, 1.0, 1.0))
]
r, g, b = ShaderParamsReader._interpolate_colors(stops, t_color)
# Original shader_params_reader 'bone' scaled to [-1,1]. Let's keep [0,1] for consistency.
color_tensor[:, 0:1], color_tensor[:, 1:2], color_tensor[:, 2:3] = r, g, b
elif color_scheme == "ocean":
if channels >= 3:
stops = [ # Based on matplotlib's ocean
(0.0, (0.0, 0.0, 0.0)), # Black
(0.33, (0.0, 0.0, 0.5)), # Dark Blue
(0.66, (0.0, 0.5, 1.0)), # Light Blue
(1.0, (0.7, 1.0, 1.0)) # Very Light Cyan/White
]
r, g, b = ShaderParamsReader._interpolate_colors(stops, t_color)
color_tensor[:, 0:1], color_tensor[:, 1:2], color_tensor[:, 2:3] = r, g, b
elif color_scheme == "terrain":
if channels >= 3:
stops = [ # Standard terrain colormap
(0.0, (0.2, 0.2, 0.6)), # Deep water blue
(0.15, (0.0, 0.5, 0.0)), # Dark Green (low land)
(0.33, (0.0, 0.8, 0.4)), # Green (land)
(0.5, (0.87, 0.87, 0.4)), # Yellowish (hills)
(0.75, (0.6, 0.4, 0.2)), # Brown (mountains)
(1.0, (1.0, 1.0, 1.0)) # White (snow peaks)
]
r, g, b = ShaderParamsReader._interpolate_colors(stops, t_color)
color_tensor[:, 0:1], color_tensor[:, 1:2], color_tensor[:, 2:3] = r, g, b
elif color_scheme == "neon":
if channels >= 3:
# Using a multi-stop lerp for vibrant neon effect
stops = [
(0.0, (1.0, 0.0, 0.5)), # Magenta
(0.33, (0.0, 1.0, 1.0)), # Cyan
(0.66, (1.0, 1.0, 0.0)), # Yellow
(1.0, (0.5, 0.0, 1.0)) # Purple
]
r, g, b = ShaderParamsReader._interpolate_colors(stops, t_color)
color_tensor[:, 0:1], color_tensor[:, 1:2], color_tensor[:, 2:3] = r, g, b
elif color_scheme == "fire":
if channels >= 3:
stops = [ # Standard fire colormap
(0.0, (0.0, 0.0, 0.0)), # Black
(0.25, (1.0, 0.0, 0.0)), # Red
(0.6, (1.0, 1.0, 0.0)), # Yellow
(1.0, (1.0, 1.0, 1.0)) # White
]
r, g, b = ShaderParamsReader._interpolate_colors(stops, t_color)
color_tensor[:, 0:1], color_tensor[:, 1:2], color_tensor[:, 2:3] = r, g, b
elif color_scheme == "fantasy":
# Fantasy colors: magical and otherworldly - keeping original SPR logic
if channels >= 3:
# Create swirling color pattern
angle = torch.atan2(noise_tensor[:, 1], noise_tensor[:, 0])
radius = torch.sqrt(noise_tensor[:, 0]**2 + noise_tensor[:, 1]**2)
# Purple/pink base
color_tensor[:, 0] = torch.sin(angle * 2.0 + radius * 3.0) * 0.5 + 0.5
# Teal/blue variations
color_tensor[:, 1] = torch.sin(angle * 3.0 - radius * 2.0) * 0.5 + 0.5
# Golden highlights
color_tensor[:, 2] = torch.sin(radius * 5.0) * 0.5 + 0.5
# Normalize to maintain proper distribution
color_tensor = (color_tensor - 0.5) * 2.0
else:
# Default case - return original noise if color scheme not implemented or recognized
print(f"WARNING: Color scheme '{color_scheme}' not recognized, using original noise")
return noise_tensor
# Blend with original based on intensity
# Ensure color_tensor values are appropriately scaled if necessary before blending.
# For now, assuming [0,1] range from most new schemes is acceptable for blending.
result = noise_tensor * (1.0 - color_intensity) + color_tensor * color_intensity
print(f"Applied {color_scheme} color scheme - result shape: {result.shape}")
return result
@staticmethod
def apply_shape_mask(coords_normalized_01, shape_type, time=0.0, base_seed=0, use_temporal_coherence=False):
"""
Apply shape mask to coordinates.
Coordinates are expected to be in the [0, 1] range.
Args:
coords_normalized_01: Coordinate tensor [batch, height, width, 2] in [0, 1] range.
shape_type: Type of shape to apply (integer or string).
time: Animation time.
base_seed: Base seed for randomness if shapes require it.
use_temporal_coherence: Flag for temporal coherence.
Returns:
Shape mask tensor [batch, height, width, 1]
"""
batch, height, width, _ = coords_normalized_01.shape
device = coords_normalized_01.device
# For shapes that assume coordinates centered at (0,0) and range approx [-0.5, 0.5] or [-1,1]
# we create centered coordinates from the [0,1] input.
centered_coords = coords_normalized_01 - 0.5 # Now in [-0.5, 0.5] range
# Distance from center for centered_coords
center_dist = torch.sqrt(centered_coords[:, :, :, 0]**2 + centered_coords[:, :, :, 1]**2) # Max dist ~0.707
# Angle from center for centered_coords
angle = torch.atan2(centered_coords[:, :, :, 1], centered_coords[:, :, :, 0])
# Default mask
mask_output = torch.ones((batch, height, width), device=device)
# Convert string shape_type to standardized string format
if isinstance(shape_type, str):
shape_type = shape_type.lower()
# Handle both numeric and string shape types
# Note: Shapes from original apply_shape_mask are adapted to the new coordinate system.
# The radius/size parameters might need adjustment if they were tuned for [-1,1] p.
if shape_type == 1 or shape_type == "circle": # Original "radial" was also circle
# Circle - centered_coords range from approx -0.5 to 0.5. center_dist max ~0.707
# To make a circle that fills most of the [0,1] original space, radius should be ~0.5
# CN version: 1.0 - torch.clamp(dist * 2, 0, 1) where dist is from center of [0,1] grid
# For coords_normalized_01, dist from center (0.5,0.5) is `center_dist_01`
center_x_01, center_y_01 = 0.5, 0.5
y_diff_01 = coords_normalized_01[:, :, :, 1] - center_y_01
x_diff_01 = coords_normalized_01[:, :, :, 0] - center_x_01
dist_01 = torch.sqrt(x_diff_01**2 + y_diff_01**2)
mask_output = 1.0 - torch.clamp(dist_01 * 2.0, 0.0, 1.0) # Match CN circle
elif shape_type == 2 or shape_type == "square":
# Square - centered_coords values are in [-0.5, 0.5]
# CN version: x_mask = torch.abs(coords_bhwc[:, :, :, 0] - 0.5) * 2
# y_mask = torch.abs(coords_bhwc[:, :, :, 1] - 0.5) * 2
# dist = torch.max(x_mask, y_mask)
# shape_mask = 1.0 - torch.clamp(dist, 0, 1)
# This uses coords_normalized_01 (same as coords_bhwc in CN)
x_mask_sq = torch.abs(coords_normalized_01[:, :, :, 0] - 0.5) * 2.0
y_mask_sq = torch.abs(coords_normalized_01[:, :, :, 1] - 0.5) * 2.0
dist_sq = torch.max(x_mask_sq, y_mask_sq)
mask_output = 1.0 - torch.clamp(dist_sq, 0.0, 1.0)
elif shape_type == "radial": # Use the same logic as radial_animated (curl_noise.py radial)
# Uses coords_normalized_01 (range [0,1])
time_tensor = torch.tensor(time, device=device, dtype=coords_normalized_01.dtype)
center_x = 0.5 + 0.2 * torch.cos(time_tensor)
center_y = 0.5 + 0.2 * torch.sin(time_tensor)
# Calculate distance from the animated center using coords_normalized_01
y_diff = coords_normalized_01[:, :, :, 1] - center_y
x_diff = coords_normalized_01[:, :, :, 0] - center_x
dist_from_anim_center = torch.sqrt(x_diff**2 + y_diff**2) * 2.0 # Multiplied by 2 like in curl_noise
mask_output = torch.clamp(1.0 - dist_from_anim_center, 0.0, 1.0)
elif shape_type == 3 or shape_type == "star": # SPR original "star"
# Star-like shape - using centered_coords
points = 5.0
star_radius = 0.25 + 0.125 * torch.cos(angle * points)
mask_output = (center_dist < star_radius).float()
elif shape_type == "linear": # Ported from curl_noise.py
# Uses original [0,1] coordinates (coords_normalized_01)
# Convert time calculation to tensor for torch.fmod
time_tensor_02 = torch.tensor(time * 0.2, device=device, dtype=coords_normalized_01.dtype)
x_offset = torch.fmod(time_tensor_02, 1.0) * 2.0 # Match JS fract animation
shifted_x = torch.fmod(coords_normalized_01[:, :, :, 0] + x_offset, 1.0)
mask_output = shifted_x # Mask values will be [0,1]
elif shape_type == "radial_animated": # Ported from curl_noise.py (its "radial" shape)
# Uses coords_normalized_01 (range [0,1])
time_tensor = torch.tensor(time, device=device, dtype=coords_normalized_01.dtype)
center_x = 0.5 + 0.2 * torch.cos(time_tensor)
center_y = 0.5 + 0.2 * torch.sin(time_tensor)
# Calculate distance from the animated center using coords_normalized_01
y_diff = coords_normalized_01[:, :, :, 1] - center_y
x_diff = coords_normalized_01[:, :, :, 0] - center_x
dist_from_anim_center = torch.sqrt(x_diff**2 + y_diff**2) * 2.0 # Multiplied by 2 like in curl_noise
mask_output = torch.clamp(1.0 - dist_from_anim_center, 0.0, 1.0)
elif shape_type == "spiral": # Ported and enhanced from curl_noise.py, uses centered_coords
# centered_coords are in [-0.5, 0.5]
# theta and r are calculated from centered_coords
# angle = atan2(centered_coords_y, centered_coords_x) - already available as 'angle'
# r = norm(centered_coords) * 2.0 - center_dist is norm(centered_coords), so r = center_dist * 2.0
r_spiral = center_dist * 2.0 # center_dist is norm of coords in [-0.5,0.5], max ~0.707. So r_spiral max ~1.414
time_tensor = torch.tensor(time, device=device, dtype=coords_normalized_01.dtype)
theta_animated = angle + time_tensor # angle is already calculated from centered_coords
mask_output = torch.fmod((theta_animated / (2.0 * math.pi) + r_spiral), 1.0)
elif shape_type == "checkerboard": # Ported from curl_noise.py
# Uses coords_normalized_01 (range [0,1])
grid_size = 8.0 # From curl_noise.py
# Convert time expressions to tensors
time_tensor_gs_02 = torch.tensor(time * grid_size * 0.2, device=device, dtype=coords_normalized_01.dtype)
time_tensor_gs_01 = torch.tensor(time * grid_size * 0.1, device=device, dtype=coords_normalized_01.dtype)
x_offset = time_tensor_gs_02
y_offset = time_tensor_gs_01
# Apply offset and scale for grid calculation
x_grid_val = (coords_normalized_01[:, :, :, 0] + x_offset / grid_size) * grid_size
y_grid_val = (coords_normalized_01[:, :, :, 1] + y_offset / grid_size) * grid_size
# Floor and compute pattern. The *0.5 and fmod 1.0 results in 0 or 0.5 values, then combined for 0, 0.5, 1.0.
# To get a binary mask (0 or 1), we can fmod the sum by 2 and then check if it's < 1, or directly use fmod 1.0 from curl_noise.
x_grid_processed = torch.floor(x_grid_val) * 0.5
y_grid_processed = torch.floor(y_grid_val) * 0.5
mask_output = torch.fmod(x_grid_processed + y_grid_processed, 1.0)
elif shape_type == "spots": # Ported from curl_noise.py (more complex version)
# Uses coords_normalized_01 (same as coords_bhwc in CN)
mask_spots_cn = torch.zeros_like(coords_normalized_01[:, :, :, 0])
num_spots_cn = 10
time_tensor_cn_spots = torch.tensor(time, device=device, dtype=coords_normalized_01.dtype)
for i in range(num_spots_cn):
# Use ShaderParamsReader.random_val
rand_x_cn = ShaderParamsReader.random_val(coords_normalized_01, base_seed, i * 78)
rand_y_cn = ShaderParamsReader.random_val(coords_normalized_01, base_seed, i * 12)
size_cn_base = (ShaderParamsReader.random_val(coords_normalized_01, base_seed, i * 93) * 0.3 + 0.1)
angle_float_cn = time + float(i) # time is already a float or tensor
angle_tensor_cn = torch.tensor(angle_float_cn, device=device, dtype=coords_normalized_01.dtype)
spot_pos_x_cn = 0.5 + torch.cos(angle_tensor_cn) * 0.4 * rand_x_cn
spot_pos_y_cn = 0.5 + torch.sin(angle_tensor_cn) * 0.4 * rand_y_cn
size_anim_angle_cn = torch.tensor(time * 2.0 + float(i), device=device, dtype=coords_normalized_01.dtype)
size_cn_final = size_cn_base * (1.0 + 0.2 * torch.sin(size_anim_angle_cn))
dist_cn_spots = torch.sqrt((coords_normalized_01[:, :, :, 0] - spot_pos_x_cn)**2 + (coords_normalized_01[:, :, :, 1] - spot_pos_y_cn)**2)
# Avoid division by zero or very small size
spot_mask_cn_indiv = torch.clamp(1.0 - dist_cn_spots / (size_cn_final + 1e-8), 0.0, 1.0)
mask_spots_cn = torch.maximum(mask_spots_cn, spot_mask_cn_indiv)
mask_output = mask_spots_cn
elif shape_type == "hexgrid": # Adapted from original, using centered_coords
# This was complex. Let's simplify for [0,1] input.
# Using coords_normalized_01 directly for hexgrid based on curl_noise's hexgrid logic
hex_uv = coords_normalized_01 * 6.0 # Scale for hex grid density
# Convert time expressions to tensors before sin/cos
time_tensor_05 = torch.tensor(time * 0.5, device=device, dtype=coords_normalized_01.dtype)
time_tensor_03 = torch.tensor(time * 0.3, device=device, dtype=coords_normalized_01.dtype)
time_tensor = torch.tensor(time, device=device, dtype=coords_normalized_01.dtype)
hex_uv_anim = hex_uv.clone() # Avoid in-place modification if hex_uv is reused
hex_uv_anim[:, :, :, 0] += torch.sin(time_tensor_05) * 0.5
hex_uv_anim[:, :, :, 1] += torch.cos(time_tensor_03) * 0.5
r_vec = torch.tensor([1.0, 1.73], device=device, dtype=coords_normalized_01.dtype).reshape(1, 1, 1, 2)
h_vec = r_vec * 0.5
a_vec = torch.fmod(hex_uv_anim, r_vec) - h_vec
b_vec = torch.fmod(hex_uv_anim + h_vec, r_vec) - h_vec
dist_hex = torch.minimum(torch.norm(a_vec, dim=-1), torch.norm(b_vec, dim=-1))
cell_size = 0.3 + 0.1 * torch.sin(time_tensor)
# Use ShaderParamsReader.smoothstep
mask_output = ShaderParamsReader.smoothstep(cell_size + 0.05, cell_size - 0.05, dist_hex)
elif shape_type == "stripes": # Ported from curl_noise.py
# Uses coords_normalized_01 (range [0,1])
freq = 10.0
# Convert time expressions to tensors
time_tensor = torch.tensor(time, device=device, dtype=coords_normalized_01.dtype)
time_tensor_02 = torch.tensor(time * 0.2, device=device, dtype=coords_normalized_01.dtype)
angle_anim = 0.5 * torch.sin(time_tensor_02)
cos_a = torch.cos(angle_anim)
sin_a = torch.sin(angle_anim)
# Rotate coordinates directly using coords_normalized_01 to match curl_noise.py
# This rotates around the (0,0) origin of the [0,1] coordinate system.
rotated_x = coords_normalized_01[:, :, :, 0] * cos_a - coords_normalized_01[:, :, :, 1] * sin_a
# rotated_y is not strictly needed for this pattern as stripes are based on rotated_x
stripes_val = torch.sin(rotated_x * freq + time_tensor)
# Apply smoothstep to create distinct stripes
mask_output = ShaderParamsReader.smoothstep(0.0, 0.1, stripes_val) * ShaderParamsReader.smoothstep(0.0, -0.1, -stripes_val)
elif shape_type == "radial_gradient_static": # Renamed from "gradient"
# This is the original radial gradient from shader_params_reader.py
# center_dist max ~0.707 for centered_coords (which are coords_normalized_01 - 0.5).
# To have gradient from center to edge of original [0,1] box, we need to normalize center_dist.
# Max distance from center of a [0,1] box is sqrt(0.5^2+0.5^2) = ~0.707.
# So center_dist / 0.707 normalizes it roughly to [0,1] for points within the box.
# Or simpler, just use center_dist directly, it gives a gradient from 0 to ~0.7
mask_output = 1.0 - torch.clamp(center_dist / 0.5, 0.0, 1.0) # Soft radial gradient, 0.5 radius
elif shape_type == "gradient": # Ported from curl_noise.py (animated directional gradient)
# Uses coords_normalized_01 (equivalent to coords_bhwc in curl_noise.py)
time_tensor_02_grad = torch.tensor(time * 0.2, device=device, dtype=coords_normalized_01.dtype)
angle_grad = time_tensor_02_grad
dir_x_grad = torch.cos(angle_grad)
dir_y_grad = torch.sin(angle_grad)
# Project centered coordinates onto the direction vector
# coords_normalized_01 are [0,1], so (coords_normalized_01 - 0.5) makes them [-0.5, 0.5]
proj_grad = (coords_normalized_01[:, :, :, 0] - 0.5) * dir_x_grad + \
(coords_normalized_01[:, :, :, 1] - 0.5) * dir_y_grad + 0.5
mask_output = proj_grad # Result is roughly in [0,1] range
elif shape_type == "vignette": # Ported from curl_noise.py
# Uses coords_normalized_01 (range [0,1])
# Convert time expressions to tensors
time_tensor_03 = torch.tensor(time * 0.3, device=device, dtype=coords_normalized_01.dtype)
time_tensor_04 = torch.tensor(time * 0.4, device=device, dtype=coords_normalized_01.dtype)
time_tensor_05 = torch.tensor(time * 0.5, device=device, dtype=coords_normalized_01.dtype)
center_x_anim = 0.5 + 0.2 * torch.sin(time_tensor_03)
center_y_anim = 0.5 + 0.2 * torch.cos(time_tensor_04)
# Calculate distance from animated center using coords_normalized_01
dist_from_anim_center_x = coords_normalized_01[:, :, :, 0] - center_x_anim
dist_from_anim_center_y = coords_normalized_01[:, :, :, 1] - center_y_anim
dist_vignette = torch.sqrt(dist_from_anim_center_x**2 + dist_from_anim_center_y**2)
radius_anim = 0.6 + 0.2 * torch.sin(time_tensor_05)
smoothness = 0.3 # As in curl_noise.py
mask_output = 1.0 - ShaderParamsReader.smoothstep(radius_anim - smoothness, radius_anim, dist_vignette)
elif shape_type == "cross": # Ported from curl_noise.py
# Uses coords_normalized_01 (range [0,1]), but calculations are around center (0.5,0.5)
# Convert time expressions to tensors
time_tensor = torch.tensor(time, device=device, dtype=coords_normalized_01.dtype)
time_tensor_02 = torch.tensor(time * 0.2, device=device, dtype=coords_normalized_01.dtype)
thickness_anim = 0.1 + 0.05 * torch.sin(time_tensor)
rotation_anim = time_tensor_02
cos_r = torch.cos(rotation_anim)
sin_r = torch.sin(rotation_anim)
# Use centered_coords for rotation calculation
# centered_coords are already coords_normalized_01 - 0.5
rotated_x = centered_coords[:, :, :, 0] * cos_r - centered_coords[:, :, :, 1] * sin_r
rotated_y = centered_coords[:, :, :, 0] * sin_r + centered_coords[:, :, :, 1] * cos_r
# Shift back to [0,1]-like domain for comparison if needed, or compare in [-0.5,0.5] domain
# The original curl_noise compared rotated_x/y against 0.5 after adding 0.5.
# Here, rotated_x/y are already centered around 0. So, we compare against 0.
# Horizontal bar (rotated)
h_bar = ShaderParamsReader.smoothstep(0.0 - thickness_anim, 0.0 - thickness_anim + 0.02, rotated_y) * \
ShaderParamsReader.smoothstep(0.0 + thickness_anim, 0.0 + thickness_anim - 0.02, rotated_y)
# Vertical bar (rotated)
v_bar = ShaderParamsReader.smoothstep(0.0 - thickness_anim, 0.0 - thickness_anim + 0.02, rotated_x) * \
ShaderParamsReader.smoothstep(0.0 + thickness_anim, 0.0 + thickness_anim - 0.02, rotated_x)
mask_output = torch.maximum(h_bar, v_bar)
elif shape_type == "triangles": # Adapted from original, using centered_coords
# Re-evaluate scaling for centered_coords.
# Let's use coords_normalized_01 for a direct port attempt of a triangle grid like curl_noise
t_tensor = torch.tensor(time * 0.2, device=device, dtype=coords_normalized_01.dtype)
t_sin_arg = t_tensor
t_cos_arg = torch.tensor(time * 0.7, device=device, dtype=coords_normalized_01.dtype)
t_border_arg = torch.tensor(time * 1.5, device=device, dtype=coords_normalized_01.dtype)
scale_factor = 5.0
uv_tri = coords_normalized_01 * scale_factor
uv_tri_anim = uv_tri.clone()
uv_tri_anim[:, :, :, 0] += torch.sin(t_sin_arg) * 0.5
uv_tri_anim[:, :, :, 1] += torch.cos(t_cos_arg) * 0.5
gv = torch.fmod(uv_tri_anim, 1.0) - 0.5 # gv is now in [-0.5, 0.5]
d1 = torch.abs(gv[:, :, :, 0] + gv[:, :, :, 1])
d2 = torch.abs(gv[:, :, :, 0] - gv[:, :, :, 1])
d3 = torch.abs(gv[:, :, :, 0]) * 0.866 + torch.abs(gv[:, :, :, 1]) * 0.5 # Approx dist for equilateral
d_tri = torch.minimum(torch.minimum(d1, d2), d3) * 0.7
border_width = 0.05 + 0.03 * torch.sin(t_border_arg)
mask_output = ShaderParamsReader.smoothstep(border_width, border_width - 0.02, d_tri)
elif shape_type == "concentric": # Ported from curl_noise.py
# Uses coords_normalized_01 (range [0,1]) for calculating distance from an animated center.
# Convert time expressions to tensors
time_tensor_03 = torch.tensor(time * 0.3, device=device, dtype=coords_normalized_01.dtype)
time_tensor_04 = torch.tensor(time * 0.4, device=device, dtype=coords_normalized_01.dtype)
time_tensor_01 = torch.tensor(time * 0.1, device=device, dtype=coords_normalized_01.dtype)
time_tensor_05 = torch.tensor(time * 0.5, device=device, dtype=coords_normalized_01.dtype)
center_x_anim = 0.5 + 0.2 * torch.sin(time_tensor_03)
center_y_anim = 0.5 + 0.2 * torch.cos(time_tensor_04)
dist_from_center = torch.sqrt((coords_normalized_01[:, :, :, 0] - center_x_anim)**2 +
(coords_normalized_01[:, :, :, 1] - center_y_anim)**2)
freq_anim = 10.0 + 5.0 * torch.sin(time_tensor_01)
phase_anim = time_tensor_05
rings_val = torch.sin(dist_from_center * freq_anim + phase_anim)
mask_output = ShaderParamsReader.smoothstep(0.0, 0.1, rings_val) * ShaderParamsReader.smoothstep(0.0, -0.1, -rings_val)
elif shape_type == "rays": # Ported from curl_noise.py
# Uses coords_normalized_01 (range [0,1]) for calculating angle and dist from an animated center.
# Convert time expressions to tensors
time_tensor_03 = torch.tensor(time * 0.3, device=device, dtype=coords_normalized_01.dtype)
time_tensor_04 = torch.tensor(time * 0.4, device=device, dtype=coords_normalized_01.dtype)
time_tensor_05 = torch.tensor(time * 0.5, device=device, dtype=coords_normalized_01.dtype)
center_x_anim = 0.5 + 0.1 * torch.sin(time_tensor_03)
center_y_anim = 0.5 + 0.1 * torch.cos(time_tensor_04)
to_center_x = coords_normalized_01[:, :, :, 0] - center_x_anim
to_center_y = coords_normalized_01[:, :, :, 1] - center_y_anim
angle_rays = torch.atan2(to_center_y, to_center_x)
freq_rays = 8.0
phase_rays = time_tensor_05
rays_val = torch.sin(angle_rays * freq_rays + phase_rays)
dist_rays = torch.sqrt(to_center_x**2 + to_center_y**2)
falloff = 1.0 - ShaderParamsReader.smoothstep(0.0, 0.8, dist_rays)
mask_output = ShaderParamsReader.smoothstep(0.0, 0.3, rays_val) * falloff
elif shape_type == "zigzag": # Ported from curl_noise.py
# Uses coords_normalized_01 (range [0,1]) and centers for rotation.
freq_zigzag = 10.0
# Convert time expressions to tensors
time_tensor = torch.tensor(time, device=device, dtype=coords_normalized_01.dtype)
time_tensor_02 = torch.tensor(time * 0.2, device=device, dtype=coords_normalized_01.dtype)
time_tensor_05 = torch.tensor(time * 0.5, device=device, dtype=coords_normalized_01.dtype)
time_tensor_03 = torch.tensor(time * 0.3, device=device, dtype=coords_normalized_01.dtype)
angle_zigzag = 0.5 * torch.sin(time_tensor_02)
cos_a = torch.cos(angle_zigzag)
sin_a = torch.sin(angle_zigzag)
# Rotate coords_normalized_01 directly to match curl_noise.py (rotation around 0,0 of the [0,1] grid)
rotated_x_norm = coords_normalized_01[:, :, :, 0] * cos_a - coords_normalized_01[:, :, :, 1] * sin_a
rotated_y_norm = coords_normalized_01[:, :, :, 0] * sin_a + coords_normalized_01[:, :, :, 1] * cos_a
# The original curl_noise performed fmod(rotated_coord * freq - time_offset, 1.0).
# Then 2.0 * fmod_result - 1.0 to bring to [-1,1], then abs for [0,1].
# Since rotated_x_norm and rotated_y_norm are in a range determined by the rotation of [0,1] coordinates,
# multiplying by freq_zigzag will expand this range before fmod.
zigzag1 = torch.abs(2.0 * torch.fmod(rotated_x_norm * freq_zigzag - time_tensor_05, 1.0) - 1.0)
zigzag2 = torch.abs(2.0 * torch.fmod(rotated_y_norm * freq_zigzag + time_tensor_03, 1.0) - 1.0)
zigzag_combined = torch.minimum(zigzag1, zigzag2)
thickness_anim = 0.3 + 0.1 * torch.sin(time_tensor)
# torch.heaviside(input, values) outputs values where input > 0, and 0 where input < 0.
# For input == 0, it outputs values[0] if it's a tensor, or just values if scalar.
# A common way to get a step is (input > threshold).float()
# curl_noise's step(edge, x) is (x >= edge).float()
# So, heaviside(zigzag - thickness, torch.tensor(0.5)) is similar to (zigzag - thickness >= 0).float()
# which is (zigzag >= thickness).float()
mask_output = (zigzag_combined >= thickness_anim).float() # More direct step function
elif shape_type == "gradient_x": # Ported from curl_noise.py
# Uses coords_normalized_01 (range [0,1])
mask_output = coords_normalized_01[:, :, :, 0]
elif shape_type == "gradient_y": # Ported from curl_noise.py
# Uses coords_normalized_01 (range [0,1])
mask_output = coords_normalized_01[:, :, :, 1]
elif shape_type == "stars": # Ported from curl_noise.py
# Uses coords_normalized_01 (same as coords_bhwc in CN)
mask_stars_cn = torch.zeros_like(coords_normalized_01[:, :, :, 0])
num_stars_cn = 20
time_tensor_cn_stars = torch.tensor(time, device=device, dtype=coords_normalized_01.dtype)
time_tensor_01_cn = torch.tensor(time * 0.1, device=device, dtype=coords_normalized_01.dtype)
time_tensor_015_cn = torch.tensor(time * 0.15, device=device, dtype=coords_normalized_01.dtype)
for i in range(num_stars_cn):
# Use ShaderParamsReader.random_val
rand_x_stars = ShaderParamsReader.random_val(coords_normalized_01, base_seed, i * 78 + 10)
rand_y_stars = ShaderParamsReader.random_val(coords_normalized_01, base_seed, i * 12 + 20)
time_sin_arg_stars = torch.tensor(float(i), device=device, dtype=coords_normalized_01.dtype) + time_tensor_01_cn
time_cos_arg_stars = torch.tensor(float(i) * 1.5, device=device, dtype=coords_normalized_01.dtype) + time_tensor_015_cn
star_pos_x_cn = torch.fmod(rand_x_stars + 0.05 * torch.sin(time_sin_arg_stars), 1.0)
star_pos_y_cn = torch.fmod(rand_y_stars + 0.05 * torch.cos(time_cos_arg_stars), 1.0)
brightness_arg_stars = torch.tensor(float(i), device=device, dtype=coords_normalized_01.dtype) + time_tensor_cn_stars * (0.5 + rand_x_stars * 0.5)
brightness_cn = 0.5 + 0.5 * torch.sin(brightness_arg_stars)
size_stars_cn = 0.01 + 0.015 * rand_y_stars * brightness_cn
dist_stars_cn = torch.sqrt((coords_normalized_01[:, :, :, 0] - star_pos_x_cn)**2 + (coords_normalized_01[:, :, :, 1] - star_pos_y_cn)**2)
# Use ShaderParamsReader.smoothstep for soft stars
star_mask_cn_indiv = ShaderParamsReader.smoothstep(size_stars_cn, size_stars_cn * 0.5, dist_stars_cn) * brightness_cn
mask_stars_cn = torch.maximum(mask_stars_cn, star_mask_cn_indiv)
mask_output = mask_stars_cn
else: # Default for unknown or "none"
if shape_type not in ["none", "0", 0]: # Only print warning for actual unknown types
print(f"Unknown shape type: {shape_type}, using default (full mask)")
# Default is full mask (all ones)
mask_output = torch.ones((batch, height, width), device=device)
# Ensure mask_output is [B, H, W, 1]
if len(mask_output.shape) == 3: # If it's [B, H, W]
mask_output = mask_output.unsqueeze(-1)
elif len(mask_output.shape) == 4 and mask_output.shape[-1] != 1: # If it's [B,H,W,C] C!=1
print(f"Warning: Shape mask generated with {mask_output.shape[-1]} channels. Taking first channel.")
mask_output = mask_output[..., 0:1]
return mask_output
# Legacy functions for backward compatibility
def get_shader_params():
"""Legacy function that calls the new class method"""
return ShaderParamsReader.get_shader_params()
def build_shader_params(base, seed, shader_type, shape_type, color_scheme, noise_scale,
octaves, warp_strength, shape_mask_strength, phase_shift,
color_intensity, use_temporal_coherence, fast_high_channel_noise):
"""
Write a node's shader inputs over `base`, in every spelling the generators read.
`base` is the saved params file; the node's own widgets win over it.
"""
params = dict(base)
params["shader_type"] = shader_type
params["shaderType"] = shader_type
params["shape_type"] = shape_type
params["shaderShapeType"] = shape_type
params["colorScheme"] = color_scheme
params["color_scheme"] = color_scheme
params["scale"] = noise_scale
params["shaderScale"] = noise_scale
params["octaves"] = float(octaves)
params["shaderOctaves"] = float(octaves)
params["warp_strength"] = warp_strength
params["shaderWarpStrength"] = warp_strength
params["shapemaskstrength"] = shape_mask_strength
params["shaderShapeStrength"] = shape_mask_strength
params["shapeMaskStrength"] = shape_mask_strength
params["shape_mask_strength"] = shape_mask_strength
params["shape_strength"] = shape_mask_strength
params["phase_shift"] = phase_shift
params["shaderPhaseShift"] = phase_shift
params["intensity"] = color_intensity
params["shaderColorIntensity"] = color_intensity
params["time"] = params.get("time", 0.0)
params["base_seed"] = seed
params["useTemporalCoherence"] = use_temporal_coherence
params["temporal_coherence"] = use_temporal_coherence
params["fast_high_channel_noise"] = fast_high_channel_noise
params["visualization_type"] = params.get("visualization_type", 3)
# Clamp octaves, seeds and enum values before they reach noise generation.
params = ShaderParamsReader.validate_and_sanitize_params(params)
# Sanitising truncates octaves to an integer; the standard pipeline
# interpolates between integer renders, so keep the requested value.
params["octaves"] = float(octaves)
return params
def test_params():
"""Test function to check if parameters are loading correctly"""
params = ShaderParamsReader.get_shader_params()
print(f"TEST: Current shader parameters: {params}")
return params
def generate_noise_tensor(shader_params, height, width, batch_size=1, device="cuda", seed=0, target_channels=None):
"""
Legacy function that uses the new class to generate noise
Args:
shader_params: Dictionary containing shader parameters
height: Height of the tensor
width: Width of the tensor
batch_size: Number of images in the batch
device: Device to create tensor on
seed: Random seed for deterministic noise generation
target_channels: Number of output channels (optional, default is 4)
Returns:
Noise tensor with shape [batch_size, channels, height, width]
where channels is determined by target_channels (default: 4)
"""
# Import ShaderToTensor for direct shader noise generation
from .shader_to_tensor import ShaderToTensor
# Make a copy of shader_params to avoid modifying the original
shader_params = shader_params.copy()
# Add target_channels to shader_params if provided
if target_channels is not None:
shader_params["target_channels"] = target_channels
print(f"Using target_channels={target_channels} in default generator")
# Check if we're using specialized shader types which are handled separately in the sampler
shader_type = shader_params.get("shader_type", "tensor_field")
if shader_type == "cellular":
print("Note: Cellular shader type detected in generate_noise_tensor, but this will be handled by the specialized cellular noise generator")
elif shader_type == "fractal":
print("Note: Fractal shader type detected in generate_noise_tensor, but this will be handled by the specialized fractal noise generator")
elif shader_type == "perlin":
print("Note: Perlin shader type detected in generate_noise_tensor, but this will be handled by the specialized perlin noise generator")
elif shader_type == "waves":
print("Note: Waves shader type detected in generate_noise_tensor, but this will be handled by the specialized waves noise generator")
elif shader_type == "gaussian":
print("Note: Gaussian shader type detected in generate_noise_tensor, but this will be handled by the specialized gaussian noise generator")
elif shader_type == "tensor_field":
print("Note: Tensor field shader type detected in generate_noise_tensor, but this will be handled by the specialized tensor field generator")
elif shader_type == "heterogeneous_fbm":
print("Note: Heterogeneous FBM shader type detected in generate_noise_tensor, but this will be handled by the specialized generator")
elif shader_type == "interference_patterns":
print("Note: Interference patterns shader type detected in generate_noise_tensor, but this will be handled by the specialized generator")
# Use the seed for deterministic generation
torch.manual_seed(seed)
# Extract parameters from shader_params
viz_type = shader_params.get("visualization_type", 3)
scale = shader_params.get("scale", 1.0)
warp_strength = shader_params.get("warp_strength", 0.5)
phase_shift = shader_params.get("phase_shift", 0.0)
time = shader_params.get("time", 0.0)
octaves = shader_params.get("octaves", 3.0)
shape_type = shader_params.get("shape_type", "none")
shape_mask_strength = shader_params.get("shapemaskstrength", 1.0)
# Generate shader noise directly using ShaderToTensor class
# This avoids using random noise as a starting point
shader_noise = ShaderToTensor.shader_noise_to_tensor(
batch_size=batch_size,
height=height,
width=width,
shader_type=shader_type,
visualization_type=viz_type,
scale=scale,
phase_shift=phase_shift,
warp_strength=warp_strength,
time=time,
device=device,
seed=seed,
octaves=octaves,
shape_type=shape_type,
shape_mask_strength=shape_mask_strength,
shader_params=shader_params
)
# Reset random seed state
torch.manual_seed(torch.seed())
# Apply color scheme transformation
color_scheme = shader_params.get("colorScheme", "none")
color_intensity = shader_params.get("shaderColorIntensity", 0.8)
if color_scheme != "none" and color_intensity > 0:
print(f"Applying color scheme: {color_scheme} with intensity: {color_intensity}")
# Track channel stats before applying color
num_channels = shader_noise.shape[1]
channel_means_before = [shader_noise[:, i].mean().item() for i in range(num_channels)]
channel_stds_before = [shader_noise[:, i].std().item() for i in range(num_channels)]
print(f"Channel means before color: {[f'{m:.4f}' for m in channel_means_before]}")
# Apply color scheme
colored_noise = ShaderParamsReader.apply_color_scheme(shader_noise, shader_params)
# Track channel stats after applying color
channel_means_after = [colored_noise[:, i].mean().item() for i in range(num_channels)]
channel_stds_after = [colored_noise[:, i].std().item() for i in range(num_channels)]
print(f"Channel means after color: {[f'{m:.4f}' for m in channel_means_after]}")
# Normalize each channel separately while preserving mean differences
# This ensures the color impact remains visible
# We only normalize the standard deviation to keep it at ~1.0
normalized_colored_noise = torch.zeros_like(colored_noise)
for i in range(num_channels):
# Only normalize the standard deviation while keeping the mean offset
channel = colored_noise[:, i:i+1]
normalized_colored_noise[:, i:i+1] = (channel - channel.mean()) / (channel.std() + 1e-8) + channel_means_after[i]
# Verify the normalization preserved color differences
final_means = [normalized_colored_noise[:, i].mean().item() for i in range(num_channels)]
final_stds = [normalized_colored_noise[:, i].std().item() for i in range(num_channels)]
print(f"Final means after normalization: {[f'{m:.4f}' for m in final_means]}")
print(f"Final stds after normalization: {[f'{s:.4f}' for s in final_stds]}")
return normalized_colored_noise
return shader_noise
# Example of use in sampling process:
#
# 1. Load shader parameters
# shader_params = ShaderParamsReader.get_shader_params()
#
# 2. Generate initial noise for sampling
# noise = comfy.sample.prepare_noise(latent_samples, seed, batch_inds)
#
# 3. Apply shader transformation to noise (implements Sα(N))
# modified_noise = ShaderParamsReader.apply_shader_to_noise(noise, shader_params)
#
# 4. Use modified noise in sampling (implements Kβ(t))
# samples = comfy.sample.sample(
# model=model,
# noise=modified_noise,
# # ... other parameters
# disable_noise=True, # Using our pre-modified noise
# )
Reference in New Issue View Git Blame Copy Permalink