Remove nodes that have been made by others already
This commit is contained in:
+12
@@ -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)
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user