Merge branch 'main' of https://github.com/kijai/ComfyUI-KJNodes

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

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",
 }

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)