diff --git a/fluid.py b/fluid.py new file mode 100644 index 0000000..c069198 --- /dev/null +++ b/fluid.py @@ -0,0 +1,67 @@ +import numpy as np +from scipy.ndimage import map_coordinates, spline_filter +from scipy.sparse.linalg import factorized + +from .numerical import difference, operator + + +class Fluid: + def __init__(self, shape, *quantities, pressure_order=1, advect_order=3): + self.shape = shape + self.dimensions = len(shape) + + # Prototyping is simplified by dynamically + # creating advected quantities as needed. + self.quantities = quantities + for q in quantities: + setattr(self, q, np.zeros(shape)) + + self.indices = np.indices(shape) + self.velocity = np.zeros((self.dimensions, *shape)) + + laplacian = operator(shape, difference(2, pressure_order)) + self.pressure_solver = factorized(laplacian) + + self.advect_order = advect_order + + def step(self): + # Advection is computed backwards in time as described in Stable Fluids. + advection_map = self.indices - self.velocity + + # SciPy's spline filter introduces checkerboard divergence. + # A linear blend of the filtered and unfiltered fields based + # on some value epsilon eliminates this error. + def advect(field, filter_epsilon=10e-2, mode='constant'): + filtered = spline_filter(field, order=self.advect_order, mode=mode) + field = filtered * (1 - filter_epsilon) + field * filter_epsilon + return map_coordinates(field, advection_map, prefilter=False, order=self.advect_order, mode=mode) + + # Apply advection to each axis of the + # velocity field and each user-defined quantity. + for d in range(self.dimensions): + self.velocity[d] = advect(self.velocity[d]) + + for q in self.quantities: + setattr(self, q, advect(getattr(self, q))) + + # Compute the jacobian at each point in the + # velocity field to extract curl and divergence. + jacobian_shape = (self.dimensions,) * 2 + partials = tuple(np.gradient(d) for d in self.velocity) + jacobian = np.stack(partials).reshape(*jacobian_shape, *self.shape) + + divergence = jacobian.trace() + + # If this curl calculation is extended to 3D, the y-axis value must be negated. + # This corresponds to the coefficients of the levi-civita symbol in that dimension. + # Higher dimensions do not have a vector -> scalar, or vector -> vector, + # correspondence between velocity and curl due to differing isomorphisms + # between exterior powers in dimensions != 2 or 3 respectively. + curl_mask = np.triu(np.ones(jacobian_shape, dtype=bool), k=1) + curl = (jacobian[curl_mask] - jacobian[curl_mask.T]).squeeze() + + # Apply the pressure correction to the fluid's velocity field. + pressure = self.pressure_solver(divergence.flatten()).reshape(self.shape) + self.velocity -= np.gradient(pressure) + + return divergence, curl, pressure \ No newline at end of file diff --git a/nodes.py b/nodes.py index 47e9448..1af61ec 100644 --- a/nodes.py +++ b/nodes.py @@ -6,6 +6,8 @@ import numpy as np from PIL import ImageColor, Image, ImageDraw, ImageFont import os import librosa +from scipy.special import erf +from .fluid import Fluid from nodes import MAX_RESOLUTION @@ -33,6 +35,84 @@ def gaussian_kernel(kernel_size: int, sigma: float, device=None): g = torch.exp(-(d * d) / (2.0 * sigma * sigma)) return g / g.sum() +class CreateFluidMask: + + RETURN_TYPES = ("IMAGE", "MASK") + FUNCTION = "createfluidmask" + CATEGORY = "KJNodes" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "invert": ("BOOLEAN", {"default": False}), + "frames": ("INT", {"default": 0,"min": 0, "max": 255, "step": 1}), + "width": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}), + "height": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}), + "inflow_count": ("INT", {"default": 3,"min": 0, "max": 255, "step": 1}), + "inflow_velocity": ("INT", {"default": 1,"min": 0, "max": 255, "step": 1}), + "inflow_radius": ("INT", {"default": 8,"min": 0, "max": 255, "step": 1}), + "inflow_padding": ("INT", {"default": 50,"min": 0, "max": 255, "step": 1}), + "inflow_duration": ("INT", {"default": 60,"min": 0, "max": 255, "step": 1}), + + }, + } + #using code from https://github.com/GregTJ/stable-fluids + def createfluidmask(self, frames, width, height, invert, inflow_count, inflow_velocity, inflow_radius, inflow_padding, inflow_duration): + out = [] + masks = [] + print(frames) + RESOLUTION = width, height + DURATION = frames + + INFLOW_PADDING = inflow_padding + INFLOW_DURATION = inflow_duration + INFLOW_RADIUS = inflow_radius + INFLOW_VELOCITY = inflow_velocity + INFLOW_COUNT = inflow_count + + print('Generating fluid solver, this may take some time.') + fluid = Fluid(RESOLUTION, 'dye') + + center = np.floor_divide(RESOLUTION, 2) + r = np.min(center) - INFLOW_PADDING + + points = np.linspace(-np.pi, np.pi, INFLOW_COUNT, endpoint=False) + points = tuple(np.array((np.cos(p), np.sin(p))) for p in points) + normals = tuple(-p for p in points) + points = tuple(r * p + center for p in points) + + inflow_velocity = np.zeros_like(fluid.velocity) + inflow_dye = np.zeros(fluid.shape) + for p, n in zip(points, normals): + mask = np.linalg.norm(fluid.indices - p[:, None, None], axis=0) <= INFLOW_RADIUS + inflow_velocity[:, mask] += n[:, None] * INFLOW_VELOCITY + inflow_dye[mask] = 1 + + + for f in range(DURATION): + print(f'Computing frame {f + 1} of {DURATION}.') + if f <= INFLOW_DURATION: + fluid.velocity += inflow_velocity + fluid.dye += inflow_dye + + curl = fluid.step()[1] + # Using the error function to make the contrast a bit higher. + # Any other sigmoid function e.g. smoothstep would work. + curl = (erf(curl * 2) + 1) / 4 + + color = np.dstack((curl, np.ones(fluid.shape), fluid.dye)) + color = (np.clip(color, 0, 1) * 255).astype('uint8') + image = np.array(color).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[None,] + mask = image[:, :, :, 0] + masks.append(mask) + out.append(image) + + if invert: + return (1.0 - torch.cat(out, dim=0),1.0 - torch.cat(masks, dim=0),) + return (torch.cat(out, dim=0),torch.cat(masks, dim=0),) + class CreateAudioMask: RETURN_TYPES = ("IMAGE",) @@ -548,6 +628,7 @@ NODE_CLASS_MAPPINGS = { "CreateTextMask": CreateTextMask, "CreateAudioMask": CreateAudioMask, "CreateFadeMask": CreateFadeMask, + "CreateFluidMask" :CreateFluidMask, } NODE_DISPLAY_NAME_MAPPINGS = { "INTConstant": "INT Constant", @@ -559,4 +640,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "CreateGradientMask": "CreateGradientMask", "CreateTextMask" : "CreateTextMask", "CreateFadeMask" : "CreateFadeMask", + "CreateFluidMask" : "CreateFluidMask", } \ No newline at end of file diff --git a/numerical.py b/numerical.py new file mode 100644 index 0000000..b5b88bc --- /dev/null +++ b/numerical.py @@ -0,0 +1,25 @@ +from functools import reduce +from itertools import cycle +from math import factorial + +import numpy as np +import scipy.sparse as sp + + +def difference(derivative, accuracy=1): + # Central differences implemented based on the article here: + # http://web.media.mit.edu/~crtaylor/calculator.html + derivative += 1 + radius = accuracy + derivative // 2 - 1 + points = range(-radius, radius + 1) + coefficients = np.linalg.inv(np.vander(points)) + return coefficients[-derivative] * factorial(derivative - 1), points + + +def operator(shape, *differences): + # Credit to Philip Zucker for figuring out + # that kronsum's argument order is reversed. + # Without that bit of wisdom I'd have lost it. + differences = zip(shape, cycle(differences)) + factors = (sp.diags(*diff, shape=(dim,) * 2) for dim, diff in differences) + return reduce(lambda a, f: sp.kronsum(f, a, format='csc'), factors) \ No newline at end of file