mirror of
https://git.datalinker.icu/kijai/ComfyUI-KJNodes.git
synced 2026-03-16 14:27:05 +08:00
Fix aspect ration on voronoi mask and add initial version of magic mask
This commit is contained in:
kijai
parent
02b00c6480
commit
3569b842df
95
magictex.py
Normal file
95
magictex.py
Normal file
@ -0,0 +1,95 @@
|
||||
"""Generates psychedelic color textures in the spirit of Blender's magic texture shader using Python/Numpy
|
||||
|
||||
https://github.com/cheind/magic-texture
|
||||
"""
|
||||
from typing import Tuple, Optional
|
||||
import numpy as np
|
||||
|
||||
|
||||
def coordinate_grid(shape: Tuple[int, int], dtype=np.float32):
|
||||
"""Returns a three-dimensional coordinate grid of given shape for use in `magic`."""
|
||||
x = np.linspace(-1, 1, shape[1], endpoint=True, dtype=dtype)
|
||||
y = np.linspace(-1, 1, shape[0], endpoint=True, dtype=dtype)
|
||||
X, Y = np.meshgrid(x, y)
|
||||
XYZ = np.stack((X, Y, np.ones_like(X)), -1)
|
||||
return XYZ
|
||||
|
||||
|
||||
def random_transform(coords: np.ndarray, rng: np.random.Generator = None):
|
||||
"""Returns randomly transformed coordinates"""
|
||||
H, W = coords.shape[:2]
|
||||
rng = rng or np.random.default_rng()
|
||||
m = rng.uniform(-1.0, 1.0, size=(3, 3)).astype(coords.dtype)
|
||||
return (coords.reshape(-1, 3) @ m.T).reshape(H, W, 3)
|
||||
|
||||
|
||||
def magic(
|
||||
coords: np.ndarray,
|
||||
depth: Optional[int] = None,
|
||||
distortion: Optional[int] = None,
|
||||
rng: np.random.Generator = None,
|
||||
):
|
||||
"""Returns color magic color texture.
|
||||
|
||||
The implementation is based on Blender's (https://www.blender.org/) magic
|
||||
texture shader. The following adaptions have been made:
|
||||
- we exchange the nested if-cascade by a probabilistic iterative approach
|
||||
|
||||
Kwargs
|
||||
------
|
||||
coords: HxWx3 array
|
||||
Coordinates transformed into colors by this method. See
|
||||
`magictex.coordinate_grid` to generate the default.
|
||||
depth: int (optional)
|
||||
Number of transformations applied. Higher numbers lead to more
|
||||
nested patterns. If not specified, randomly sampled.
|
||||
distortion: float (optional)
|
||||
Distortion of patterns. Larger values indicate more distortion,
|
||||
lower values tend to generate smoother patterns. If not specified,
|
||||
randomly sampled.
|
||||
rng: np.random.Generator
|
||||
Optional random generator to draw samples from.
|
||||
|
||||
Returns
|
||||
-------
|
||||
colors: HxWx3 array
|
||||
Three channel color image in range [0,1]
|
||||
"""
|
||||
rng = rng or np.random.default_rng()
|
||||
if distortion is None:
|
||||
distortion = rng.uniform(1, 4)
|
||||
if depth is None:
|
||||
depth = rng.integers(1, 5)
|
||||
|
||||
H, W = coords.shape[:2]
|
||||
XYZ = coords
|
||||
x = np.sin((XYZ[..., 0] + XYZ[..., 1] + XYZ[..., 2]) * distortion)
|
||||
y = np.cos((-XYZ[..., 0] + XYZ[..., 1] - XYZ[..., 2]) * distortion)
|
||||
z = -np.cos((-XYZ[..., 0] - XYZ[..., 1] + XYZ[..., 2]) * distortion)
|
||||
|
||||
if depth > 0:
|
||||
x *= distortion
|
||||
y *= distortion
|
||||
z *= distortion
|
||||
y = -np.cos(x - y + z)
|
||||
y *= distortion
|
||||
|
||||
xyz = [x, y, z]
|
||||
fns = [np.cos, np.sin]
|
||||
for _ in range(1, depth):
|
||||
axis = rng.choice(3)
|
||||
fn = fns[rng.choice(2)]
|
||||
signs = rng.binomial(n=1, p=0.5, size=4) * 2 - 1
|
||||
|
||||
xyz[axis] = signs[-1] * fn(
|
||||
signs[0] * xyz[0] + signs[1] * xyz[1] + signs[2] * xyz[2]
|
||||
)
|
||||
xyz[axis] *= distortion
|
||||
|
||||
x, y, z = xyz
|
||||
x /= 2 * distortion
|
||||
y /= 2 * distortion
|
||||
z /= 2 * distortion
|
||||
c = 0.5 - np.stack((x, y, z), -1)
|
||||
np.clip(c, 0, 1.0)
|
||||
return c
|
||||
98
nodes.py
98
nodes.py
@ -1957,7 +1957,7 @@ class OffsetMask:
|
||||
|
||||
batch_size, height, width = mask.shape
|
||||
|
||||
if angle is not 0 and incremental:
|
||||
if angle != 0 and incremental:
|
||||
for i in range(batch_size):
|
||||
rotation_angle = angle * (i+1)
|
||||
mask[i] = TF.rotate(mask[i].unsqueeze(0), rotation_angle).squeeze(0)
|
||||
@ -2128,7 +2128,7 @@ class CreateVoronoiMask:
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"frames": ("INT", {"default": 1,"min": 1, "max": 4096, "step": 1}),
|
||||
"frames": ("INT", {"default": 16,"min": 2, "max": 4096, "step": 1}),
|
||||
"num_points": ("INT", {"default": 15,"min": 1, "max": 4096, "step": 1}),
|
||||
"line_width": ("INT", {"default": 4,"min": 1, "max": 4096, "step": 1}),
|
||||
"speed": ("FLOAT", {"default": 0.5,"min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
@ -2142,9 +2142,15 @@ class CreateVoronoiMask:
|
||||
batch_size = frames
|
||||
out = []
|
||||
|
||||
# Create start and end points for each point
|
||||
# Calculate aspect ratio
|
||||
aspect_ratio = frame_width / frame_height
|
||||
|
||||
# Create start and end points for each point, considering the aspect ratio
|
||||
start_points = np.random.rand(num_points, 2)
|
||||
start_points[:, 0] *= aspect_ratio
|
||||
|
||||
end_points = np.random.rand(num_points, 2)
|
||||
end_points[:, 0] *= aspect_ratio
|
||||
|
||||
for i in range(batch_size):
|
||||
# Interpolate the points' positions based on the current frame
|
||||
@ -2152,15 +2158,18 @@ class CreateVoronoiMask:
|
||||
t = np.clip(t, 0, 1) # ensure t is in [0, 1]
|
||||
points = (1 - t) * start_points + t * end_points # lerp
|
||||
|
||||
# Adjust points for aspect ratio
|
||||
points[:, 0] *= aspect_ratio
|
||||
|
||||
vor = Voronoi(points)
|
||||
|
||||
# Create a blank image with a white background
|
||||
fig, ax = plt.subplots()
|
||||
plt.subplots_adjust(left=0, right=1, bottom=0, top=1)
|
||||
ax.set_xlim([0, 1]); ax.set_ylim([0, 1])
|
||||
ax.set_xlim([0, aspect_ratio]); ax.set_ylim([0, 1]) # adjust x limits
|
||||
ax.axis('off')
|
||||
ax.margins(0, 0)
|
||||
fig.set_size_inches(frame_width/100, frame_height/100)
|
||||
fig.set_size_inches(aspect_ratio * frame_height/100, frame_height/100) # adjust figure size
|
||||
ax.fill_between([0, 1], [0, 1], color='white')
|
||||
|
||||
# Plot each Voronoi ridge
|
||||
@ -2181,6 +2190,83 @@ class CreateVoronoiMask:
|
||||
|
||||
return (torch.stack(out, dim=0), 1.0 - torch.stack(out, dim=0),)
|
||||
|
||||
from mpl_toolkits.axes_grid1 import ImageGrid
|
||||
|
||||
from .magictex import *
|
||||
|
||||
class CreateMagicMask:
|
||||
|
||||
RETURN_TYPES = ("MASK", "MASK",)
|
||||
RETURN_NAMES = ("mask", "mask_inverted",)
|
||||
FUNCTION = "createmagicmask"
|
||||
CATEGORY = "KJNodes/masking/generate"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"frames": ("INT", {"default": 16,"min": 2, "max": 4096, "step": 1}),
|
||||
"depth": ("INT", {"default": 2,"min": 1, "max": 50, "step": 1}),
|
||||
"distortion": ("FLOAT", {"default": 1.0,"min": 0.0, "max": 10.0, "step": 0.01}),
|
||||
"seed": ("INT", {"default": 123,"min": 0, "max": 99999999, "step": 1}),
|
||||
"transitions": ("INT", {"default": 2,"min": 1, "max": 20, "step": 1}),
|
||||
"frame_width": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}),
|
||||
"frame_height": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}),
|
||||
},
|
||||
}
|
||||
|
||||
def createmagicmask(self, frames, transitions, depth, distortion, seed, frame_width, frame_height):
|
||||
rng = np.random.default_rng(seed)
|
||||
out = []
|
||||
coords = coordinate_grid((frame_width, frame_height))
|
||||
|
||||
# Calculate the number of frames for each transition
|
||||
frames_per_transition = frames // transitions
|
||||
|
||||
# Generate a base set of parameters
|
||||
base_params = {
|
||||
"coords": random_transform(coords, rng),
|
||||
"depth": depth,
|
||||
"distortion": distortion,
|
||||
}
|
||||
for t in range(transitions):
|
||||
# Generate a second set of parameters that is at most max_diff away from the base parameters
|
||||
params1 = base_params.copy()
|
||||
params2 = base_params.copy()
|
||||
|
||||
params1['coords'] = random_transform(coords, rng)
|
||||
params2['coords'] = random_transform(coords, rng)
|
||||
|
||||
for i in range(frames_per_transition):
|
||||
# Compute the interpolation factor
|
||||
alpha = i / frames_per_transition
|
||||
|
||||
# Interpolate between the two sets of parameters
|
||||
params = params1.copy()
|
||||
params['coords'] = (1 - alpha) * params1['coords'] + alpha * params2['coords']
|
||||
|
||||
tex = magic(**params)
|
||||
|
||||
fig = plt.figure(figsize=(10, 10))
|
||||
ax = fig.add_subplot(111)
|
||||
plt.subplots_adjust(left=0, right=1, bottom=0, top=1)
|
||||
|
||||
ax.get_yaxis().set_ticks([])
|
||||
ax.get_xaxis().set_ticks([])
|
||||
ax.imshow(tex, aspect='auto')
|
||||
|
||||
fig.canvas.draw()
|
||||
img = np.array(fig.canvas.renderer._renderer)
|
||||
|
||||
plt.close(fig)
|
||||
|
||||
pil_img = Image.fromarray(img).convert("L")
|
||||
mask = torch.tensor(np.array(pil_img)) / 255.0
|
||||
|
||||
out.append(mask)
|
||||
|
||||
return (torch.stack(out, dim=0), 1.0 - torch.stack(out, dim=0),)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"INTConstant": INTConstant,
|
||||
"FloatConstant": FloatConstant,
|
||||
@ -2220,6 +2306,7 @@ NODE_CLASS_MAPPINGS = {
|
||||
"WidgetToString": WidgetToString,
|
||||
"CreateShapeMask": CreateShapeMask,
|
||||
"CreateVoronoiMask": CreateVoronoiMask,
|
||||
"CreateMagicMask": CreateMagicMask,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"INTConstant": "INT Constant",
|
||||
@ -2259,4 +2346,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"WidgetToString": "WidgetToString",
|
||||
"CreateShapeMask": "CreateShapeMask",
|
||||
"CreateVoronoiMask": "CreateVoronoiMask",
|
||||
"CreateMagicMask": "CreateMagicMask",
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user