AEmotionStudio-ComfyUI-ShaderNoiseKSampler/shader_params_reader.py

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
    @staticmethod
    def valid_shader_types():
        """
        The names the generator registry can build, plus the spellings mapped
        onto them below. A literal set here drifted: it never gained
        temporal_coherent, so a legacy-mode run asking for it was quietly
        handed tensor_field instead.
        """
        from .shaders.registry import list_shaders

        return set(list_shaders()) | {"curl", "interference_patterns"}

    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 == "interference_patterns": st = "interference"

                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"
                        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"
                            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
# )