Source code for noisemaker.generators

"""Noise generation interface for Noisemaker"""

from __future__ import annotations

import tempfile
from functools import partial
from typing import TYPE_CHECKING, Any, Callable

import tensorflow as tf

import noisemaker.ai as ai
import noisemaker.oklab as oklab
import noisemaker.simplex as simplex
import noisemaker.util as util
import noisemaker.value as value
from noisemaker.constants import (
    ColorSpace,
    InterpolationType,
    OctaveBlending,
    ValueDistribution,
    ValueMask,
)

if TYPE_CHECKING:
    from noisemaker.composer import Preset


[docs] def basic( freq: int | list[int], shape: list[int], ridges: bool = False, sin: float = 0.0, spline_order: InterpolationType = InterpolationType.bicubic, distrib: ValueDistribution = ValueDistribution.simplex, corners: bool = False, mask: ValueMask | None = None, mask_inverse: bool = False, mask_static: bool = False, lattice_drift: float = 0.0, color_space: ColorSpace = ColorSpace.hsv, hue_range: float = 0.125, hue_rotation: float | None = None, saturation: float = 1.0, hue_distrib: ValueDistribution | None = None, brightness_distrib: ValueDistribution | None = None, brightness_freq: int | list[int] | None = None, saturation_distrib: ValueDistribution | None = None, speed: float = 1.0, time: float = 0.0, octave_effects: list[Callable] | None = None, octave: int = 1, ) -> tf.Tensor: """ Generate a single layer of scaled noise. .. noisemaker-live:: :generator: basic :seed: 12345 :width: 512 :height: 256 :lazy: Args: freq: Base noise frequency. Int, or list of ints for each spatial dimension shape: Shape of noise. For 2D noise, this is [height, width, channels] ridges: "Crease" at midpoint values: (1 - abs(n * 2 - 1)) sin: Apply sin function to noise basis spline_order: Spline point count. 0=Constant, 1=Linear, 2=Cosine, 3=Bicubic distrib: Type of noise distribution. See :class:`ValueDistribution` enum corners: If True, pin values to corners instead of image center mask: Optional mask to apply mask_inverse: Invert the mask mask_static: If True, don't animate the mask lattice_drift: Push away from underlying lattice color_space: Color space to use (HSV, RGB, etc) hue_range: HSV hue range hue_rotation: HSV hue bias saturation: HSV saturation hue_distrib: Override ValueDistribution for hue saturation_distrib: Override ValueDistribution for saturation brightness_distrib: Override ValueDistribution for brightness brightness_freq: Override frequency for brightness speed: Displacement range for Z/W axis (simplex and periodic only) time: Time argument for Z/W axis (simplex and periodic only) octave_effects: Effects to apply per octave octave: Current octave number Returns: Generated noise tensor """ freq_list: list[int] if isinstance(freq, int): freq_list = value.freq_for_shape(freq, shape) else: freq_list = freq color_space = value.coerce_enum(color_space, ColorSpace) common_value_params: dict[str, Any] = { "corners": corners, "mask": mask, "mask_inverse": mask_inverse, "mask_static": mask_static, "speed": speed, "spline_order": spline_order, "time": time, } tensor = value.values( freq=freq_list, shape=shape, distrib=distrib, corners=corners, mask=mask, mask_inverse=mask_inverse, mask_static=mask_static, speed=speed, spline_order=spline_order, time=time, ) if lattice_drift: tensor = value.refract( tensor, shape, time=time, speed=speed, displacement=lattice_drift / min(freq_list[0], freq_list[1]), warp_freq=freq_list, spline_order=spline_order, signed_range=False, ) if octave_effects is not None: for effect_or_preset in octave_effects: tensor = _apply_octave_effect_or_preset(effect_or_preset, tensor, shape, time, speed, octave) # Preserve alpha channel for color space conversions alpha = None if shape[2] == 4: alpha = tensor[:, :, 3] tensor = tf.stack([tensor[:, :, 0], tensor[:, :, 1], tensor[:, :, 2]], 2) elif shape[2] == 2: alpha = tensor[:, :, 1] tensor = tf.stack([tensor[:, :, 0]], 2) original_color_space = color_space if color_space == ColorSpace.oklab: L = tensor[:, :, 0] a = tensor[:, :, 1] * -0.509 + 0.276 b = tensor[:, :, 2] * -0.509 + 0.198 tensor = value.clamp01(oklab.oklab_to_rgb(tf.stack([L, a, b], 2))) color_space = ColorSpace.rgb if color_space == ColorSpace.rgb: tensor = tf.image.rgb_to_hsv([tensor])[0] color_space = ColorSpace.hsv if color_space == ColorSpace.hsv: # Use 1 channel for per-channel noise generation, if any common_value_params["shape"] = [shape[0], shape[1], 1] # tweak hue if hue_distrib: h = tf.squeeze(value.values(freq=freq_list, distrib=hue_distrib, **common_value_params)) else: if original_color_space == ColorSpace.hsv: if hue_rotation is None: hue_rotation = simplex.random(time=time, speed=speed) else: # Avoid hard edges on color models that don't wrap hue from 1 to 0 naturally hue_range = 1.0 hue_rotation = 0.0 h = (tensor[:, :, 0] * hue_range + hue_rotation) % 1.0 # tweak saturation if saturation_distrib: s = tf.squeeze(value.values(freq=freq_list, distrib=saturation_distrib, **common_value_params)) else: s = tensor[:, :, 1] s *= saturation # tweak brightness if brightness_distrib or brightness_freq: if isinstance(brightness_freq, int): brightness_freq = value.freq_for_shape(brightness_freq, shape) v = tf.squeeze(value.values(freq=brightness_freq or freq_list, distrib=brightness_distrib or ValueDistribution.simplex, **common_value_params)) else: v = tensor[:, :, 2] if ridges and spline_order: # ridges don't work with spline_order == 0 v = value.ridge(v) if sin: v = value.normalize(tf.sin(sin * v)) tensor = tf.image.hsv_to_rgb([tf.stack([h, s, v], 2)])[0] if color_space == ColorSpace.grayscale: if ridges and spline_order: # ridges don't work with spline_order == 0 tensor = value.ridge(tensor) if sin: tensor = tf.sin(sin * tensor) # re-insert the alpha channel if shape[2] == 4: tensor = tf.stack([tensor[:, :, 0], tensor[:, :, 1], tensor[:, :, 2], alpha], 2) elif shape[2] == 2: tensor = tf.stack([tensor[:, :, 0], alpha], 2) return tensor
[docs] def multires( preset, seed: int, freq: int | list[int] = 3, shape: list[int] | None = None, octaves: int = 1, ridges: bool = False, sin: float = 0.0, spline_order: InterpolationType = InterpolationType.bicubic, distrib: ValueDistribution = ValueDistribution.simplex, corners: bool = False, mask: ValueMask | None = None, mask_inverse: bool = False, mask_static: bool = False, lattice_drift: float = 0.0, with_supersample: bool = False, color_space: ColorSpace = ColorSpace.hsv, hue_range: float = 0.125, hue_rotation: float | None = None, saturation: float = 1.0, hue_distrib: ValueDistribution | None = None, saturation_distrib: ValueDistribution | None = None, brightness_distrib: ValueDistribution | None = None, brightness_freq: int | list[int] | None = None, octave_blending: OctaveBlending = OctaveBlending.falloff, octave_effects: list[Callable] | None = None, post_effects: list[Callable] | None = None, with_alpha: bool = False, final_effects: list[Callable] | None = None, with_upscale: bool = False, with_fxaa: bool = False, time: float = 0.0, speed: float = 1.0, tensor: tf.Tensor | None = None, ) -> tf.Tensor: """ Generate multi-resolution value noise. For each octave: freq increases, amplitude decreases. .. noisemaker-live:: :generator: multires :seed: 42 :octaves: 5 :ridges: True :width: 512 :height: 256 :lazy: Args: preset: The Preset object being rendered seed: The generation seed to use freq: Bottom layer frequency. Int, or list of ints for each spatial dimension shape: Shape of noise. For 2D noise, this is [height, width, channels] octaves: Octave count. Number of multi-res layers. Typically 1-8 ridges: Per-octave "crease" at midpoint values: (1 - abs(n * 2 - 1)) sin: Apply sin function to noise basis spline_order: Spline point count. 0=Constant, 1=Linear, 2=Cosine, 3=Bicubic distrib: Type of noise distribution. See :class:`ValueDistribution` enum corners: If True, pin values to corners instead of image center mask: Optional mask to apply mask_inverse: Invert the mask mask_static: If True, don't animate the mask lattice_drift: Push away from underlying lattice with_supersample: Use x2 supersampling color_space: Color space to use (HSV, RGB, etc) hue_range: HSV hue range hue_rotation: HSV hue bias saturation: HSV saturation hue_distrib: Override ValueDistribution for HSV hue saturation_distrib: Override ValueDistribution for HSV saturation brightness_distrib: Override ValueDistribution for HSV brightness brightness_freq: Override frequency for HSV brightness octave_blending: Method for flattening octave values octave_effects: A list of composer lambdas to invoke per-octave post_effects: A list of composer lambdas to invoke after flattening layers with_alpha: Include alpha channel final_effects: A list of composer lambdas to invoke after everything else with_upscale: AI: x2 upscale final results with_fxaa: Apply FXAA to results speed: Displacement range for Z/W axis (simplex and periodic only) time: Time argument for Z/W axis (simplex and periodic only) tensor: Optional input tensor to start with Returns: Generated multi-resolution noise tensor """ if seed: value.set_seed(seed) # Normalize input color_space = value.coerce_enum(color_space, ColorSpace) octave_blending = value.coerce_enum(octave_blending, OctaveBlending) # At this point, shape must be defined if shape is None: raise ValueError("shape must be provided") original_shape = shape.copy() if shape[-1] is None: shape = util.shape_from_params(shape[1], shape[0], color_space, with_alpha) if isinstance(freq, int): freq = value.freq_for_shape(freq, shape) if with_supersample: shape[0] *= 2 shape[1] *= 2 if octave_blending == OctaveBlending.alpha and shape[2] in (1, 3): # Make sure there's an alpha channel shape[2] += 1 if tensor is None: tensor = tf.zeros(shape, dtype=tf.float32) for octave in range(1, octaves + 1): multiplier = 2**octave freq_list = freq if isinstance(freq, list) else [freq, freq] base_freq = [int(f * 0.5 * multiplier) for f in freq_list] if all(base_freq[i] > shape[i] for i in range(len(base_freq))): break layer = basic( base_freq, shape, ridges=ridges, sin=sin, spline_order=spline_order, corners=corners, distrib=distrib, mask=mask, mask_inverse=mask_inverse, mask_static=mask_static, lattice_drift=lattice_drift, color_space=color_space, hue_range=hue_range, hue_rotation=hue_rotation, saturation=saturation, hue_distrib=hue_distrib, brightness_distrib=brightness_distrib, brightness_freq=brightness_freq, saturation_distrib=saturation_distrib, octave_effects=octave_effects, octave=octave, time=time, speed=speed, ) if octave_blending == OctaveBlending.reduce_max: tensor = tf.maximum(tensor, layer) elif octave_blending == OctaveBlending.alpha: a = tf.expand_dims(layer[:, :, -1], -1) tensor = (tensor * (1.0 - a)) + layer * a else: # falloff tensor += layer / multiplier # If the original shape did not include an alpha channel, reduce masked values to 0 (black) if octave_blending == OctaveBlending.alpha and original_shape[2] in (1, 3): a = tensor[:, :, -1] if original_shape[2] == 1: tensor = tf.expand_dims(tensor[:, :, 0] * a, -1) elif original_shape[2] == 3: tensor = tf.stack([tensor[:, :, 0], tensor[:, :, 1], tensor[:, :, 2]], 2) * tf.expand_dims(a, -1) shape = original_shape else: if octave_effects is not None: for effect_or_preset in octave_effects: tensor = _apply_octave_effect_or_preset(effect_or_preset, tensor, shape, time, speed, 1) tensor = value.normalize(tensor) final = [] if tensor.shape != shape: value.resample(tensor, shape) if post_effects is not None: for effect_or_preset in post_effects: tensor, f = _apply_post_effect_or_preset(effect_or_preset, tensor, shape, time, speed) final += f for effect_or_preset in final + (final_effects or []): tensor = _apply_final_effect_or_preset(effect_or_preset, tensor, shape, time, speed) tensor = value.normalize(tensor) if with_fxaa: tensor = value.fxaa(tensor, shape) if with_supersample: tensor = value.proportional_downsample(tensor, shape, original_shape) if with_upscale: with tempfile.TemporaryDirectory() as tmp: tmp_path = f"{tmp}/temp.png" util.save(tensor, tmp_path) try: tensor = ai.x4_upscale(tmp_path) except Exception as e: util.logger.error(f"preset.upscale() failed: {e}\nSeed: {seed}") return tensor
def _apply_octave_effect_or_preset( effect_or_preset: Callable | Preset, tensor: tf.Tensor, shape: list[int], time: float, speed: float, octave: int ) -> tf.Tensor: """ Helper function to either invoke an octave effect or unroll a preset. Args: effect_or_preset: Effect function or Preset to apply tensor: Input tensor shape: Tensor shape time: Time parameter speed: Speed parameter octave: Current octave number Returns: Processed tensor """ if callable(effect_or_preset): # Check if this is a functools.partial object with displacement keyword if hasattr(effect_or_preset, "keywords") and hasattr(effect_or_preset, "func"): if "displacement" in effect_or_preset.keywords: kwargs = dict(effect_or_preset.keywords) kwargs["displacement"] /= 2**octave effect_or_preset = partial(effect_or_preset.func, **kwargs) return effect_or_preset(tensor=tensor, shape=shape, time=time, speed=speed) else: # Is a Preset. Unroll me. for e_or_p in effect_or_preset.octave_effects: tensor = _apply_octave_effect_or_preset(e_or_p, tensor, shape, time, speed, octave) return tensor def _apply_post_effect_or_preset(effect_or_preset: Callable | Preset, tensor: tf.Tensor, shape: list[int], time: float, speed: float) -> tuple[tf.Tensor, list]: """ Helper function to either invoke a post effect or unroll a preset. Args: effect_or_preset: Effect function or Preset to apply tensor: Input tensor shape: Tensor shape time: Time parameter speed: Speed parameter Returns: Tuple of (processed tensor, list of final effects to apply later) """ if callable(effect_or_preset): return effect_or_preset(tensor=tensor, shape=shape, time=time, speed=speed), [] else: # Is a Preset. Unroll me. final = [] # Post effects may also define "final" effects. Collect them and return them so we # can tack them on at the end after everything is said and done final += effect_or_preset.final_effects for e_or_p in effect_or_preset.post_effects: tensor, f = _apply_post_effect_or_preset(e_or_p, tensor, shape, time, speed) final += f return tensor, final def _apply_final_effect_or_preset(effect_or_preset: Callable | Preset, tensor: tf.Tensor, shape: list[int], time: float, speed: float) -> tf.Tensor: """ Helper function to either invoke a final effect or unroll a preset. Args: effect_or_preset: Effect function or Preset to apply tensor: Input tensor shape: Tensor shape time: Time parameter speed: Speed parameter Returns: Processed tensor """ if callable(effect_or_preset): return effect_or_preset(tensor=tensor, shape=shape, time=time, speed=speed) else: # Is a Preset. Unroll me. for e_or_p in effect_or_preset.post_effects + effect_or_preset.final_effects: tensor = _apply_final_effect_or_preset(e_or_p, tensor, shape, time, speed) return tensor