Remove nodes that have been made by others already

This commit is contained in:
Dion Timmer
2023-03-31 15:05:08 -04:00
parent ca845bce15
commit 56e4a85558
16 changed files with 179 additions and 440 deletions
+12
View File
@@ -0,0 +1,12 @@
import os
import importlib
NODE_CLASS_MAPPINGS = {}
for node in os.listdir(os.path.dirname(__file__) + '\\nodes'):
if node.startswith('DT_'):
node = node.split('.')[0]
node_import = importlib.import_module('custom_nodes.VextraNodes.nodes.' + node)
print('Imported node: ' + node)
# get class node mappings from py file
NODE_CLASS_MAPPINGS.update(node_import.NODE_CLASS_MAPPINGS)
-106
View File
@@ -1,106 +0,0 @@
import torch
import numpy as np
from PIL import Image
import subprocess
import sys
try:
import blend_modes
except ModuleNotFoundError:
# install pixelsort in current venv
subprocess.check_call([sys.executable, "-m", "pip", "install", "blend-modes"])
import blend_modes
class Blend():
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
"""
Input Types
"""
return {
"required": {
"images_1": ("IMAGE",),
"images_2": ("IMAGE",),},
"optional": {
"blend_mode": ([
"soft_light",
"lighten_only",
"dodge",
"addition",
"darken_only",
"multiply",
"hard_light",
"difference",
"subtract",
"grain_extract",
"grain_merge",
"divide",
"overlay",
"normal",
],),
"blend_opacity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "apply_blend"
CATEGORY = "VextraNodes"
def tensor_to_pil(self, img):
if img is not None:
i = 255. * img.cpu().numpy().squeeze()
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
return img
def hack_alpha_channel(self, pil_image):
# Create a new image with the same size and mode as the original image and fill it with opaque white
new_image = Image.new('RGBA', pil_image.size, (255, 255, 255, 255))
# Paste the original image onto the new image
new_image.paste(pil_image, (0, 0))
return new_image
def apply_blend(self, images_1, images_2, blend_mode, blend_opacity):
#create empty tensor with the same shape as images
total_images = []
blend_fn = getattr(blend_modes, blend_mode)
if len(images_1) > len(images_2):
raise Exception("BLEND: Second set of images cannot be less than the first set of images!")
for i, image_1 in enumerate(images_1):
image = self.tensor_to_pil(image_1)
image = self.hack_alpha_channel(image)
image_2 = self.tensor_to_pil(images_2[i])
image_2 = self.hack_alpha_channel(image_2)
if image.size != image_2.size:
raise Exception("BLEND: Images must be the same size!")
image = np.array(image)
image = image.astype(float)
image_2 = np.array(image_2)
image_2 = image_2.astype(float)
out_image = blend_fn(image, image_2, blend_opacity)
out_image = Image.fromarray(out_image.astype(np.uint8))
# convert to tensor
out_image = np.array(out_image.convert("RGB")).astype(np.float32) / 255.0
out_image = torch.from_numpy(out_image).unsqueeze(0)
total_images.append(out_image)
total_images = torch.cat(total_images, 0)
return (total_images,)
NODE_CLASS_MAPPINGS = {
"Blend": Blend,
}
-236
View File
@@ -1,236 +0,0 @@
import torch
import numpy as np
from PIL import Image
import math
class Chromatic_Aberration():
"""
This node provides a simple interface to apply PixelSort blur to the output image.
"""
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
"""
Input Types
"""
return {
"required": {
"images": ("IMAGE",),},
"optional": {
"chromatic_strength": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 10.0, "step": 0.1}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "do_chromatic"
CATEGORY = "VextraNodes"
def tensor_to_pil(self, img):
if img is not None:
i = 255. * img.cpu().numpy().squeeze()
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
return img
def do_chromatic(self, images, chromatic_strength=0):
#create empty tensor with the same shape as images
total_images = []
for image in images:
image = self.tensor_to_pil(image)
image = add_chromatic(image, chromatic_strength)
# convert to tensor
out_image = np.array(image.convert("RGB")).astype(np.float32) / 255.0
out_image = torch.from_numpy(out_image).unsqueeze(0)
total_images.append(out_image)
total_images = torch.cat(total_images, 0)
return (total_images,)
def add_chromatic(im, strength: float = 0):
if strength == 0:
return im
r, g, b = im.split()
rdata = np.asarray(r)
rfinal = r
gfinal = g
bfinal = b
# enlarge the green and blue channels slightly, blue being the most enlarged
gfinal = gfinal.resize((round((1 + 0.018 * strength) * rdata.shape[1]),
round((1 + 0.018 * strength) * rdata.shape[0])), Image.ANTIALIAS)
bfinal = bfinal.resize((round((1 + 0.044 * strength) * rdata.shape[1]),
round((1 + 0.044 * strength) * rdata.shape[0])), Image.ANTIALIAS)
rwidth, rheight = rfinal.size
gwidth, gheight = gfinal.size
bwidth, bheight = bfinal.size
rhdiff = (bheight - rheight) // 2
rwdiff = (bwidth - rwidth) // 2
ghdiff = (bheight - gheight) // 2
gwdiff = (bwidth - gwidth) // 2
# Centre the channels
im = Image.merge("RGB", (
rfinal.crop((-rwdiff, -rhdiff, bwidth - rwdiff, bheight - rhdiff)),
gfinal.crop((-gwdiff, -ghdiff, bwidth - gwdiff, bheight - ghdiff)),
bfinal))
# Crop the image to the original image dimensions
return im.crop((rwdiff, rhdiff, rwidth + rwdiff, rheight + rhdiff))
NODE_CLASS_MAPPINGS = {
"Chromatic Aberration": Chromatic_Aberration
}
def cartesian_to_polar(data: np.ndarray) -> np.ndarray:
"""Returns the polar form of <data>
"""
width = data.shape[1]
height = data.shape[0]
assert (width > 2)
assert (height > 2)
assert (width % 2 == 1)
assert (height % 2 == 1)
perimeter = 2 * (width + height - 2)
halfdiag = math.ceil(((width ** 2 + height ** 2) ** 0.5) / 2)
halfw = width // 2
halfh = height // 2
ret = np.zeros((halfdiag, perimeter, 3))
# Don't want to deal with divide by zero errors...
ret[0:(halfw + 1), halfh] = data[halfh, halfw::-1]
ret[0:(halfw + 1), height + width - 2 +
halfh] = data[halfh, halfw:(halfw * 2 + 1)]
ret[0:(halfh + 1), height - 1 + halfw] = data[halfh:(halfh * 2 + 1), halfw]
ret[0:(halfh + 1), perimeter - halfw] = data[halfh::-1, halfw]
# Divide the image into 8 triangles, and use the same calculation on
# 4 triangles at a time. This is possible due to symmetry.
# This section is also responsible for the corner pixels
for i in range(0, halfh):
slope = (halfh - i) / (halfw)
diagx = ((halfdiag ** 2) / (slope ** 2 + 1)) ** 0.5
unit_xstep = diagx / (halfdiag - 1)
unit_ystep = diagx * slope / (halfdiag - 1)
for row in range(halfdiag):
ystep = round(row * unit_ystep)
xstep = round(row * unit_xstep)
if ((halfh >= ystep) and halfw >= xstep):
ret[row, i] = data[halfh - ystep, halfw - xstep]
ret[row, height - 1 - i] = data[halfh + ystep, halfw - xstep]
ret[row, height + width - 2 +
i] = data[halfh + ystep, halfw + xstep]
ret[row, height + width + height - 3 -
i] = data[halfh - ystep, halfw + xstep]
else:
break
# Remaining 4 triangles
for j in range(1, halfw):
slope = (halfh) / (halfw - j)
diagx = ((halfdiag ** 2) / (slope ** 2 + 1)) ** 0.5
unit_xstep = diagx / (halfdiag - 1)
unit_ystep = diagx * slope / (halfdiag - 1)
for row in range(halfdiag):
ystep = round(row * unit_ystep)
xstep = round(row * unit_xstep)
if (halfw >= xstep and halfh >= ystep):
ret[row, height - 1 + j] = data[halfh + ystep, halfw - xstep]
ret[row, height + width - 2 -
j] = data[halfh + ystep, halfw + xstep]
ret[row, height + width + height - 3 +
j] = data[halfh - ystep, halfw + xstep]
ret[row, perimeter - j] = data[halfh - ystep, halfw - xstep]
else:
break
return ret
def polar_to_cartesian(data: np.ndarray, width: int, height: int) -> np.ndarray:
"""Returns the cartesian form of <data>.
<width> is the original width of the cartesian image
<height> is the original height of the cartesian image
"""
assert (width > 2)
assert (height > 2)
assert (width % 2 == 1)
assert (height % 2 == 1)
perimeter = 2 * (width + height - 2)
halfdiag = math.ceil(((width ** 2 + height ** 2) ** 0.5) / 2)
halfw = width // 2
halfh = height // 2
ret = np.zeros((height, width, 3))
def div0():
# Don't want to deal with divide by zero errors...
ret[halfh, halfw::-1] = data[0:(halfw + 1), halfh]
ret[halfh, halfw:(halfw * 2 + 1)] = data[0:(halfw + 1),
height + width - 2 + halfh]
ret[halfh:(halfh * 2 + 1), halfw] = data[0:(halfh + 1), height - 1 + halfw]
ret[halfh::-1, halfw] = data[0:(halfh + 1), perimeter - halfw]
div0()
# Same code as above, except the order of the assignments are switched
# Code blocks are split up for easier profiling
def part1():
for i in range(0, halfh):
slope = (halfh - i) / (halfw)
diagx = ((halfdiag ** 2) / (slope ** 2 + 1)) ** 0.5
unit_xstep = diagx / (halfdiag - 1)
unit_ystep = diagx * slope / (halfdiag - 1)
for row in range(halfdiag):
ystep = round(row * unit_ystep)
xstep = round(row * unit_xstep)
if ((halfh >= ystep) and halfw >= xstep):
ret[halfh - ystep, halfw - xstep] = \
data[row, i]
ret[halfh + ystep, halfw - xstep] = \
data[row, height - 1 - i]
ret[halfh + ystep, halfw + xstep] = \
data[row, height + width - 2 + i]
ret[halfh - ystep, halfw + xstep] = \
data[row, height + width + height - 3 - i]
else:
break
part1()
def part2():
for j in range(1, halfw):
slope = (halfh) / (halfw - j)
diagx = ((halfdiag ** 2) / (slope ** 2 + 1)) ** 0.5
unit_xstep = diagx / (halfdiag - 1)
unit_ystep = diagx * slope / (halfdiag - 1)
for row in range(halfdiag):
ystep = round(row * unit_ystep)
xstep = round(row * unit_xstep)
if (halfw >= xstep and halfh >= ystep):
ret[halfh + ystep, halfw - xstep] = \
data[row, height - 1 + j]
ret[halfh + ystep, halfw + xstep] = \
data[row, height + width - 2 - j]
ret[halfh - ystep, halfw + xstep] = \
data[row, height + width + height - 3 + j]
ret[halfh - ystep, halfw - xstep] = \
data[row, perimeter - j]
else:
break
part2()
# Repairs black/missing pixels in the transformed image
def set_zeros():
zero_mask = ret[1:-1, 1:-1] == 0
ret[1:-1, 1:-1] = np.where(zero_mask, (ret[:-2, 1:-1] + ret[2:, 1:-1]) / 2, ret[1:-1, 1:-1])
set_zeros()
return ret
-96
View File
@@ -1,96 +0,0 @@
import torch
import numpy as np
from PIL import Image
class Displacement_Map():
"""
This node provides a simple interface to apply PixelSort blur to the output image.
"""
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
"""
Input Types
"""
return {
"required": {
"images": ("IMAGE",),
"displacement_maps": ("IMAGE",),},
"optional": {
"scale": ("FLOAT", {"default": 5.0, "min": 1.0, "max": 500.0, "step": 0.1}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "do_displace"
CATEGORY = "VextraNodes"
def tensor_to_pil(self, img):
if img is not None:
i = 255. * img.cpu().numpy().squeeze()
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
return img
def do_displace(self, images, displacement_maps, scale):
#create empty tensor with the same shape as images
total_images = []
if len(images) > len(displacement_maps):
raise Exception("Number of images must be equal or less than the number of displacement maps!")
for i, image in enumerate(images):
displacement_map = displacement_maps[i]
displacement_map = self.tensor_to_pil(displacement_map)
image = self.tensor_to_pil(image)
if displacement_map.size != image.size:
raise Exception("Displacement map and image must be the same size!")
image = apply_displacement_map(image, displacement_map, scale)
# convert to tensor
out_image = np.array(image.convert("RGB")).astype(np.float32) / 255.0
out_image = torch.from_numpy(out_image).unsqueeze(0)
total_images.append(out_image)
total_images = torch.cat(total_images, 0)
return (total_images,)
NODE_CLASS_MAPPINGS = {
"Displacement Map": Displacement_Map
}
def apply_displacement_map(image, displacement_map, scale):
# Convert PIL images to NumPy arrays
image_array = np.array(image)
displacement_map_array = np.array(displacement_map)
# Get the dimensions of the image
height, width, _ = image_array.shape
# Calculate the displacement offsets based on the scale factor
displacement_offsets = (displacement_map_array / 255 - 0.5) * scale
# Create arrays for the X and Y coordinates of the pixels
x_coords, y_coords = np.meshgrid(np.arange(width), np.arange(height))
# Apply the displacement offsets to the X and Y coordinates
x_displaced = (x_coords + displacement_offsets[..., 0]).clip(0, width - 1).astype(int)
y_displaced = (y_coords + displacement_offsets[..., 1]).clip(0, height - 1).astype(int)
# Create a new array with the same shape as the original image and copy the displaced pixels
displaced_image_array = np.zeros_like(image_array)
displaced_image_array[y_coords, x_coords] = image_array[y_displaced, x_displaced]
# Convert the displaced image array back to a PIL image
displaced_image = Image.fromarray(displaced_image_array)
return displaced_image
@@ -24,7 +24,7 @@ class FontText():
"size": ("INT", {"default": 50, "min": 2, "max": 1000, "step": 1}),
"x": ("INT", {"default": 50, "min": 2, "max": 10000, "step": 1}),
"y": ("INT", {"default": 50, "min": 2, "max": 10000, "step": 1}),
"text": ("STRING", {"default": "Hello World"}),
"text": ("STRING", {"default": "Hello World", "multiline": True}),
"color": ("STRING", {"default": 'rgba(255, 255, 255, 255)'}),
"anchor": (["Bottom Left Corner", "Center"],),
"rotate": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 360.0, "step": 0.1}),
@@ -43,7 +43,7 @@ class FontText():
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
return img
def do_font(self, images, font_ttf, size, x, y, text, color, anchor, rotate, color_mode):
def do_font(self, images, font_ttf, size, x, y, color, anchor, rotate, color_mode, text):
#create empty tensor with the same shape as images
total_images = []
center_anchor = True if anchor == 'Center' else False
+69
View File
@@ -0,0 +1,69 @@
import torch
import numpy as np
from PIL import Image
import subprocess
import sys
try:
from glitch_this import ImageGlitcher
except ModuleNotFoundError:
# install pixelsort in current venv
subprocess.check_call([sys.executable, "-m", "pip", "install", "glitch-this"])
from glitch_this import ImageGlitcher
class GlitchThis():
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
"""
Input Types
"""
return {
"required": {
"images": ("IMAGE",),},
"optional": {
"glitch_amount": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.01}),
"color_offset": (['Disable', 'Enable'],),
"scan_lines": (['Disable', 'Enable'],),
"seed": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "apply_glitch"
CATEGORY = "VextraNodes"
def tensor_to_pil(self, img):
if img is not None:
i = 255. * img.cpu().numpy().squeeze()
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
return img
def string2bool(self, v):
return v == 'Enable'
def apply_glitch(self, images, glitch_amount=1, color_offset='Disable', scan_lines='Disable', seed=0):
color_offset = self.string2bool(color_offset)
scan_lines = self.string2bool(scan_lines)
glitcher = ImageGlitcher()
#create empty tensor with the same shape as images
total_images = []
for image in images:
image = self.tensor_to_pil(image)
image = glitcher.glitch_image(image, glitch_amount, color_offset=color_offset, scan_lines=scan_lines, seed=seed)
# convert to tensor
out_image = np.array(image.convert("RGB")).astype(np.float32) / 255.0
out_image = torch.from_numpy(out_image).unsqueeze(0)
total_images.append(out_image)
total_images = torch.cat(total_images, 0)
return (total_images,)
NODE_CLASS_MAPPINGS = {
"GlitchThis Effect": GlitchThis,
}
+39
View File
@@ -0,0 +1,39 @@
import os
import torchvision.transforms.functional as TF
from PIL import Image
class PictureIndex:
def __init__(self) -> None:
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"path": ("STRING", {"default": ""}),
"index": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "doStuff"
CATEGORY = "VextraNodes"
def doStuff(self, path, index):
if not os.path.exists(path):
raise Exception("Path does not exist")
images = []
for image in os.listdir(path):
if any(image.endswith(ext) for ext in [".png", ".jpg", ".jpeg"]):
images.append(image)
image = Image.open(os.path.join(path, images[index]))
image = TF.to_tensor(image)
print(image)
return (image,)
NODE_CLASS_MAPPINGS = {
"Load Picture Index": PictureIndex,
}
+57
View File
@@ -0,0 +1,57 @@
from transformers import GPT2Tokenizer, GPT2LMHeadModel, set_seed
import os
import random
script_path = os.path.dirname(os.path.realpath(__file__))
comfy_path = script_path.split('custom_nodes')[0]
if not os.path.exists(f'{comfy_path}custom_nodes/VextraNodes/binary/distilgpt2-stable-diffusion-v2'):
print('Downloading model...')
os.system(f'git clone https://huggingface.co/FredZhang7/distilgpt2-stable-diffusion-v2 {comfy_path}custom_nodes/VextraNodes/binary/distilgpt2-stable-diffusion-v2')
tokenizer = GPT2Tokenizer.from_pretrained('distilgpt2')
tokenizer.add_special_tokens({'pad_token': '[PAD]'})
model = GPT2LMHeadModel.from_pretrained(f'{comfy_path}custom_nodes/VextraNodes/binary/distilgpt2-stable-diffusion-v2')
class GetPrompt():
"""
This node provides a simple interface to apply PixelSort blur to the output image.
"""
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
"""
Input Types
"""
return {
"required": {
"temperature": ("FLOAT", {"default": 0.9, "min": 0.1, "max": 1.0, "step": 0.1}),
"top_k": ("INT", {"default": 8, "min": 1, "max": 100, "step": 1}),
"max_length": ("INT", {"default": 80, "min": 1, "max": 1000, "step": 1}),
"repetition_penalty": ("FLOAT", {"default": 1.2, "min": 0, "max": 10, "step": 0.1}),
"seed": ("INT", {"default": -1, "min": -1, "max": 1000000000, "step": 1}),
"prompt": ("STRING", {"default": '', "multiline": True}),
},
"optional": {
},
}
RETURN_TYPES = ("STRING",)
FUNCTION = "get_prompt_from_model"
CATEGORY = "VextraNodes"
def get_prompt_from_model(self, temperature=0.9, top_k=8, max_length=80, repetition_penalty=1.2, seed=-1, prompt=''):
seed = int(seed) if seed != -1 else random.randint(1, 1000000)
set_seed(seed)
num_return_sequences=1 # the number of results to generate
input_ids = tokenizer(prompt, return_tensors='pt').input_ids
output = model.generate(input_ids, do_sample=True, temperature=temperature, top_k=top_k, max_length=max_length, num_return_sequences=num_return_sequences, repetition_penalty=repetition_penalty, penalty_alpha=0.6, no_repeat_ngram_size=1, early_stopping=True)
return (tokenizer.decode(output[0], skip_special_tokens=True), )
NODE_CLASS_MAPPINGS = {
"Prettify Prompt Using distilgpt2": GetPrompt
}