nodes/TashaSkyUp-EternalKernelPytorchNodes
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Remove non-pytorch nodes

Browse Source
This commit is contained in:
Tasha Upchurch
2025-06-22 01:04:19 -06:00
parent 056faf8ac3
commit 0ee6e2820a
49 changed files with 2 additions and 31333 deletions
Download Patch File Download Diff File Expand all files Collapse all files
3d.py
-159
View File
@@ -1,159 +0,0 @@
NODE_CLASS_MAPPINGS = {}
NODE_DISPLAY_NAME_MAPPINGS = {}
def ETK_3d_base(cls):
import functools
cls.CATEGORY = "ETK/3d"
# Add spaces to the camel case class name
pretty_name = cls.__name__
for i in range(1, len(pretty_name)):
if pretty_name[i].isupper():
pretty_name = pretty_name[:i] + " " + pretty_name[i:]
cls.DISPLAY_NAME = pretty_name
NODE_CLASS_MAPPINGS[cls.__name__] = cls
NODE_DISPLAY_NAME_MAPPINGS[cls.DISPLAY_NAME] = pretty_name
return cls
@ETK_3d_base
class TripoSRNode:
def __init__(self):
import os
from folder_paths import output_directory
self.output_dir = output_directory
self.repo_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "repos", "img23d")
self.repo_url = "https://github.com/VAST-AI-Research/TripoSR.git"
self.venv_dir = os.path.join(self.repo_dir, "venv")
self.python_executable = os.path.join(self.venv_dir, "Scripts", "python")
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"output_dir": ("STRING", {"default": "outputs/TripoSR"}),
"mc_resolution": ("INT", {"default": 256}),
}
}
RETURN_TYPES = ("STRING",)
FUNCTION = "generate_3d_models"
CATEGORY = "3D"
def generate_3d_models(self, image, output_dir=None, mc_resolution=256):
import asyncio
if output_dir is not None:
#self.output_dir = f"{self.output_dir}/{output_dir}"
self.final_out_dir = f"{self.output_dir}/{output_dir}"
# Create a new asyncio loop
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
# Run the async function within the loop
output_paths = loop.run_until_complete(self.async_generate_3d_models(image, mc_resolution))
return (output_paths,)
async def async_generate_3d_models(self, image, mc_resolution):
import os
import numpy as np
from PIL import Image
import asyncio
# Check if the TripoSR repository is installed
if not os.path.exists(self.repo_dir):
# Clone the repository if it doesn't exist
await self.clone_repository()
# Create a virtual environment and install requirements
await self.create_venv_and_install_requirements()
# Create a temporary directory for saving input images
temp_dir = "temp_images"
os.makedirs(temp_dir, exist_ok=True)
# Save each image in the batch as a separate file
# each image is a torch tensor with shape (b,h,w,c)
image_paths = []
for i in range(image.shape[0]):
image_np = image[i].cpu().numpy()
image_pil = Image.fromarray((image_np * 255).astype(np.uint8))
image_path = os.path.join(temp_dir, f"image_{i}.png")
image_pil.save(image_path)
image_paths.append(image_path)
# Run inference using the TripoSR CLI for each image simultaneously
output_paths = await asyncio.gather(*[self.run_inference(image_path,mc_resolution) for image_path in image_paths])
# Clean up the temporary image files
for image_path in image_paths:
os.remove(image_path)
return output_paths
async def run_inference(self, image_path, mc_resolution):
import asyncio
import os
# Create a virtual environment and install requirements if they don't exist
command = [
self.python_executable,
os.path.join(self.repo_dir, "run.py"),
image_path,
"--mc-resolution",
f"{mc_resolution}",
"--output-dir",
self.final_out_dir,
]
process = await asyncio.create_subprocess_exec(*command)
await process.wait()
output_path = self.get_output_path(image_path)
return output_path
async def clone_repository(self):
import asyncio
# Clone the TripoSR repository
command = ["git", "clone", self.repo_url, self.repo_dir]
process = await asyncio.create_subprocess_exec(*command)
await process.wait()
async def create_venv_and_install_requirements(self):
import asyncio
import os
# Create a virtual environment
command = ["python", "-m", "venv", self.venv_dir]
# use os to execute the command
os.system(" ".join(command))
# install pip "python -m ensurepip"
command = [self.python_executable, "-m", "ensurepip"]
os.system(" ".join(command))
# install wheel
command = [self.python_executable, "-m", "pip", "install", "wheel"]
os.system(" ".join(command))
# install torch
# TODO: detect what version of torch to install based on what is installed on system
command = [self.python_executable, "-m", "pip", "install", "torch",
"--index-url https://download.pytorch.org/whl/cu121"]
#command = [self.python_executable, "-m", "pip", "install", "torch",
# "--index-url https://download.pytorch.org/whl/cu118"]
os.system(" ".join(command))
# Install requirements
command = [self.python_executable, "-m", "pip", "install", "-r",
os.path.join(self.repo_dir, "requirements.txt")]
os.system(" ".join(command))
def get_output_path(self, image_path):
import os
# Extract the base name of the input image file
base_name = os.path.splitext(os.path.basename(image_path))[0]
# Construct the output path based on the base name and output directory
output_path = os.path.join(self.output_dir, base_name, "obj", "mesh_norm.obj")
return output_path
AVF/__init__.py
-6
View File
@@ -1,6 +0,0 @@
from .. import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
from .audio_nodes import *
from .video_nodes import *
from .common import run_ffmpeg_command
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "run_ffmpeg_command"]
AVF/audio_nodes.py
-342
View File
@@ -1,342 +0,0 @@
# audio_module.py
from .common import ETK_AVF_base, duplicate_file, run_ffmpeg_command
@ETK_AVF_base
class AudioFolderProvider:
"""
Given a folder path, creates the folder or empties it entirely returns the folder path as
a string but as the node type "AUDIOFOLDER"
"""
@classmethod
def INPUT_TYPES(cls):
ret = {
"required": {
"folder_path": ("STRING", {"default": ""}),
"clear_folder": ("BOOLEAN", {"default": True}),
},
"optional": {
"activator:": ("*", {"default": None})
}
}
return ret
RETURN_TYPES = ("AUDIOFOLDER",)
FUNCTION = "audio_folder_provider"
def audio_folder_provider(self, **kwargs):
import shutil
import os
# check if the folder_path is none
if kwargs["folder_path"] is None:
raise ValueError("folder_path must be specified")
# set folder exists flag by using the os.path.exists function
folder_exists = os.path.exists(kwargs["folder_path"])
if not folder_exists:
os.makedirs(kwargs["folder_path"])
folder_exists = True
# set files in folder flag by using the os.listdir function
files_in_folder = os.listdir(kwargs["folder_path"])
# set accessible flag by using the os.access function
accessible = os.access(kwargs["folder_path"], os.W_OK)
if not accessible:
raise ValueError("Folder for audio files is not accessible")
if kwargs["clear_folder"]:
# if the folder exists, delete all files in the folder
if folder_exists:
for file in files_in_folder:
# use shutil to remove the file
fl = os.path.join(kwargs["folder_path"], file)
if os.path.isfile(fl):
os.remove(fl)
# use shutil to remove the entire path
check_this_path = file
if len(check_this_path) > 3:
shutil.rmtree(check_this_path, ignore_errors=True)
else:
raise ValueError("Folder path is too short")
return (kwargs["folder_path"],)
@ETK_AVF_base
class DictToComfyAudio:
"""
Given a dictionary, creates a comfy audio node
"""
@classmethod
def INPUT_TYPES(cls):
ret = {"required": {"dict": ("DICT", {"default": {}})}}
return ret
RETURN_TYPES = ("AUDIO",)
FUNCTION = "dict_to_comfy_audio"
def dict_to_comfy_audio(self, **kwargs):
return (kwargs["dict"],)
@ETK_AVF_base
class LoadAudioFromServer:
"""
Given a path on the server loads the audio file from the server and returns it as an AUDIO node
"""
@classmethod
def INPUT_TYPES(cls):
ret = {"required": {"server_path": ("STRING", {"default": ""})}}
return ret
RETURN_TYPES = ("AUDIO",)
FUNCTION = "load_audio_from_server"
def load_audio_from_server(self, **kwargs):
file_on_server = kwargs["server_path"]
from comfy_extras.nodes_audio import LoadAudio
f = LoadAudio().load
r = f(file_on_server)[0]
return (r,)
@ETK_AVF_base
class AudioDefinitionProvider:
"""
Given a AUDIOFOLDER node, and inputs for sample rate, and channels, creates a AUDIOFOLDERDEF variable
"""
@classmethod
def INPUT_TYPES(cls):
ret = {
"required": {
"audio_folder": ("AUDIOFOLDER", {"default": None}),
"sample rate": ("INT", {"default": 44100}),
"channels": ("INT", {"default": 2}),
"audio format": (["WAV", "MP3"], {"default": "WAV"}),
},
}
return ret
RETURN_TYPES = ("AUDIO_FOLDER_DEF",)
FUNCTION = "audio_definition_provider"
def get_num_audio_files_that_match(self, audio_folder):
audio_files = self.get_all_audio_files(audio_folder)
return len(audio_files)
def get_all_audio_files(self, audio_folder):
import os
audio_files = os.listdir(audio_folder)
audio_files = sorted([os.path.join(audio_folder, x) for x in audio_files if "audio_" in x])
return audio_files
def audio_definition_provider(self, **kwargs):
from addict import Dict as adict
import os
ret = adict()
ret["audio_folder"] = kwargs["audio_folder"]
ret["sample rate"] = kwargs.get("sample rate") or kwargs.get("sample_rate")
ret["audio format"] = kwargs.get("audio format") or kwargs.get("audio_format")
ret["channels"] = kwargs.get("channels")
# a function that gets the frame count from the disk
ret.get_file_count = lambda: self.get_num_audio_files_that_match(ret["audio_folder"])
# a function that gets the next file name and full path dont forget to zfill
ret.get_next_file_name = lambda: f"{ret['audio_folder']}{os.sep}audio_{str(ret.get_file_count()).zfill(10)}.{ret['audio format'].lower()}"
ret.get_last_file_name = lambda: f"{ret['audio_folder']}{os.sep}audio_{str(ret.get_file_count() - 1).zfill(10)}.{ret['audio format'].lower()}"
ret.get_all_audio_files = lambda: self.get_all_audio_files(ret["audio_folder"])
return (ret,)
@ETK_AVF_base
class AddAudioToAudioFolder:
"""
Given a AUDIO_FOLDER_DEF node, and a list of audio files, adds the audio files to the audio folder
"""
@classmethod
def INPUT_TYPES(cls):
ret = {
"required": {
"audio_folder_def": ("AUDIO_FOLDER_DEF", {"default": None}),
"AUDIO": ("LIST", {"default": None}),
"delete after": ("BOOLEAN", {"default": False}),
},
}
return ret
RETURN_TYPES = ("AUDIO_FOLDER_DEF",)
FUNCTION = "add_audio_to_audio_folder"
def add_audio_to_audio_folder(self, **kwargs):
import os
audio_folder_def = kwargs["audio_folder_def"]
audio_files = kwargs["AUDIO"]
delete = kwargs["delete after"]
for audio_file in audio_files:
# use ffmpeg_exe to convert the audio file to the correct format which is audio_folder_def["audio format"]
# use the audio_folder_def["audio format"] to determine the output file name
# use the audio_folder_def["audio folder"] to determine the output folder
# use the audio_folder_def["sample rate"] to determine the sample rate
# use the audio_folder_def["channels"] to determine the number of channels
# get the output file name
output_file_name = audio_folder_def.get_next_file_name()
input_file_name = audio_file
CMD_AS_LIST = [
"-y",
"-i",
input_file_name,
"-ar",
str(audio_folder_def['sample rate']),
"-ac",
str(audio_folder_def['channels']),
output_file_name
]
run_ffmpeg_command(CMD_AS_LIST)
if delete:
os.remove(input_file_name)
return (audio_folder_def,)
@ETK_AVF_base
class AddSilenceToAudioFolder:
"""Create and add a new audio file that is silence for the specified duration"""
@classmethod
def INPUT_TYPES(cls):
ret = {
"required": {
"audio_folder_def": ("AUDIO_FOLDER_DEF", {"default": None}),
"duration": ("FLOAT", {"default": 1.0}),
},
}
return ret
RETURN_TYPES = ("AUDIO_FOLDER_DEF",)
FUNCTION = "add_silence_to_audio_folder"
def add_silence_to_audio_folder(self, **kwargs):
audio_folder_def = kwargs["audio_folder_def"]
duration = kwargs["duration"]
# get the output file name
output_file_name = audio_folder_def.get_next_file_name()
sample_rate = audio_folder_def["sample rate"]
channels = audio_folder_def["channels"]
fmt = audio_folder_def["audio format"]
# create the silence audio file
if fmt == "WAV":
fmt_part = "pcm_s16le"
elif fmt == "MP3":
fmt_part = "mp3"
else:
fmt_part = "pcm_s16le"
# Create the FFmpeg command for generating silence
CMD_AS_LIST = [
"-f",
"lavfi",
"-i",
f"aevalsrc=0:sample_rate={sample_rate}:channel_layout={channels},atrim=0:{duration}",
"-c:a",
fmt_part,
output_file_name
]
run_ffmpeg_command(CMD_AS_LIST)
return (audio_folder_def,)
@ETK_AVF_base
class GetAudioFilesLengthsFromAudioFolder:
"""Get the lengths of all audio files in a given audio folder"""
@classmethod
def INPUT_TYPES(cls):
return {"required": {"audio_folder_def": ("AUDIO_FOLDER_DEF", {"default": None})}}
RETURN_TYPES = ("AUDIO_FOLDER_DEF", "LIST",)
RETURN_NAMES = ("audio_folder_def", "[lens(seconds)]",)
FUNCTION = "get_audio_files_lengths"
def get_audio_files_lengths(self, audio_folder_def):
import os
import torchaudio
audio_folder = audio_folder_def["audio_folder"]
audio_files = os.listdir(audio_folder)
audio_files.sort()
audio_files_lengths = []
for audio_file in audio_files:
if audio_file[-3:].lower() not in ["wav", "mp3"]:
continue
audio_path = os.path.join(audio_folder, audio_file)
waveform, sample_rate = torchaudio.load(audio_path)
duration = waveform.shape[1] / sample_rate
audio_files_lengths.append(duration)
return (audio_folder_def, audio_files_lengths,)
@ETK_AVF_base
class AudioFolderRenderAudio:
"""Joins all the audio files"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"audio_folder_def": ("AUDIO_FOLDER_DEF", {"default": None}),
"output_path": ("STRING", {"multiline": False, "default": ""}),
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("output_path",)
FUNCTION = "join_audio_files"
CATEGORY = "ETK/audio"
def join_audio_files(self, **kwargs):
import librosa
import copy
import soundfile as sf
import numpy as np
audio_folder_def = kwargs['audio_folder_def']
kwargs = copy.deepcopy(kwargs)
paths = audio_folder_def.get_all_audio_files()
output_path = kwargs['output_path']
# Read the first file to get metadata
audio_data_list = []
for path in paths:
audio_data, sr = librosa.load(path, sr=None)
audio_data_list.append(audio_data)
# Concatenate all audio data
joined_audio_data = np.concatenate(audio_data_list)
# TODO: need to implement writing in the correct format
# Save the concatenated audio
sf.write(output_path, joined_audio_data, sr)
return (output_path,)
AVF/common.py
-45
View File
@@ -1,45 +0,0 @@
# common.py
from .. import NODE_CLASS_MAPPINGS
from .. import NODE_DISPLAY_NAME_MAPPINGS
def ETK_AVF_base(cls):
cls.CATEGORY = "ETK/AVF"
# Add spaces to the camel case class name
pretty_name = cls.__name__
for i in range(1, len(pretty_name)):
if pretty_name[i].isupper():
pretty_name = pretty_name[:i] + " " + pretty_name[i:]
cls.DISPLAY_NAME = pretty_name
NODE_CLASS_MAPPINGS[cls.__name__] = cls
NODE_DISPLAY_NAME_MAPPINGS[cls.DISPLAY_NAME] = pretty_name
return cls
import threading
def duplicate_file(src_file, dest_files):
import shutil
threads = []
for dest_file in dest_files:
thread = threading.Thread(target=shutil.copy2, args=(src_file, dest_file))
thread.start()
threads.append(thread)
for thread in threads:
thread.join()
def run_ffmpeg_command(args):
"""
Executes a command using the ffmpeg executable.
Parameters:
args (list): List of command-line arguments for ffmpeg.
"""
import subprocess
import imageio_ffmpeg as ffmpeg
if isinstance(args, str):
args = args.split()
cmd = [ffmpeg.get_ffmpeg_exe()] + args
subprocess.run(cmd)
AVF/video_nodes.py
-695
View File
@@ -1,695 +0,0 @@
# video_module.py
from .common import ETK_AVF_base, duplicate_file, run_ffmpeg_command
@ETK_AVF_base
class VideoFolderFFMPEGVideoInterpolate:
"""
Use run_ffmpeg to interpolate the frames in the video folder
"""
@classmethod
def INPUT_TYPES(cls):
ret = {
"required": {
"video_folder_def": ("VIDEO_FOLDER_DEF", {"default": None}),
"last_n_frames": ("INT", {"default": 2}),
"target_n_frames": ("INT", {"default": 3}),
},
}
return ret
RETURN_TYPES = ("VIDEO_FOLDER_DEF",)
RETURN_NAMES = ("video_folder_def",)
FUNCTION = "video_folder_ffmpeg_video_interpolate"
def video_folder_ffmpeg_video_interpolate(self, **kwargs):
import os
import shutil
video_folder_def = kwargs["video_folder_def"]
last_n_frames = kwargs["last_n_frames"]
target_n_frames = kwargs["target_n_frames"]
# Get the current frame count
current_frame_count = video_folder_def.get_frame_count()
# Calculate the target frame rate for interpolation
target_fps = video_folder_def["fps"] * (target_n_frames / last_n_frames)
source_fps = video_folder_def["fps"]
# Get the file names for the last n frames
last_frame_names = [video_folder_def.get_last_file_name(offset=i) for i in range(last_n_frames)]
last_frame_names = list(set(last_frame_names))
last_frame_names.sort()
# Create a temporary folder off of the video frame folder's base path
temp_folder = os.path.join(video_folder_def["video_folder"], "temp")
os.makedirs(temp_folder, exist_ok=True)
# Move the input frames into the new folder
for frame_name in last_frame_names:
# Do this in a way that overwrites the files and MOVES them
shutil.move(frame_name, os.path.join(temp_folder, os.path.basename(frame_name)))
# New input frame names list
last_frame_names = [os.path.join(temp_folder, os.path.basename(frame_name)) for frame_name in last_frame_names]
last_frame_names.sort()
# Create a file with the list of input frame paths in the temporary folder
input_frame_list_file = os.path.join(temp_folder, "input_frames.txt")
input_frame_list_file = os.path.abspath(input_frame_list_file)
with open(input_frame_list_file, 'w') as file:
# Need to do this in a manner that puts the first frame at the top of the file.
file.write('\n'.join([f"file '{frame_name}'" for frame_name in last_frame_names]))
# Get the output frame names for the interpolated frames
output_frame_names = [video_folder_def.get_next_file_name(i) for i in range(target_n_frames)]
temp_output_frame_pattern = os.path.join(video_folder_def["video_folder"], "temp",
"tmp_frame_%010d." + video_folder_def["frame format"].lower())
temp_output_frame_pattern = os.path.abspath(temp_output_frame_pattern)
# Create the FFmpeg command for motion interpolation using the frame list file
ffmpeg_interpolate_cmd = [
"-f", "concat", "-safe", "0",
"-r", str(source_fps),
"-i", input_frame_list_file,
"-vsync", "1",
"-filter_complex",
f"[0:v]setpts=N/TB,minterpolate=fps={target_fps}:mi_mode=blend,setpts=PTS-STARTPTS[v]", # Reset timestamps
"-map", "[v]", "-q:v", "2",
"-r", str(target_fps),
"-frames:v", str(target_n_frames),
"-threads", "0",
"-y", temp_output_frame_pattern
]
run_ffmpeg_command(ffmpeg_interpolate_cmd)
# Move the interpolated frames to the video folder with the correct names
for i, frame_name in enumerate(output_frame_names):
temp_frame_name = os.path.join(temp_folder,
f"tmp_frame_{i + 1:010d}.{video_folder_def['frame format'].lower()}")
try:
shutil.move(temp_frame_name, frame_name)
except FileNotFoundError:
pass
# Remove the temporary folder and frame list file
os.remove(input_frame_list_file)
# Remove this without generating an error even if it is not empty
try:
os.rmdir(temp_folder)
except OSError:
shutil.rmtree(temp_folder)
return (video_folder_def,)
@ETK_AVF_base
class VideoDefinitionProvider:
"""
Given a VIDEOFOLDER node, and inputs for height, width, and fps, creates a VIDEOFOLDERDEF variable
"""
@classmethod
def INPUT_TYPES(cls):
ret = {
"required": {
"video_folder": ("VIDEOFOLDER", {"default": None}),
"height": ("INT", {"default": 1080}),
"width": ("INT", {"default": 1920}),
"fps": ("FLOAT", {"default": 30}),
"frame format": (["JPG", "PNG"], {"default": "JPG"}),
},
}
return ret
RETURN_TYPES = ("VIDEO_FOLDER_DEF",)
FUNCTION = "video_definition_provider"
def video_definition_provider(self, **kwargs):
from addict import Dict as adict
import os
ret = adict()
ret["video_folder"] = kwargs["video_folder"]
ret["height"] = kwargs["height"]
ret["width"] = kwargs["width"]
ret["fps"] = kwargs["fps"]
ret["frame format"] = kwargs.get("frame format") or kwargs.get("frame_format")
# A function that gets the frame count from the disk
ret.get_frame_count = lambda: len([f for f in os.listdir(ret["video_folder"]) if "frame_" in f])
# A function that gets the next file name and full path, don't forget to add 10 zeros to the frame count
ret.get_next_file_name = lambda x=0: os.path.join(
ret['video_folder'],
f"frame_{str(ret.get_frame_count() + x).zfill(10)}.{ret['frame format'].lower()}"
)
ret.get_last_file_name = lambda offset=0: os.path.join(
ret['video_folder'],
f"frame_{str(max(ret.get_frame_count() - offset - 1, 0)).zfill(10)}.{ret['frame format'].lower()}"
)
ret.get_current_file_name = lambda: os.path.join(
ret['video_folder'],
f"frame_{str(ret.get_frame_count() - 1).zfill(10)}.{ret['frame format'].lower()}"
)
def reserve_next_frame():
fn = ret.get_next_file_name()
open(fn, "w").close()
return fn
ret.reserve_next_frame = reserve_next_frame
return (ret,)
@ETK_AVF_base
class VideoFolderProvider:
"""
Given a folder path, creates the folder or empties it entirely returns the folder path as
a string but as the node type "VIDEOFOLDER"
"""
@classmethod
def INPUT_TYPES(cls):
ret = {
"required": {
"folder_path": ("STRING", {"default": ""}),
"clear_folder": ("BOOLEAN", {"default": True}),
},
"optional": {
"activator:": ("*", {"default": None})
}
}
return ret
RETURN_TYPES = ("VIDEOFOLDER",)
FUNCTION = "video_folder_provider"
def video_folder_provider(self, **kwargs):
import shutil
import os
# Check if the folder_path is none
if kwargs["folder_path"] is None:
raise ValueError("folder_path must be specified")
# Set folder exists flag by using the os.path.exists function
folder_exists = os.path.exists(kwargs["folder_path"])
if not folder_exists:
os.makedirs(kwargs["folder_path"])
folder_exists = True
# Set files in folder flag by using the os.listdir function
files_in_folder = os.listdir(kwargs["folder_path"])
# Set accessible flag by using the os.access function
accessible = os.access(kwargs["folder_path"], os.W_OK)
if not accessible:
raise ValueError("Folder for video files is not accessible")
if kwargs["clear_folder"]:
# If the folder exists, delete all files in the folder
if folder_exists:
for file in files_in_folder:
# Use shutil to remove the file
fl = os.path.join(kwargs["folder_path"], file)
if os.path.isfile(fl):
os.remove(fl)
# Use shutil to remove the entire path
check_this_path = file
if len(check_this_path) > 3:
shutil.rmtree(check_this_path, ignore_errors=True)
else:
raise ValueError("Folder path is too short")
return (kwargs["folder_path"],)
@ETK_AVF_base
class AddFramesToVideoFolder:
"""
Given a VIDEO_FOLDER_DEF node, and a list of frames, adds the frames to the video folder
"""
@classmethod
def INPUT_TYPES(cls):
ret = {
"required": {
"video_folder_def": ("VIDEO_FOLDER_DEF", {"default": None}),
"IMAGE": ("IMAGE", {"default": None}),
"Rescale method": (
["NEAREST", "BOX", "BILINEAR", "HAMMING", "BICUBIC", "LANCZOS"],
{"default": "LANCZOS"}),
"Release vram?": ("BOOLEAN", {"default": True}),
},
}
return ret
RETURN_TYPES = ("VIDEO_FOLDER_DEF",)
FUNCTION = "add_frames_to_video_folder"
@staticmethod
def save_or_copy_frame(current_frame, previous_frame, frame_fp, previous_frame_fp):
import shutil
from PIL import Image
import torch
if previous_frame is None or not torch.equal(current_frame, previous_frame):
# If the current frame is different from the previous frame, save it to disk
current_frame = current_frame
# expected shape is H, W, C uint8 0-255
# image.fromarray expects uint8 0-255 and shape H, W, C
needs_dtype_conversion = current_frame.dtype != torch.uint8
if len(current_frame.shape) == 3 and current_frame.dtype == torch.uint8:
needs_dtype_conversion = False
elif len(current_frame.shape) == 2:
# assume grayscale H,W
current_frame = current_frame.unsqueeze(2).repeat(1, 1, 3)
elif len(current_frame.shape) == 4 and current_frame.shape[0] == 1:
# assume B, H, W, C
current_frame = current_frame.squeeze(0)
elif len(current_frame.shape) == 3 and current_frame.shape[-1] == 4:
# assume H, W, C with alpha
current_frame = current_frame[..., :3]
else: # try to move forward
pass
# assume H, W, C
if needs_dtype_conversion:
current_frame = current_frame.mul(255).byte()
ndarr = current_frame.cpu().numpy()
im = Image.fromarray(ndarr)
im.save(frame_fp)
previous_frame = current_frame
else:
# If the current frame is the same as the previous frame, copy the previously saved frame file
shutil.copy2(previous_frame_fp, frame_fp)
return previous_frame
@staticmethod
def process_frames(frames, rescale_method, video_folder_def):
import gc
import torch
num_frames = frames.shape[0]
previous_frame = None
frame_fp = None
frames_are_all_same = torch.all(frames[0] == frames)
if frames_are_all_same and frames.shape[0] > 1:
# Get a list of file names using reserve_next_frame
frame_f_names = [video_folder_def.reserve_next_frame() for _ in range(num_frames)]
# Save the first frame
AddFramesToVideoFolder.save_or_copy_frame(frames[0], previous_frame, frame_f_names[0], None)
# Copy the first frame to the rest of the frames using os by calling the cli
duplicate_file(frame_f_names[0], frame_f_names[1:])
else:
for i in range(num_frames):
next_fn = video_folder_def.get_next_file_name()
previous_frame_fp = frame_fp
frame_fp = next_fn
current_frame = frames[i]
previous_frame = AddFramesToVideoFolder.save_or_copy_frame(current_frame, previous_frame, frame_fp,
previous_frame_fp)
del frames
gc.collect()
return None
def add_frames_to_video_folder(self, **kwargs):
from PIL import Image
import torch
import gc
from torchvision.transforms.functional import resize
import psutil
from custom_nodes.EternalKernelLiteGraphNodes.modules.image_utils import lanczos_resize
initial_memory = psutil.virtual_memory().available
print(f"Initial free memory: {initial_memory / (1024 * 1024):.2f} MB")
# Get the video folder path
video_folder = kwargs["video_folder_def"]["video_folder"]
# Get the frame format
frame_format = kwargs["video_folder_def"]["frame format"]
# Get the height
height = kwargs["video_folder_def"]["height"]
# Get the width
width = kwargs["video_folder_def"]["width"]
# Get the fps
fps = kwargs["video_folder_def"]["fps"]
# Get the frames
frames = kwargs["IMAGE"]
# Get the rescale method
rescale_method = kwargs.get("Rescale method", None)
if not rescale_method:
rescale_method = kwargs.get("Rescale_method", None)
free_vram = kwargs.get("Release vram?", None)
if not free_vram:
free_vram = kwargs.get("Release_vram_", True)
orig_frames = frames
# num_files = len(glob(os.path.join(video_folder, f"frame_*.{frame_format.lower()}")))
# Rescale if needed to the correct size
if height != frames.shape[1] or width != frames.shape[2]:
# First since transforms.Resize requires a tensor of shape B,C, H, W
# We need to permute the tensor to B, C, H, W
frames = frames.permute(0, 3, 1, 2)
frames = frames.to("cuda")
# Get the right value for interpolation
interpolation_methods = {
"NEAREST": Image.Resampling.NEAREST,
"BOX": Image.Resampling.BOX,
"BILINEAR": Image.Resampling.BILINEAR,
"HAMMING": Image.Resampling.HAMMING,
"BICUBIC": Image.Resampling.BICUBIC,
"LANCZOS": Image.Resampling.LANCZOS
}
rescale_type = interpolation_methods[rescale_method]
# rescale_method = transforms.Resize((height, width), interpolation=rescale_type) # this doesnt work so well lets use a different function
if rescale_type != Image.Resampling.LANCZOS:
rescale_method_fn = lambda x: resize(x, (height, width), interpolation=rescale_type)
frames = rescale_method_fn(frames)
else:
rescale_method_fn = lambda x: lanczos_resize(x, (width, height))
frames = rescale_method_fn(frames)
# Now we need to permute back to B, H, W, C
frames = frames.permute(0, 2, 3, 1)
self.process_frames(frames, rescale_method_fn, kwargs["video_folder_def"])
else:
self.process_frames(frames, rescale_method, kwargs["video_folder_def"])
gc.collect()
if free_vram:
import gc
if isinstance(frames, torch.Tensor):
frames.to("cpu").detach().numpy()
del frames
if isinstance(orig_frames, torch.Tensor):
orig_frames = orig_frames.to("cpu").detach().numpy()
del orig_frames
torch.cuda.empty_cache()
gc.collect()
initial_memory = psutil.virtual_memory().available
print(f"After adding frames free memory: {initial_memory / (1024 * 1024):.2f} MB")
# Return the VIDEO_FOLDER_DEF
return (kwargs["video_folder_def"],)
@ETK_AVF_base
class HoldFrameForSecsInVideoFolder:
"""
Given a VIDEO_FOLDER_DEF node, and a list of frames, first uses RepeatFrames to repeat the frames
for the correct amount and then adds the frames to the video folder using AddFramesToVideoFolder
"""
@classmethod
def INPUT_TYPES(cls):
ret = {
"required": {
"video_folder_def": ("VIDEO_FOLDER_DEF", {"default": None}),
"IMAGE": ("IMAGE", {"default": None}),
"seconds": ("FLOAT", {"default": 1.0}),
"Rescale method": (
["NEAREST", "BOX", "BILINEAR", "HAMMING", "BICUBIC", "LANCZOS"],
{"default": "LANCZOS"}),
"Release vram?": ("BOOLEAN", {"default": True}),
},
}
return ret
RETURN_TYPES = ("VIDEO_FOLDER_DEF",)
FUNCTION = "hold_frame_for_secs_in_video_folder"
def saver(self, x, kwargs):
# This just saves in an async way
rescale_method = kwargs.get("Rescale method", None)
if rescale_method is None:
rescale_method = kwargs.get("Rescale_method", None)
rel_vram = kwargs.get("Release vram?", None)
if rel_vram is None:
rel_vram = kwargs.get("Release_vram_", True)
AddFramesToVideoFolder().add_frames_to_video_folder(
video_folder_def=kwargs["video_folder_def"],
IMAGE=x,
Rescale_method=rescale_method,
Release_vram_=rel_vram
)
def hold_frame_for_secs_in_video_folder(self, **kwargs):
video_folder_ref = kwargs["video_folder_def"]
video_folder = kwargs["video_folder_def"]["video_folder"]
# Get the frame format
frame_format = kwargs["video_folder_def"]["frame format"]
# Get the height
height = kwargs["video_folder_def"]["height"]
# Get the width
width = kwargs["video_folder_def"]["width"]
# Get the fps
fps = kwargs["video_folder_def"]["fps"]
# Get the frames
frames = kwargs["IMAGE"]
# Get the number of frames
num_frames = frames.shape[0]
# Get the rescale method
rescale_method = kwargs.get("Rescale method", None)
if rescale_method is None:
rescale_method = kwargs.get("Rescale_method", None)
# Get the number of seconds
seconds = kwargs["seconds"]
# Get the number of frames to hold
num_frames_to_hold = int(fps * seconds)
num_frames_in_this_stack = frames.shape[0]
# Get the number of repeats needed
num_repeats_needed = num_frames_to_hold // num_frames_in_this_stack
# Get the number of frames left over
num_frames_left_over = num_frames_to_hold % num_frames_in_this_stack
# Write the first set of frames to disk
self.saver(frames, kwargs)
# Get the name of the last frame
source_last_frame_name = video_folder_ref.get_last_file_name()
# From the last frame find the rest of the frames names/paths
source_last_frame_num = int(source_last_frame_name.split("_")[-1].split(".")[0])
source_first_frame_num = source_last_frame_num - num_frames_in_this_stack + 1
template = source_last_frame_name.replace(str(source_last_frame_num).zfill(10), "{}")
range_of_source_frame_numbers = list(range(source_first_frame_num, source_last_frame_num + 1))
source_frame_names = [template.format(str(x).zfill(10)) for x in range_of_source_frame_numbers]
target_done_frame_num = source_last_frame_num + num_frames_to_hold - num_frames_in_this_stack
num_frames_remaining = num_frames_to_hold - num_frames_in_this_stack
target_start_frame_num = source_last_frame_num + 1
target_frame_range = list(range(target_start_frame_num, target_done_frame_num + 1))
target_frame_names = [template.format(str(x).zfill(10)) for x in target_frame_range]
set_todo = {fn: [] for fn in source_frame_names}
# We can now repeat the frames
for i, target_frame_name in enumerate(target_frame_names):
# Keep track of what stack frame we are on
stack_frame_idx = i % num_frames_in_this_stack
stack_frame_name_to_use = source_frame_names[stack_frame_idx]
set_todo[stack_frame_name_to_use].append(target_frame_name)
# Now we can repeat the frames
for frame_name in set_todo:
src_frame = frame_name
target_frames = set_todo[src_frame]
duplicate_file(src_frame, target_frames)
return (kwargs["video_folder_def"],)
@ETK_AVF_base
class Listed_HoldFramesForSecsInVideoFolder:
"""This just calls HoldFrameForSecsInVideoFolder multiple times for a list of frames and a list of durations"""
@classmethod
def INPUT_TYPES(cls):
ret = {
"required": {
"video_folder_def": ("VIDEO_FOLDER_DEF", {"default": None}),
"IMAGE": ("LIST", {"default": None}),
"seconds": ("LIST", {"default": None}),
"Rescale method": (
["NEAREST", "BOX", "BILINEAR", "HAMMING", "BICUBIC", "LANCZOS"],
{"default": "LANCZOS"}),
"Release vram?": ("BOOLEAN", {"default": True}),
},
}
return ret
RETURN_TYPES = ("VIDEO_FOLDER_DEF",)
FUNCTION = "listed_hold_frames_for_secs_in_video_folder"
def listed_hold_frames_for_secs_in_video_folder(self, **kwargs):
# The plan is to call HoldFrameForSecsInVideoFolder for each frame and duration
# and then return the video_folder_def
if "IMAGE" not in kwargs or "seconds" not in kwargs:
return (kwargs["video_folder_def"],)
for frame, duration in zip(kwargs["IMAGE"], kwargs["seconds"]):
HoldFrameForSecsInVideoFolder().hold_frame_for_secs_in_video_folder(
video_folder_def=kwargs["video_folder_def"],
IMAGE=frame,
seconds=duration,
Rescale_method=kwargs["Rescale method"],
Release_vram_=kwargs["Release vram?"])
return (kwargs["video_folder_def"],)
@ETK_AVF_base
class VideoFolderDefInfo:
"""
Returns everything in the VIDEO_FOLDER_DEF, and also returns information about the status of the folder
"""
@classmethod
def INPUT_TYPES(cls):
ret = {
"required": {
"video_folder_def": ("VIDEO_FOLDER_DEF", {"default": None}),
},
}
return ret
RETURN_TYPES = ("VIDEO_FOLDER_DEF", "STRING", "INT", "INT", "FLOAT", "STRING", "INT", "LIST",)
RETURN_NAMES = ("VFD", "video_folder", "height", "width", "fps", "frame format", "num_files", "list of files")
OUTPUT_NODE = True
FUNCTION = "video_folder_def_info"
def video_folder_def_info(self, **kwargs):
import os
video_folder = kwargs["video_folder_def"]["video_folder"]
height = kwargs["video_folder_def"]["height"]
width = kwargs["video_folder_def"]["width"]
fps = kwargs["video_folder_def"]["fps"]
frame_format = kwargs["video_folder_def"]["frame format"]
files = os.listdir(video_folder)
num_files = len(files)
return (kwargs["video_folder_def"], video_folder, height, width, fps, frame_format, num_files, files,)
@ETK_AVF_base
class FadeToInVideoFolder:
"""
Given a VIDEO_FOLDER_DEF node, and a stack of frames, fades the frames in and adds them to the video folder
"""
@classmethod
def INPUT_TYPES(cls):
ret = {
"required": {
"video_folder_def": ("VIDEO_FOLDER_DEF", {"default": None}),
"IMAGE": ("IMAGE", {"default": None}),
"seconds": ("FLOAT", {"default": 1.0}),
"transition ms": ("FLOAT", {"default": 0.1}),
"Rescale method": (
["NEAREST", "BOX", "BILINEAR", "HAMMING", "BICUBIC", "LANCZOS"],
{"default": "LANCZOS"}),
"Release vram?": ("BOOLEAN", {"default": True}),
},
}
return ret
RETURN_TYPES = ("VIDEO_FOLDER_DEF",)
FUNCTION = "fade_to_in_video_folder"
def fade_to_in_video_folder(self, **kwargs):
from PIL import Image
import torch
# Use HoldFrameForSecsInVideoFolder to hold the frame after creating and AddFramesToVideoFolder the transition
transition_length_secs = int(kwargs["transition ms"] * kwargs["video_folder_def"]["fps"]) / 1000
hold_frame_length_secs = int(kwargs["seconds"] * kwargs["video_folder_def"]["fps"]) - transition_length_secs
transition_frame_count = int(kwargs["video_folder_def"]["fps"] * transition_length_secs)
base_frame_fn = kwargs["video_folder_def"].get_last_file_name()
# Load it into a pytorch tensor (B, C, H, W)
base_frame = Image.open(base_frame_fn)
base_frame = torch.tensor(base_frame).permute(2, 0, 1).unsqueeze(0).to("cuda")
transition_stack = torch.zeros(
(transition_frame_count, base_frame.shape[1], base_frame.shape[2], base_frame.shape[3]),
dtype=torch.float16, device="cuda"
)
# Fade from base_frame to the first frame in the stack over the transition length
new_frame = kwargs["IMAGE"][0].to("cuda")
for i in range(transition_frame_count):
alpha = i / transition_frame_count
interpolated_frame = (1 - alpha) * base_frame + alpha * new_frame
interpolated_frame = interpolated_frame.to("cpu")
transition_stack[i] = interpolated_frame
# Add to the video folder
AddFramesToVideoFolder().add_frames_to_video_folder(
video_folder_def=kwargs["video_folder_def"],
IMAGE=transition_stack,
Rescale_method=kwargs["Rescale method"],
Release_vram_=kwargs["Release vram?"]
)
# Hold the frame
HoldFrameForSecsInVideoFolder().hold_frame_for_secs_in_video_folder(
video_folder_def=kwargs["video_folder_def"],
IMAGE=kwargs["IMAGE"],
seconds=hold_frame_length_secs,
Rescale_method=kwargs["Rescale method"],
Release_vram_=kwargs["Release vram?"]
)
return (kwargs["video_folder_def"],)
__init__.py
+2 -30
View File
@@ -2,12 +2,8 @@ NODE_CLASS_MAPPINGS = {}
NODE_DISPLAY_NAME_MAPPINGS = {}
WEB_DIRECTORY = "./web"
# In /custom_nodes/ETK/__init__.py
import sys
from types import ModuleType
import os
import importlib.util
def patch_validate_inputs():
@@ -46,32 +42,8 @@ def patch_validate_inputs():
# Apply the patch when this module is imported
patch_validate_inputs()
UNIT_MODE = os.environ.get("UNIT_TEST") or importlib.util.find_spec('comfy') is None
if not UNIT_MODE:
from . import AVF
# Update NODE_CLASS_MAPPINGS with AVF nodes
NODE_CLASS_MAPPINGS.update(AVF.NODE_CLASS_MAPPINGS)
# Update NODE_DISPLAY_NAME_MAPPINGS with AVF nodes
NODE_DISPLAY_NAME_MAPPINGS.update(AVF.NODE_DISPLAY_NAME_MAPPINGS)
from . import basic
from . import image
from . import video
from . import functional
from . import audio
from . import git
from . import server_endpoints
from . import torchvision_nodes
from . import utils
from . import pytorch_nodes
from . import youtube_nodes
from . import torchvision_nodes
from . import hf_diffusers_nodes
#from . import audio_video_folder
# Import PyTorch nodes to populate node mappings
from . import pytorch_nodes
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
audio.py
-742
View File
@@ -1,742 +0,0 @@
import shutil
import struct
import io
import copy
import imageio_ffmpeg as ffmpeg
import joblib
from .config import config_settings
temp_dir = config_settings["tmp_dir"]
from . import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
# NODE_CLASS_MAPPINGS = {}
# NODE_DISPLAY_NAME_MAPPINGS = {}
def create_vorbis_comment_block(comment_dict, last_block):
vendor_string = b'ComfyUI'
vendor_length = len(vendor_string)
comments = []
for key, value in comment_dict.items():
comment = f"{key}={value}".encode('utf-8')
comments.append(struct.pack('<I', len(comment)) + comment)
user_comment_list_length = len(comments)
user_comments = b''.join(comments)
comment_data = struct.pack('<I', vendor_length) + vendor_string + struct.pack('<I', user_comment_list_length) + user_comments
if last_block:
id = b'\x84'
else:
id = b'\x04'
comment_block = id + struct.pack('>I', len(comment_data))[1:] + comment_data
return comment_block
# Create a memory object for caching
# memory = joblib.Memory(temp_dir, verbose=0)
def insert_or_replace_vorbis_comment(flac_io, comment_dict):
if len(comment_dict) == 0:
return flac_io
flac_io.seek(4)
blocks = []
last_block = False
while not last_block:
header = flac_io.read(4)
last_block = (header[0] & 0x80) != 0
block_type = header[0] & 0x7F
block_length = struct.unpack('>I', b'\x00' + header[1:])[0]
block_data = flac_io.read(block_length)
if block_type == 4 or block_type == 1:
pass
else:
header = bytes([(header[0] & (~0x80))]) + header[1:]
blocks.append(header + block_data)
blocks.append(create_vorbis_comment_block(comment_dict, last_block=True))
new_flac_io = io.BytesIO()
new_flac_io.write(b'fLaC')
for block in blocks:
new_flac_io.write(block)
new_flac_io.write(flac_io.read())
return new_flac_io
def ETK_audio_base(cls):
cls.CATEGORY = "ETK/audio"
NODE_CLASS_MAPPINGS[cls.__name__] = cls
return cls
@ETK_audio_base
class GetAudioMetaData:
""" Returns a dictionary of audio metadata """
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"path": ("STRING", {"multiline": False, "default": ""}),
}
}
RETURN_TYPES = ("DICT",)
FUNCTION = 'get_audio_metadata'
CATEGORY = 'ETK/audio'
def get_audio_metadata(self, **kwargs):
import audioread
kwargs = copy.deepcopy(kwargs)
path = kwargs.get('path')
metadata = {}
with audioread.audio_open(path) as f:
metadata['duration'] = f.duration
metadata['samplerate'] = f.samplerate
metadata['channels'] = f.channels
return (metadata,)
@ETK_audio_base
class AudioFileListFromPathPattern:
""" Returns a list of audio files from a path pattern """
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"path": ("STRING", {"multiline": False, "default": ""}),
}
}
RETURN_TYPES = ("LIST", "LIST", "LIST",)
RETURN_NAMES = ("list_of_fpaths", "list_of_bytes", "list_of_metadata")
FUNCTION = 'get_audio_file_list_from_path_pattern'
CATEGORY = 'ETK/audio'
def get_audio_file_list_from_path_pattern(self, **kwargs):
import glob
kwargs = copy.deepcopy(kwargs)
path = kwargs.get('path')
l_fpaths = glob.glob(path)
l_bytes = []
l_metadata = []
for fpath in l_fpaths:
if fpath[-4:] in [".wav", ".mp3", ".flac", ".ogg", ".m4a", ".aac", ".wma", ".aiff", ".alac", ".au", ".raw",
".mp2", ]:
with open(fpath, "rb") as f:
l_bytes.append(f.read())
l_metadata.append(GetAudioMetaData().get_audio_metadata(path=fpath)[0])
return (l_fpaths, l_bytes, l_metadata,)
import numpy as np
@ETK_audio_base
class JoinAudioFiles:
""" Joins a list of audio files into a single audio file """
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"paths": ("LIST", {"default": []}),
"output_path": ("STRING", {"multiline": False, "default": ""}),
}
}
RETURN_TYPES = ("STRING", "BYTES", "DICT",)
FUNCTION = 'join_audio_files'
CATEGORY = 'ETK/audio'
def join_audio_files(self, **kwargs):
import librosa
import soundfile as sf
kwargs = copy.deepcopy(kwargs)
paths = kwargs['paths']
output_path = kwargs['output_path']
# Read the first file to get metadata
audio_data_list = []
for path in paths:
audio_data, sr = librosa.load(path, sr=None)
audio_data_list.append(audio_data)
# Concatenate all audio data
joined_audio_data = np.concatenate(audio_data_list)
# Save the concatenated audio
sf.write(output_path, joined_audio_data, sr)
# Get metadata
metadata = {
'duration': librosa.get_duration(y=joined_audio_data, sr=sr),
'samplerate': sr
}
# Get bytes
with open(output_path, "rb") as f:
file_bytes = f.read()
return (output_path, file_bytes, metadata,)
@ETK_audio_base
class ExtractAudioFromVideo:
""" Extracts audio from a video file and saves it as an audio file with a specified format and bitrate """
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"video_path": ("STRING", {"multiline": False, "default": ""}),
"format": (["wav", "mp3"], {"default": "wav"}),
"output_path": ("STRING", {"multiline": False, "default": ""}),
},
"optional": {
"bitrate (MP3)": ("INT", {"default": 192}) # Default bitrate value (change as needed)
}
}
RETURN_TYPES = ("STRING", "BYTES")
RETURN_NAMES = ("output_wave_file_path", "audio_bytes")
FUNCTION = 'extract_audio_from_video'
CATEGORY = 'ETK/audio'
def extract_audio_from_video(self, **kwargs):
import moviepy.editor as mp
import io
kwargs = copy.deepcopy(kwargs)
video_path = kwargs.get('video_path')
output_path = kwargs.get('output_path', "tmp.wav")
audio_format = kwargs.get('format', 'wav')
bitrate = kwargs.get('bitrate', None) # Default bitrate is None
# Load the video clip
video_clip = mp.VideoFileClip(video_path)
# Determine the appropriate codec based on the audio format
if audio_format == 'wav':
codec = "pcm_s16le" # For WAV format
elif audio_format == 'mp3':
codec = "mp3" # For MP3 format
# Extract audio and save it with the specified format and bitrate
audio_clip = video_clip.audio
audio_clip.write_audiofile(output_path, codec=codec, bitrate=bitrate)
# Read the generated audio file as bytes
audio_bytes = None
with open(output_path, 'rb') as audio_file:
audio_bytes = audio_file.read()
return (output_path, audio_bytes,)
# import torchaudio
# from audiocraft.models import MusicGen
# from audiocraft.data.audio import audio_write
#
# model = MusicGen.get_pretrained('melody')
# model.set_generation_params(duration=8) # generate 8 seconds.
#
# descriptions = ['happy rock', 'energetic EDM', 'sad jazz']
#
# melody, sr = torchaudio.load('./assets/bach.mp3')
# # generates using the melody from the given audio and the provided descriptions.
# wav = model.generate_with_chroma(descriptions, melody[None].expand(3, -1, -1), sr)
#
# for idx, one_wav in enumerate(wav):
# # Will save under {idx}.wav, with loudness normalization at -14 db LUFS.
# audio_write(f'{idx}', one_wav.cpu(), model.sample_rate, strategy="loudness")
@ETK_audio_base
class GenerateMusic:
""" Generates music from a given melody and descriptions """
@classmethod
def INPUT_TYPES(cls):
return {
"optional": {
"description": ("STRING", {"default": ""}),
"[description]": ("LIST", {"default": None}),
},
"required": {
"model_name": (["facebook/musicgen-melody-large",
"facebook/musicgen-melody-small",
"facebook/musicgen-melody-medium",
"facebook/musicgen-stereo-melody-large",
"facebook/musicgen-stereo-melody-small",
"facebook/musicgen-stereo-melody-medium",
"facebook/musicgen-large",
"facebook/musicgen-small",
"facebook/musicgen-medium",
"facebook/musicgen-stereo-large",
"facebook/musicgen-stereo-small",
"facebook/musicgen-stereo-medium",
],),
"duration": ("INT", {"default": None}),
"[duration]": ("LIST", {"default": None}),
"file": ("STRING", {"default": ""}),
"seed": ("INT", {"default": -1}),
}
}
RETURN_TYPES = ("STRING", "LIST",)
RETURN_NAMES = ("file", "files_list",)
FUNCTION = 'generate_music'
CATEGORY = 'ETK/audio'
def create_command(self, kwargs):
import os
musicgencli_path = os.path.join(os.path.dirname(__file__), "modules", "musicgencli.py")
ukwargs = copy.deepcopy(kwargs)
nl = "\n"
ukwargs["description"] = kwargs.get("[description]", [])
ukwargs.pop("[description]")
ukwargs["duration"] = kwargs.get("[duration]", [])
ukwargs.pop("[duration]")
ukwargs["description"] = str([i.replace("'", "") for i in ukwargs["description"]])
command = " ".join([f"--{k} \"{str(v).replace(nl, ' ')}\"" for k, v in ukwargs.items()])
command = f"python {musicgencli_path} {command}"
return command
def generate_music(self, **kwargs):
# use the cli to call the cli music generator at ./modules/musicgencli.py
import os
import subprocess
command = self.create_command(kwargs)
# Execute the command caputring the output
try:
output = subprocess.check_output(command, shell=True)
except subprocess.CalledProcessError as e:
print(f"Error: {e}")
# use the os to find the new files
new_files = []
for file in os.listdir(os.path.dirname(kwargs['file'])):
if file.endswith(".wav"):
new_files.append(file)
return (kwargs['file'], new_files,)
def generate_music_old(self, **kwargs):
from audiocraft.models import MusicGen
from audiocraft.data.audio import audio_write
import torch
torch.cuda.empty_cache()
import random
# Set seed, -1 means random seed
seed = kwargs.get('seed', -1)
if kwargs.get('seed') == -1:
seed = random.randint(0, 100000)
old_seed_values = set_all_seeds(seed)
kwargs = copy.deepcopy(kwargs)
description = kwargs.get('description')
descriptions = kwargs.get('[description]')
if description and not descriptions:
descriptions = [description]
file = kwargs.get('file')
# add .wave if not there
if not file.endswith('.wav'):
file = file + ".wav"
model_name = kwargs.get('model_name', 'melody')
duration = kwargs.get('duration', 10)
durations = kwargs.get('[duration]', None)
if duration:
durations = [duration]
model = MusicGen.get_pretrained(model_name)
o_files = []
for i, description in enumerate(descriptions):
ii = str(i).zfill(1 + len(str(1 + len(descriptions))))
file_to_use = file.replace(".wav", f"_{ii}.wav")
duration_to_use = durations[i]
model.set_generation_params(duration=duration_to_use)
wav = model.generate([description])
# q = MusicGen.set_generation_params()
audio_write(file_to_use.replace(".wav", ""), wav[0].cpu(), model.sample_rate, strategy="loudness")
o_files.append(file_to_use)
# Reset seed
reset_seeds(old_seed_values)
torch.cuda.empty_cache()
del model
torch.cuda.empty_cache()
return (o_files[0], o_files,)
def set_all_seeds(seed):
import random
import os
import numpy as np
import torch
old_values = {}
old_values['random'] = random.getstate()
old_values['numpy'] = np.random.get_state()
old_values['torch'] = torch.get_rng_state()
old_values['torch_cuda'] = torch.cuda.get_rng_state()
old_values['torch_cuda_all'] = torch.cuda.get_rng_state_all()
old_values['torch_deterministic'] = torch.backends.cudnn.deterministic
# From https://gist.github.com/gatheluck/c57e2a40e3122028ceaecc3cb0d152ac
random.seed(seed)
os.environ["PYTHONHASHSEED"] = str(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.backends.cudnn.deterministic = True
return old_values
def reset_seeds(old_values):
import random
import numpy as np
import torch
random.setstate(old_values['random'])
np.random.set_state(old_values['numpy'])
torch.set_rng_state(old_values['torch'])
torch.cuda.set_rng_state(old_values['torch_cuda'])
torch.cuda.set_rng_state_all(old_values['torch_cuda_all'])
torch.backends.cudnn.deterministic = old_values['torch_deterministic']
@ETK_audio_base
class XttsNode:
""" Converts text to speech """
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"device": ("STRING", {"default": "cuda"}),
},
"optional": {
"text": ("STRING", {"multiline": False, "default": ""}),
"[text]": ("LIST", {"multiline": False, "default": ""}),
"speaker_wav": ("STRING", {"multiline": False, "default": ""}),
"language_code": (
['en', 'es', 'fr', 'de', 'it', 'pt', 'pl', 'tr', 'ru', 'nl', 'cs', 'ar', 'zh-cn', 'hu', 'ko', 'ja',
'hi'], {"default": "en"}),
"output_path": ("STRING", {"multiline": False, "default": ""}),
"AUDIO_FOLDER_DEF": ("AUDIO_FOLDER_DEF", {"default": None}),
}
}
RETURN_TYPES = ("STRING", "LIST", "STRING", "AUDIO_FOLDER_DEF",)
RETURN_NAMES = ("output file", "output files list", "output folder", "AUDIO_FOLDER_DEF",)
FUNCTION = 'xtts'
CATEGORY = 'ETK/audio'
def xtts(self=None, **kwargs):
import os
import subprocess
import sys
import shutil
import tempfile
import uuid
from .config import this_file_path
# Handle audio_folder_def
audio_folder_def = kwargs.get("AUDIO_FOLDER_DEF", None)
# Get user-specified output path
base_output_path = kwargs.get('output_path', '.')
# Generate a random temporary folder using tempfile.mkdtemp (persistent until explicitly removed)
temp_folder = tempfile.mkdtemp(dir=base_output_path)
python_exe = sys.executable
# Path to the xttscli.py script
xttscli_path = os.path.join(this_file_path, "modules", "xttscli.py")
# Check if '[text]' is in kwargs and if it is a list
if '[text]' in kwargs and isinstance(kwargs['[text]'], list):
texts = kwargs['[text]']
else:
texts = [kwargs.get('text', "")]
# Create a temporary file in the given folder with the speech separated by new lines
fpn = os.path.join(temp_folder, "temp.txt")
with open(fpn, "w") as f:
for text in texts:
f.write(text + "\n")
# Create the command
command = [
python_exe, xttscli_path,
"--file", fpn,
"--folder", temp_folder,
"--lang", kwargs.get('lang', 'en'),
"--speaker", kwargs.get('speaker_wav', ''),
"--device", kwargs.get('device', 'cuda:0'),
]
# Write the command to "temp_command.txt" in the same folder as this file is in
with open(os.path.join(this_file_path, "temp_command.txt"), "w") as f:
f.write(" ".join(command))
# Execute the command in its own full process and receive the output
try:
# Capture both stdout and stderr
output = subprocess.check_output(command, stderr=subprocess.STDOUT)
generated = output.decode("utf-8")
if "Generated files: (" in generated:
# Extract generated files
generated = str(generated)
generated = generated.split(r"Generated files: (")[1]
generated_files = generated.split(")")[0]
generated_files = eval(generated_files)
# Get the output folder
output_folder = os.path.dirname(kwargs.get('file', "output.wav"))
# Move the files to the correct location
out_files = []
for i, src_file in enumerate(generated_files):
if audio_folder_def:
target_file = audio_folder_def.get_next_file_name()
else:
target_file = os.path.join(temp_folder, f"output_{i}.wav")
shutil.copy(src_file, target_file)
out_files.append(target_file)
else:
raise Exception(f"Error with xttscli, output was: {generated}")
except subprocess.CalledProcessError as e:
# Output the captured stdout + stderr
print(f"Command failed with error: {e.output.decode('utf-8')}")
raise e
# Return values in a consistent manner
return (
out_files[0],
out_files,
temp_folder,
audio_folder_def if audio_folder_def else None
)
import moviepy.editor as mpy
class MovieCrossfadeMultiAudio:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"video_path1": ("STRING", {}),
"video_path2": ("STRING", {}),
"crossfade_duration": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1})
}
}
RETURN_TYPES = ("STRING",)
FUNCTION = "crossfade_videos"
CATEGORY = "video_processing"
def crossfade_videos(self, video_path1, video_path2, crossfade_duration):
# Load video clips
clip1 = mpy.VideoFileClip(video_path1)
clip2 = mpy.VideoFileClip(video_path2)
# Extract primary audio tracks (assuming the first track is primary)
audio1 = clip1.audio
audio2 = clip2.audio
# Crossfade audio tracks
audio_crossfade = mpy.CompositeAudioClip([audio1, audio2.set_start(clip1.duration - crossfade_duration)])
# Crossfade video clips
video_crossfade = mpy.concatenate_videoclips([clip1, clip2.set_start(clip1.duration - crossfade_duration)],
method="compose", padding=-crossfade_duration)
# Set the crossfaded audio to the crossfaded video
final_clip = video_crossfade.set_audio(audio_crossfade)
# Output attributes
self.video = final_clip
return {
"video": self.video
}
@ETK_audio_base
class EtkSaveAudio:
def __init__(self):
self.type = "output"
self.prefix_append = ""
@classmethod
def INPUT_TYPES(s):
return {"required": {"audio": ("AUDIO",),
"filename": ("STRING", {"default": ""})},
}
RETURN_TYPES = ("STRING",)
FUNCTION = "save_audio"
OUTPUT_NODE = True
def save_audio(self, audio, filename=""):
import torchaudio
import io
results = list()
filename += self.prefix_append
for (batch_number, waveform) in enumerate(audio["waveform"].cpu()):
file = f"{filename}_{batch_number:05}.flac"
buff = io.BytesIO()
torchaudio.save(buff, waveform, audio["sample_rate"], format="FLAC")
buff = insert_or_replace_vorbis_comment(buff, {})
with open(file, 'wb') as f:
f.write(buff.getbuffer())
results.append({
"filename": file,
"type": self.type
})
return {"ui": {"audio": results}, "result": (filename,)}
@ETK_audio_base
class PreviewAudio(EtkSaveAudio):
def __init__(self):
import folder_paths
self.output_dir = folder_paths.get_temp_directory()
self.type = "temp"
self.prefix_append = folder_paths.get_temp_directory() + "_temp_" + ''.join(
random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5))
@classmethod
def INPUT_TYPES(s):
return {"required":
{"audio": ("AUDIO",), },
}
@ETK_audio_base
class AdjustAudioVolumeForFiles():
"""Adjusts the volume of a list of audio files."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"audio_files": ("LIST", {"default": []}),
"volume": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1}),
},
"optional": {
"output_folder": ("STRING", {"default": ""}),
}
}
RETURN_TYPES = ("LIST",)
RETURN_NAMES = ("modified_audio_paths",)
CATEGORY = "audio_processing"
OUTPUT_NODE = True
FUNCTION = "handler"
def handler(self, **kwargs):
import os
from os.path import splitext, join
audio_files = kwargs["audio_files"]
volume = kwargs["volume"]
output_folder = kwargs.get("output_folder", "")
modified_audio_paths = []
ffmpeg_exe = ffmpeg.get_ffmpeg_exe()
if ffmpeg_exe is None:
raise RuntimeError("FFmpeg could not be found.")
for audio_file in audio_files:
output_file = splitext(audio_file)[0] + f"_adjusted_volume.mp3"
if output_folder:
output_file = join(output_folder, os.path.basename(output_file))
command = [
ffmpeg_exe, '-i', audio_file,
'-filter:a', f'volume={volume}',
'-y', output_file
]
cmd = ' '.join(command)
print(cmd)
os.system(cmd)
modified_audio_paths.append(output_file)
return (modified_audio_paths,)
if __name__ == '__main__':
# do some testing
wave_file = """C:\\Users\\Tasha\\Nextcloud\\ETK\\YT\\Monsters\\The Haunter of the Ring\\The Haunter of the Ring real.wav"""
xtts_node = XttsNode()
xtts_node.xtts(speaker_wav=wave_file,
text="Hello world! I'm alive, you had better believe it! muahahaha ---- I'm alive! HAHAHAHA! hahahahahahahaha",
language_code="en",
output_path="output.wav")
# use windows to play the file
import os
os.system("start output.wav")
# Paths to the video files to crossfade
video_path1 = 'path/to/your/first/video.mp4'
video_path2 = 'path/to/your/second/video.mp4'
# Duration of the crossfade in seconds
crossfade_duration = 2.0 # Example duration
# Create an instance of the MovieCrossfadeMultiAudio node
crossfade_node = MovieCrossfadeMultiAudio()
# Execute the node's crossfade function
output_video = crossfade_node.crossfade_videos(video_path1, video_path2, crossfade_duration)
# Save the output video to a file
output_video_path = 'path/to/your/output/crossfaded_video.mp4'
output_video['video'].write_videofile(output_video_path, codec='libx264', audio_codec='aac')
print(f'Crossfaded video saved to {output_video_path}')
audio_video_folder.py
-1081
View File
File diff suppressed because it is too large Load Diff
basic.py
-748
View File
@@ -1,748 +0,0 @@
class MultilineStringNode:
"""Noe that displays multiline text box"""
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"Text": ("STRING", {
"default": "",
"multiline": True,
}),
}
}
RETURN_TYPES = ("STRING", "INT", "FLOAT", "LIST")
FUNCTION = "get_value"
CATEGORY = "text"
def get_value(self, Text):
return (Text,)
class TupleNode:
"""Node that displays a tuple"""
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"Tuple": ("STRING", {
"default": "(None,None,)",
}),
}
}
RETURN_TYPES = ("TUPLE",)
FUNCTION = "get_value"
CATEGORY = "text"
def get_value(self, Tuple):
# evaluate the string safely step by step
# type check
if not isinstance(Tuple, str):
raise TypeError("Tuple must be a string")
# format check
if not Tuple.startswith("(") or not Tuple.endswith(")"):
raise ValueError("Tuple must start with '(' and end with ')'")
# remove the brackets
Tuple = Tuple.strip("()")
# check for empty tuple
if len(Tuple) == 0:
return ((),)
# check for commas
if "," not in Tuple:
raise ValueError("Tuple must have at least one comma")
# split on commas
Tuple = Tuple.split(",")
# valid values are string,int,float,None, list,tuple,dict
# construct the tuple
ret = []
for val in Tuple:
val = val.strip()
if val == "None":
ret.append(None)
elif val.startswith("[") and val.endswith("]"):
# list
val = val.strip("[]")
val = val.split(",")
ret.append(val)
elif val.startswith("(") and val.endswith(")"):
# tuple
val = val.strip("()")
val = val.split(",")
ret.append(tuple(val))
elif val.startswith("{") and val.endswith("}"):
# dict
val = val.strip("{}")
val = val.split(",")
ret.append(dict(val))
elif val.startswith('"') and val.endswith('"'):
# string
ret.append(val.strip('"'))
elif "." in val:
# float
ret.append(float(val))
elif val == "":
# empty
pass
else:
# int
ret.append(int(val))
return (tuple(ret),)
class StrBreakout:
""" node that breaks out a string into a:
- Boolean
- int
- float
- [string]
- splitlines()
- split(",")
the first value is a description of the transformations tried
"""
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"Text": ("STRING", {
"default": "",
}),
"list_dtype": (["str", "int", "float", ],),
}
}
RETURN_TYPES = ("STRING", "BOOL", "INT", "FLOAT", "LIST", "LIST", "LIST",)
RETURN_NAMES = ("Notes", "Bool", "Int", "Float", "List by chr", "list by lines", "list by commas",)
FUNCTION = "get_value"
CATEGORY = "text"
def get_value(self, Text, list_dtype="str"):
out_bool = False
out_int = 0
out_float = 0.0
out_list = []
out_lines = []
out_comma = []
notes = []
try:
out_bool = bool(Text)
notes.append(f"bool worked")
except Exception as e:
notes.append(f"bool failed: {e}")
try:
out_int = int(Text)
notes.append(f"int worked")
except Exception as e:
notes.append(f"int failed: {e}")
pass
try:
out_float = float(Text)
notes.append(f"float worked")
except Exception as e:
pass
try:
out_list = list(Text)
notes.append(f"list worked")
except Exception as e:
notes.append(f"list failed: {e}")
pass
try:
out_lines = Text.splitlines()
notes.append(f"splitlines worked")
except Exception as e:
notes.append(f"splitlines failed: {e}")
pass
try:
out_comma = Text.split(",")
notes.append(f"split(',') worked")
except Exception as e:
notes.append(f"split(',') failed: {e}")
pass
if list_dtype == "int":
list_dtype = int
elif list_dtype == "float":
list_dtype = float
else:
list_dtype = str
try:
out_list = [list_dtype(val) for val in out_list]
notes.append(f"list dtype worked")
except Exception as e:
notes.append(f"list dtype failed: {e}")
pass
try:
out_lines = [list_dtype(val) for val in out_lines]
notes.append(f"splitlines dtype worked")
except Exception as e:
notes.append(f"splitlines dtype failed: {e}")
pass
try:
out_comma = [list_dtype(val) for val in out_comma]
notes.append(f"split(',') dtype worked")
except Exception as e:
notes.append(f"split(',') dtype failed: {e}")
pass
return ("\n".join(notes), out_bool, out_int, out_float, out_list, out_lines, out_comma,)
class IntBreakout:
""" node that breaks out an int into a:
- Boolean
- str(int)
- float
- [int]
- list(range(int))
- [None] * int
the first value is a description of the transformations tried
"""
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"Int": ("INT", {
"default": 0,
}),
}
}
RETURN_TYPES = ("STRING", "BOOL", "STRING", "FLOAT", "LIST", "LIST", "LIST",)
FUNCTION = "get_value"
CATEGORY = "text"
def get_value(self, Int):
out_bool = False
out_str = ""
out_float = 0.0
out_list = []
out_range = []
out_none = []
notes = []
try:
out_bool = bool(Int)
notes.append(f"bool worked")
except Exception as e:
notes.append(f"bool failed: {e}")
try:
out_str = str(Int)
notes.append(f"str worked")
except Exception as e:
notes.append(f"str failed: {e}")
pass
try:
out_float = float(Int)
notes.append(f"float worked")
except Exception as e:
pass
try:
out_list = list(Int)
notes.append(f"list worked")
except Exception as e:
notes.append(f"list failed: {e}")
pass
try:
out_range = list(range(Int))
notes.append(f"range worked")
except Exception as e:
notes.append(f"range failed: {e}")
pass
try:
out_none = [None] * Int
notes.append(f"[None] * worked")
except Exception as e:
notes.append(f"[None] * failed: {e}")
pass
return ("\n".join(notes), out_bool, out_str, out_float, out_list, out_range, out_none,)
class SelectStrFromDict:
"""
Selects a string from a dictionary
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"dict": ("DICT", {
"default": {},
}),
"key": ("STRING", {
"default": "",
}),
}
}
RETURN_TYPES = ("STRING",)
FUNCTION = "get_value"
CATEGORY = "text"
def get_value(self, dict, key):
if key in dict:
return (dict[key],)
return ("",)
class StringToAny:
"""takes STRING, and returns the value coded as *"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"string": ("STRING", {"forceInput": True}),
}
}
RETURN_TYPES = ("*", "STRING",)
RETURN_NAMES = ("*", "debug info",)
FUNCTION = "get_value"
CATEGORY = "conversion"
def get_value(self, string):
try:
ret_str = str(string)
dbg = f"found STRING: {string}"
except Exception as e:
ret_str = None
dbg = f"failed to verify STRING: {string} error: {e}"
return (ret_str, dbg,)
class ValueToAny:
"""takes INT,STRING,FLOAT,BOOL, and returns the value coded as *"""
@classmethod
def INPUT_TYPES(cls):
ret = {
"required": {},
"optional": {
"INT": ("INT", {"forceInput": True}),
"FLOAT": ("FLOAT", {"forceInput": True}),
"STRING": ("STRING", {"forceInput": True}),
"BOOL": ("BOOL", {"forceInput": True}),
}}
return ret
RETURN_TYPES = ("*", "STRING",)
RETURN_NAMES = ("*", "debug info",)
FUNCTION = "get_value"
CATEGORY = "conversion"
def get_value(self, INT=None, FLOAT=None, STRING=None, BOOL=None):
if INT:
dbg = f"found INT: {INT}"
return (INT, dbg,)
if FLOAT:
dbg = f"found FLOAT: {FLOAT}"
return (FLOAT, dbg,)
if STRING:
dbg = f"found STRING: {STRING}"
return (STRING, dbg,)
if BOOL:
dbg = f"found BOOL: {BOOL}"
return (BOOL, dbg,)
dbg = f"found nothing, will return the first value that is not None; from INT,FLOAT,STRING,BOOL"
return (None, dbg)
class AnyToAny:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"A": ("*", {"forceInput": True}),
}
}
RETURN_TYPES = ("*",)
FUNCTION = "get_value"
CATEGORY = "conversion"
RETURN_NAMES = ("*",)
def get_value(self, A):
return (A,)
class IterableToAny:
"""takes LIST,DICT, and returns the value coded as *"""
@classmethod
def INPUT_TYPES(cls):
return {
"optional": {
"LIST": ("LIST", {"default": None, }),
"DICT": ("DICT", {"default": None, }),
}
}
RETURN_TYPES = ("*", "STRING",)
RETURN_NAMES = ("*", "debug info",)
FUNCTION = "get_value"
CATEGORY = "conversion"
def get_value(self, LIST=None, DICT=None):
if LIST:
if not isinstance(LIST, list):
return (LIST, f"LIST is not a list: {LIST} is type: {type(LIST)}")
dbg = f"found LIST: {LIST}"
return (LIST, dbg,)
if DICT:
if not isinstance(DICT, dict):
return (LIST, f"LIST is not a list: {LIST} is type: {type(LIST)}")
dbg = f"found DICT: {DICT}"
return (DICT, dbg,)
dbg = f"found nothing, will return the first value that is not None; from LIST,DICT"
return (None, dbg)
class CallableToAny:
"""takes FUNC, (others soon) and returns the value coded as *"""
@classmethod
def INPUT_TYPES(cls):
return {
"optional": {
"FUNC": ("FUNC", {"default": None, }),
}
}
RETURN_TYPES = ("*", "STRING",)
RETURN_NAMES = ("*", "debug info",)
FUNCTION = "get_value"
CATEGORY = "conversion"
def get_value(self, FUNC=None):
if FUNC:
dbg = f"found FUNC: {FUNC}"
return (FUNC, dbg,)
dbg = f"found nothing, will return the first value that is not None; from FUNC"
return (None, dbg)
class ObjectToAny:
""" takes non-primitive objects like TENSORS and returns the value coded as *"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"label": ("STRING", {"default": "0", })},
"optional": {
"LATENT": ("LATENT", {"default": None, }),
"CONDITIONING": ("CONDITIONING", {"default": None, }),
"IMAGE": ("IMAGE", {"default": None, }),
"MASK": ("MASK", {"default": None, }),
"CLIP": ("CLIP", {"default": None, }),
"MODEL": ("MODEL", {"default": None, }),
}
}
RETURN_TYPES = ("*", "STRING",)
RETURN_NAMES = ("*", "debug info",)
FUNCTION = "get_value"
CATEGORY = "conversion"
def get_value(self, label,
LATENT=None,
CONDITIONING=None,
IMAGE=None,
MASK=None,
CLIP=None,
MODEL=None,
):
if LATENT:
dbg = f"found LATENT: {LATENT}"
return (LATENT, dbg,)
if CONDITIONING:
dbg = f"found CONDITIONING: {CONDITIONING}"
return (CONDITIONING, dbg,)
if IMAGE is not None:
dbg = f"found IMAGE: {IMAGE}"
return (IMAGE, dbg,)
if MASK:
dbg = f"found MASK: {MASK}"
return (MASK, dbg,)
if MODEL:
dbg = f"found MODEL: {MODEL}"
return (MODEL, dbg,)
if CLIP:
dbg = f"found CLIP: {CLIP}"
return (CLIP, dbg,)
dbg = f"found nothing, will return the first value that is not None; from LATENT,CONDITIONING,IMAGE,MASK"
return (None, dbg)
from torch import Tensor
from comfy.model_patcher import ModelPatcher
COMFY_PYTHON_CLAS_NAMES = {
"INT": int,
"FLOAT": float,
"STRING": str,
"BOOLEAN": bool,
"TUPLE": tuple,
"LIST": list,
"DICT": dict,
"FUNC": callable,
"LATENT": dict,
"CONDITIONING": dict,
"IMAGE": Tensor,
"MASK": Tensor,
"LLLM_MESSAGES": list,
"LITELLM_MODEL": str,
"MODEL": ModelPatcher,
}
def value_try(val, cls):
try:
return cls(val)
except:
return None
class AnyToValue:
"""
takes * and returns it in the node systems common value types
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"A": ("*", {
"default": None,
}),
}}
RETURN_TYPES = ("INT", "FLOAT", "STRING", "BOOLEAN", "TUPLE",)
FUNCTION = "get_value"
CATEGORY = "conversion"
def get_value(self, A):
int_o = value_try(A, int)
float_o = value_try(A, float)
str_o = value_try(A, str)
bool_o = value_try(A, bool)
tuple_o = value_try(A, tuple)
return (int_o, float_o, str_o, bool_o, tuple_o,)
class AnyToIterable:
"""
takes * and returns it in the node systems common value types
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"Any": ("*", {
"default": None,
}),
}}
RETURN_TYPES = ("LIST", "DICT",)
FUNCTION = "get_value"
CATEGORY = "conversion"
def get_value(self, Any):
list_o = None
dict_o = None
if isinstance(Any, list):
list_o = Any
if isinstance(Any, dict):
dict_o = Any
return (list_o, dict_o,)
class AnyToObject:
"""
like AnyToIterable but for Image,latent,conditioning,model
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"x": ("*", {
"default": None,
}),
}}
RETURN_TYPES = ("IMAGE", "LATENT", "CONDITIONING", "MODEL")
FUNCTION = "get_value"
CATEGORY = "conversion"
def get_value(self, x):
image_o = x
latent_o = x
conditioning_o = x
model_o = x
return (image_o, latent_o, conditioning_o, model_o,)
class TimeThis:
"""
Pass through node that records the current time to a given variable name
then when called again reports the time difference
"""
start_times = {}
end_times = {}
last_activator = None
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"var_name": ("STRING", {"default": "timer0", }),
"pass_through": ("*", {"default": None, }),
"start": ("BOOLEAN", {"default": True})
},
"optional": {
"activator": ("*")
}
}
INTERNAL_STATE_DISPLAY = "text_display"
RETURN_TYPES = ("STRING", "FLOAT", "OBJECT")
RETURN_NAMES = ("Notes", "Time", "Object")
FUNCTION = "time_this"
CATEGORY = "utils"
def time_this(self, var_name, pass_through, start, activator=None):
import time
if start:
self.start_times[var_name] = time.time()
print(f"Timer started for {var_name}")
ret = ("Timer started", 0.0, pass_through)
else:
if var_name not in self.start_times:
return ("Timer not started", 0.0, pass_through)
self.end_times[var_name] = time.time()
time_diff = self.end_times[var_name] - self.start_times[var_name]
print(f"Timer ended for {var_name} time taken: {time_diff}")
ret = (f"Time taken: {time_diff}", time_diff, pass_through)
ret = {"ui": {"text": (ret[0],)}, "result": ret}
return ret
class Text_ETK:
"""
to provide text to the model
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"text": ("STRING", {"multiline": True}),
}
}
RETURN_TYPES = ("STRING",)
FUNCTION = "fun"
CATEGORY = "text"
def fun(self, text):
return (text,)
class StringsToList:
"""
take two strings and return them as a list
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"list": ("LIST", {"default": [], "forceInput": True}),
"string1": ("STRING", {"default": "", "forceInput": True}),
},
"optional": {
"string2": ("STRING", {"default": "", "forceInput": True}),
}
}
RETURN_TYPES = ("LIST",)
FUNCTION = "fun"
CATEGORY = "text"
def fun(self, list, string1, string2=""):
return (list + [string1, string2],)
# from __init__ import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
from . import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
NODE_CLASS_MAPPINGS["Text_ETK"] = Text_ETK
NODE_CLASS_MAPPINGS["Text _O"] = Text_ETK
NODE_CLASS_MAPPINGS["MultilineStringNode"] = MultilineStringNode
NODE_CLASS_MAPPINGS["StrBreakout"] = StrBreakout
NODE_CLASS_MAPPINGS["IntBreakout"] = IntBreakout
NODE_CLASS_MAPPINGS["ValueToAny"] = ValueToAny
NODE_CLASS_MAPPINGS["IterableToAny"] = IterableToAny
NODE_CLASS_MAPPINGS["CallableToAny"] = CallableToAny
NODE_CLASS_MAPPINGS["ObjectToAny"] = ObjectToAny
NODE_CLASS_MAPPINGS["ETK_Tuple"] = TupleNode
NODE_CLASS_MAPPINGS["AnyToValue"] = AnyToValue
NODE_CLASS_MAPPINGS["AnyToIterable"] = AnyToIterable
NODE_CLASS_MAPPINGS["TimeThis"] = TimeThis
NODE_CLASS_MAPPINGS["AnyToObject"] = AnyToObject
NODE_CLASS_MAPPINGS["SelectStrFromDict"] = SelectStrFromDict
NODE_CLASS_MAPPINGS["StringToAny"] = StringToAny
NODE_CLASS_MAPPINGS["StringsToList"] = StringsToList
NODE_CLASS_MAPPINGS["AnyToAny"] = AnyToAny
descript.py
-87
View File
@@ -1,87 +0,0 @@
import copy
NODE_CLASS_MAPPINGS = {}
NODE_DISPLAY_NAME_MAPPINGS = {}
def ETK_descript_base(cls):
cls.FUNCTION = "func"
cls.CATEGORY = "ETK/Descript"
NODE_CLASS_MAPPINGS[cls.__name__] = cls
return cls
@ETK_descript_base
class DecodeTranscriptNode:
"""
Decode one of descript's transcript formatted files
String input is the files contents
output is 3 lists of strings (speaker,duration,paragraph,)
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"string": ("STRING", {"multiline": True,
"default": ""
}),
}
}
RETURN_TYPES = ("LIST", "LIST", "LIST")
RETURN_NAMES = ("speakers", "durations", "paragraphs")
OUTPUT_NODE = True
def func(self, string):
"""
input is like:
[00:00:00] jackie: Let's talk about a
[00:00:01] jackie: thing
:param string:
:return:
"""
paras = string.split("\n\n")
speakers = []
durations = []
paragraphs = []
for para in paras:
if para.strip() == "":
continue
# first find the first ] and split on that
part1 = para.find("]")
time_part = para[:part1 + 1]
remainder = para[part1 + 1:]
# now find first : in the remainder
part2 = remainder.find(":")
speaker_part = remainder[:part2 + 1].replace(":", "").strip()
paragraph_part = remainder[part2 + 1:]
speakers.append(speaker_part.strip())
# now just remove [ and ] from time_part
time_part = time_part.replace("[", "").replace("]", "").strip()
# split it on :
time_part = time_part.split(":")
# now convert to seconds
time_part = int(time_part[0]) * 60 * 60 + int(time_part[1]) * 60 + int(time_part[2])
# save here for post-processing later
durations.append(time_part)
paragraphs.append(paragraph_part.strip())
import copy
start_seconds = copy.copy(durations)
durations = []
for i in range(len(start_seconds) - 1):
durations.append(start_seconds[i + 1] - start_seconds[i])
# now append the last value
durations.append(int(sum(durations) / len(durations)))
return (speakers, durations, paragraphs,)
functional.py
-774
View File
@@ -1,774 +0,0 @@
import hashlib
from . import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
# NODE_CLASS_MAPPINGS = {}
# NODE_DISPLAY_NAME_MAPPINGS = {}
_list = {"list": ("LIST",)}
_func = {"func": ("FUNC", {"default": None})}
_func_extra = {"func_extra": ("FUNC",)}
_text = {"code": ("STRING", {"multiline": True, "default": "y mapped to output"})}
_image = {"image": ("IMAGE",)}
_string = {"string": ("STRING", {"default": ""})}
_any = {"any": ("*", {"default": ""})}
optional = lambda x: {"optional": {k: v for k, v in x.items()}}
required = lambda x: {"required": {k: v for k, v in x.items()}}
both = lambda a, b: {**a, **b}
many = lambda *dicts: {k: v for d in dicts for k, v in d.items()} # non tested
def ETK_functional_base(cls):
cls.FUNCTION = "func"
cls.CATEGORY = "ETK/functional"
NODE_CLASS_MAPPINGS[cls.__name__] = cls
return cls
def new_exec(c, g, l):
global old_exec
import os
# if GDE_SECURE_MODE is set to NO_EXEC then we will not execute any code
# this is a security feature to prevent malicious code from being executed
# this is set in the GDE settings
GDE_SECURE_MODE = os.environ.get("GDE_SECURE_MODE", None)
if GDE_SECURE_MODE == "NO_EXEC":
raise Exception("GDE_SECURE_MODE is set to NO_EXEC, not executing code")
old_exec(c, g, l)
return locals
old_exec = exec
exec = new_exec
@ETK_functional_base
class FuncBase:
"""
Base class for func nodes
Allows for a function to be passed in as a callable
and
for a string to be passed in as code to be executed
also
returns the function as a callable
"""
def __init__(self):
from custom_nodes.EternalKernelLiteGraphNodes.image import common_ksampler, TinyTxtToImg
self.locals = locals()
self.globals = globals()
@classmethod
def INPUT_TYPES(cls):
req = required(_text)
opt = optional(
both(
both(
_func,
_func_extra
),
_list
)
)
use = both(req, opt)
return use
RETURN_TYPES = ("FUNC",)
FUNCTION = "func"
CATEGORY = "ETK/func"
def func(self, **kwargs):
import inspect
self.globals = globals()
self.locals = locals()
self.locals["x1"] = kwargs.get("func", None)
self.locals["x2"] = kwargs.get("func_extra", None)
self.locals["x_list"] = kwargs.get("list", None)
self.IN_FUNC = kwargs.get("func", None)
self.EXTRA_IN_FUNC = kwargs.get("func_extra", None)
self.exec_code = kwargs.get("code", None)
self.ARGS = kwargs
self.CODE = inspect.getsource(self.func)
exec(self.exec_code, self.globals, self.locals)
return (self.func,)
@ETK_functional_base
class FuncRender:
"""
as input takes a function and a string of code
optionally takes a list of arguments
returns rendered results of types
float, string, int, list
"""
def __init__(self):
from custom_nodes.EternalKernelLiteGraphNodes.image import common_ksampler, TinyTxtToImg
self.locals = locals()
self.globals = globals()
@classmethod
def INPUT_TYPES(cls):
use = both(
required(
_text),
optional(both(_list, _func)
)
)
return use
RETURN_TYPES = ("FLOAT", "STRING", "INT", "LIST")
FUNCTION = "func"
CATEGORY = "ETK/func"
def func(self, **kwargs):
y_float = None
y_string = None
y_int = None
func = kwargs.get("func", None)
code = kwargs.get("code", None)
globals = self.globals
locals = self.locals
locals["y_float"] = y_float
locals["y_string"] = y_string
locals["y_int"] = y_int
locals["y_list"] = None
locals["x"] = func
locals["x_list"] = kwargs.get("list", None)
# code should do something like:
# func like : lambda x:["what"]*x
# code like : 'f=[x for x in x(2)]'
exec(code, globals, locals)
self.globals = globals
self.locals = locals
return (locals["y_float"], locals["y_string"], locals["y_int"], locals["y_list"],)
@ETK_functional_base
class FuncRenderImage:
def __init__(self):
from custom_nodes.EternalKernelLiteGraphNodes.image import common_ksampler, TinyTxtToImg
self.locals = locals()
self.globals = globals()
@classmethod
def INPUT_TYPES(cls):
use = both(
required(
both(_func, _text)),
optional(_list)
)
return use
RETURN_TYPES = ("IMAGE",)
FUNCTION = "func"
CATEGORY = "ETK/func"
OUTPUT_NODE = True
def func(self, func, code, **kwargs):
globals = self.globals
locals = self.locals
locals["x"] = func
locals["x_list"] = kwargs.get("list", None)
# code should do something like:
# func like : lambda x:["what"]*x
# code like : 'f=[x for x in x(2)]'
exec(code, globals, locals)
self.globals = globals
self.locals = locals
return (locals["y"],)
@ETK_functional_base
class ExecWidget:
import torch
"""runs eval on the given text"""
@classmethod
def INPUT_TYPES(cls):
return {"optional":
{
"code_str_in": ("STRING", {"multiline": True, "default": ""}),
"image1_in": ("IMAGE",),
"float1_in": ("FLOAT", {"multiline": False}),
"string1_in": ("STRING", {"multiline": False}),
"int1_in": ("INT", {"multiline": False}),
},
"required": {
"text_to_eval": ("STRING",
{"multiline": True,
"default": "int_out=int_out\n"
"float_out=float_out\n"
"string_out=string_out\n"
"image_out=image_out\n"
}),
}
}
CATEGORY = "text"
RETURN_TYPES = ("STRING", "IMAGE", "FLOAT", "INT", "TUPLE", "CODE")
FUNCTION = "pre_exec_handler"
INTERNAL_STATE_DISPLAY_CODE = True
OUTPUT_NODE = True
def pre_exec_handler(self, **kwargs):
image1_in = kwargs.get("image1_in", None)
float1_in = kwargs.get("float1_in", 0.0)
string1_in = kwargs.get("string1_in", "")
int1_in = kwargs.get("int1_in", 0)
tuple1_in = kwargs.get("tuple1_in", None)
name = kwargs.get("name", "exec_func")
text_to_eval = kwargs.get("text_to_eval", "")
code_str_in = kwargs.get("code_str_in", "")
if code_str_in != "":
text_to_eval = code_str_in
else:
code_str_in = None
ret = self.exec_handler(text_to_eval, image1_in, float1_in, string1_in, int1_in, tuple1_in, name)
ret_formatted = {"ui": {"text": [ret[-1]]}, "result": ret}
return ret_formatted
def exec_handler(self, text_to_eval, image_obj: torch.Tensor = None, float_obj: float = 0.0, string_obj: str = "",
int_obj: int = 0, tuple_obj: tuple = None, name: str = "exec_func"):
"""
>>> ExecWidget().exec_handler("2 + 3")
'5'
"""
new_locals = {"float_out": float_obj,
"image_out": image_obj,
"string_out": string_obj,
"int_out": int_obj,
"tuple_out": tuple_obj,
}
try:
if isinstance(text_to_eval, list):
text_to_eval = text_to_eval[0]
# evaluate for code tag
if "<code>" in text_to_eval and "</code>" in text_to_eval:
start = text_to_eval.find("<code>")
end = text_to_eval.find("</code>")
text_to_eval = text_to_eval[start + 6:end]
# also evaluate for ```python tag
if "```python" in text_to_eval and "```" in text_to_eval:
start = text_to_eval.find("```python")
end = text_to_eval.find("```", start + 8)
text_to_eval = text_to_eval[start + 9:end]
if "```" in text_to_eval:
start = text_to_eval.find("```")
end = text_to_eval.find("```", start + 2)
text_to_eval = text_to_eval[start + 3:end]
exec(text_to_eval, globals(), new_locals)
source_code = text_to_eval
except Exception as e:
string_obj = str(e)
source_code = text_to_eval
string_obj = new_locals.get("string_out", string_obj)
image_obj = new_locals.get("image_out", image_obj)
float_obj = new_locals.get("float_out", float_obj)
int_obj = new_locals.get("int_out", int_obj)
tuple_obj = new_locals.get("tuple_out", tuple_obj)
return (string_obj, image_obj, float_obj, int_obj, tuple_obj, text_to_eval)
@classmethod
def IS_CHANGED_NOPE(s, text_to_eval, image1_in: torch.Tensor = None, float1_in: float = 0.0, string1_in: str = "",
int1_in: int = 0, name: str = "exec_func"):
m = hashlib.sha256()
m.update(name.encode("utf-8"))
return m.digest().hex()
@ETK_functional_base
class FuncListToList:
"""runs eval on the given text, with a list as input and output"""
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
req = required(
both(_list,
_text)
)
opt = optional(_func)
use = both(req, opt)
return use
CATEGORY = "ETK/func"
RETURN_TYPES = ("LIST",)
FUNCTION = "func"
def func(self, **kwargs):
func = kwargs.get("func", None)
code = kwargs.get("code", None)
my_globals = globals()
my_locals = locals()
my_locals["x"] = func
my_locals["x_list"] = kwargs.get("list", None)
exec(code, my_globals, my_locals)
return (my_locals["y"],)
@ETK_functional_base
class FuncStrToStr:
"""runs eval on the given text, with a string as input and outputs a string"""
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
req = required(both(_string, _text))
return req
CATEGORY = "ETK/func"
RETURN_TYPES = ("STRING",)
FUNCTION = "func"
def func(self, **kwargs):
text = kwargs.get("string", None)
x = text
code = kwargs.get("code", None)
my_globals = globals()
my_locals = locals()
exec(code, my_globals, my_locals)
return (my_locals["y"],)
@ETK_functional_base
class FuncStrToList:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
out = {"required": {
"string": ("STRING", {"default": "", }),
"code": ("STRING", {"multiline": True, "default": "y mapped to output"})
}}
return out
CATEGORY = "ETK/func"
RETURN_TYPES = ("LIST",)
FUNCTION = "func"
def func(self, **kwargs):
text = kwargs.get("string", None)
code = kwargs.get("code", None)
x = text
my_globals = globals()
my_locals = locals()
exec(code, my_globals, my_locals)
return (my_locals["y"],)
@ETK_functional_base
class FuncListToStr:
# class FuncStrToList:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
req = required(both(_list, _text))
return req
CATEGORY = "ETK/func"
RETURN_TYPES = ("STRING",)
FUNCTION = "func"
def func(self, **kwargs):
lst = kwargs.get("list", None)
code = kwargs.get("code", None)
x = lst
x_list = lst
my_globals = globals()
my_locals = locals()
exec(code, my_globals, my_locals)
return (my_locals["y"],)
@ETK_functional_base
class GetFirstCodeBlock():
""" returns the first code block in the text which is surrounded with ``` or
```python or
<code> tags. If no code block is found, returns an empty string"""
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
req = required(_text)
return req
CATEGORY = "ETK/func"
RETURN_TYPES = ("STRING",)
FUNCTION = "func"
def func(self, **kwargs):
text = kwargs.get("code", None)
# find the code block
if "```python" in text and "```" in text:
start = text.find("```python")
end = text.find("```", start + 8)
text = text[start + 9:end]
elif "```" in text:
start = text.find("```")
end = text.find("```", start + 2)
text = text[start + 3:end]
elif "<code>" in text and "</code>" in text:
start = text.find("<code>")
end = text.find("</code>")
text = text[start + 6:end]
else:
text = ""
return (text,)
@ETK_functional_base
class FuncAnysToImage:
"""runs eval on the given text, with 5 anys as input and outputs an image"""
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
req = optional(
both(
both({"any1": ("*", {"default": None})}, {"any2": ("*", {"default": None})}),
both({"any3": ("*", {"default": None})}, {"any4": ("*", {"default": None})})
),
)
req["required"] = {}
req["required"].update(_text)
return req
CATEGORY = "ETK/func"
RETURN_TYPES = ("IMAGE",)
FUNCTION = "func"
def func(self, **kwargs):
import gc
import torch
code = kwargs.get("code", None)
any1 = kwargs.get("any1", None)
any2 = kwargs.get("any2", None)
any3 = kwargs.get("any3", None)
any4 = kwargs.get("any4", None)
my_globals = globals().copy()
my_locals = {}
my_locals["any1"] = any1
my_locals["any2"] = any2
my_locals["any3"] = any3
my_locals["any4"] = any4
exec(code, my_globals, my_locals)
if "y" in my_locals.keys():
y = my_locals["y"]
elif "y" in my_globals.keys():
y = my_globals["y"]
else:
y = None
del my_locals
del my_globals
del code
del any1, any2, any3, any4
gc.collect()
torch.cuda.empty_cache()
return (y,)
@ETK_functional_base
class FuncAnysToLatent:
"""runs eval on the given text, with 5 anys as input and outputs a latent"""
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
req = optional(
both(
both({"any1": ("*", {"default": None})}, {"any2": ("*", {"default": None})}),
both({"any3": ("*", {"default": None})}, {"any4": ("*", {"default": None})})
),
)
req["required"] = {}
req["required"].update(_text)
return req
CATEGORY = "ETK/func"
RETURN_TYPES = ("LATENT",)
FUNCTION = "func"
def func(self, **kwargs):
code = kwargs.get("code", None)
any1 = kwargs.get("any1", None)
any2 = kwargs.get("any2", None)
any3 = kwargs.get("any3", None)
any4 = kwargs.get("any4", None)
my_globals = globals()
my_locals = locals()
my_locals["any1"] = any1
my_locals["any2"] = any2
my_locals["any3"] = any3
my_locals["any4"] = any4
exec(code, my_globals, my_locals)
if "y" in my_locals.keys():
return (my_locals["y"],)
elif "y" in my_globals.keys():
return (my_globals["y"],)
else:
return (None,)
@ETK_functional_base
class FuncAnysToString:
"""runs eval on the given text, with 5 anys as input and outputs a string"""
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
req = optional(
both(
both({"any1": ("*", {"default": None})}, {"any2": ("*", {"default": None})}),
both({"any3": ("*", {"default": None})}, {"any4": ("*", {"default": None})})
),
)
req["required"] = {}
req["required"].update(_text)
return req
CATEGORY = "ETK/func"
RETURN_TYPES = ("STRING",)
FUNCTION = "func"
def func(self, **kwargs):
code = kwargs.get("code", None)
any1 = kwargs.get("any1", None)
any2 = kwargs.get("any2", None)
any3 = kwargs.get("any3", None)
any4 = kwargs.get("any4", None)
my_globals = globals()
my_locals = locals()
my_locals["any1"] = any1
my_locals["any2"] = any2
my_locals["any3"] = any3
my_locals["any4"] = any4
exec(code, my_globals, my_locals)
if "y" in my_locals.keys():
return (my_locals["y"],)
elif "y" in my_globals.keys():
return (my_globals["y"],)
else:
return (None,)
@ETK_functional_base
class FuncAnysToList:
"""runs eval on the given text, with 5 anys as input and outputs a list"""
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
req = optional(
both(
both({"any1": ("*", {"default": None})}, {"any2": ("*", {"default": None})}),
both({"any3": ("*", {"default": None})}, {"any4": ("*", {"default": None})})
),
)
req["required"] = {}
req["required"].update(_text)
return req
CATEGORY = "ETK/func"
RETURN_TYPES = ("LIST",)
FUNCTION = "func"
def func(self, **kwargs):
code = kwargs.get("code", None)
any1 = kwargs.get("any1", None)
any2 = kwargs.get("any2", None)
any3 = kwargs.get("any3", None)
any4 = kwargs.get("any4", None)
my_globals = globals()
my_locals = {}
my_locals["any1"] = any1
my_locals["any2"] = any2
my_locals["any3"] = any3
my_locals["any4"] = any4
exec(code, my_globals, my_locals)
if "y" in my_locals.keys():
return (my_locals["y"],)
elif "y" in my_globals.keys():
return (my_globals["y"],)
else:
return (None,)
@ETK_functional_base
class FuncAnysToDict:
"""runs eval on the given text, with 5 anys as input and outputs a list"""
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
req = optional(
both(
both({"any1": ("*", {"default": None})}, {"any2": ("*", {"default": None})}),
both({"any3": ("*", {"default": None})}, {"any4": ("*", {"default": None})})
),
)
req["required"] = {}
req["required"].update(_text)
return req
CATEGORY = "ETK/func"
RETURN_TYPES = ("DICT",)
FUNCTION = "func"
def func(self, **kwargs):
code = kwargs.get("code", None)
any1 = kwargs.get("any1", None)
any2 = kwargs.get("any2", None)
any3 = kwargs.get("any3", None)
any4 = kwargs.get("any4", None)
my_globals = globals()
my_locals = {}
my_locals["any1"] = any1
my_locals["any2"] = any2
my_locals["any3"] = any3
my_locals["any4"] = any4
exec(code, my_globals, my_locals)
if "y" in my_locals.keys():
return (my_locals["y"],)
elif "y" in my_globals.keys():
return (my_globals["y"],)
else:
return (None,)
@ETK_functional_base
class FuncAnysToCond:
"""runs eval on the given text, with 5 anys as input and outputs a cond"""
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
req = optional(
both(
both({"any1": ("*", {"default": None})}, {"any2": ("*", {"default": None})}),
both({"any3": ("*", {"default": None})}, {"any4": ("*", {"default": None})})
),
)
req["required"] = {}
req["required"].update(_text)
return req
CATEGORY = "ETK/func"
RETURN_TYPES = ("CONDITIONING",)
FUNCTION = "func"
def func(self, **kwargs):
import gc
code = kwargs.get("code", None)
any1 = kwargs.get("any1", None)
any2 = kwargs.get("any2", None)
any3 = kwargs.get("any3", None)
any4 = kwargs.get("any4", None)
my_globals = globals()
my_locals = {}
my_locals["any1"] = any1
my_locals["any2"] = any2
my_locals["any3"] = any3
my_locals["any4"] = any4
exec(code, my_globals, my_locals)
if "y" in my_locals.keys():
y = (my_locals["y"],)
elif "y" in my_globals.keys():
y = my_globals["y"],
else:
y = (None,)
del my_locals
gc.collect()
return (y,)
if __name__ == "__main__":
test = {}
pass
git.py
-185
View File
@@ -1,185 +0,0 @@
from git import Repo
import copy
import re
from . import NODE_CLASS_MAPPINGS,NODE_DISPLAY_NAME_MAPPINGS
#NODE_CLASS_MAPPINGS = {}
#NODE_DISPLAY_NAME_MAPPINGS = {}
def ETK_git_base(cls):
cls.FUNCTION = "func"
cls.CATEGORY = "ETK/git"
NODE_CLASS_MAPPINGS[cls.__name__] = cls
return cls
@ETK_git_base
class PublicGitRepo:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"repo_url": ("STRING", {"multiline": False, "default": ""}),
"token": ("STRING", {"multiline": False, "default": ""}),
}
}
RETURN_TYPES = ("bool",)
FUNCTION = 'validate_git_repo_url'
CATEGORY = 'ETK/git'
def validate_git_repo_url(self, **kwargs):
kwargs = copy.deepcopy(kwargs)
repo_url = kwargs.get('repo_url')
token = kwargs.get('token')
# Regular expression to match GitHub repository URLs
pattern = r'https://github\.com/[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+/?'
match = re.match(pattern, repo_url)
# The URL is valid if there's a match and it covers the whole URL
return match is not None and match.group() == repo_url
@ETK_git_base
class CloneRepoNode:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"repo_url": ("STRING", {"multiline": False, "default": ""}),
"path": ("STRING", {"multiline": False, "default": ""}),
# "path_to_key": ("STRING", {"multiline": False, "default": ""}),
}
}
RETURN_TYPES = ("STRING", "STRING", "LIST", "STRING",)
RETURN_NAMES = ("path_passthrough", "file_structure", "file_list", "STD_OUT+ERR",)
CATEGORY = 'ETK/git'
def func(self, **kwargs):
kwargs = copy.deepcopy(kwargs)
repo_url = kwargs.get('repo_url')
path = kwargs.get('path')
# token = kwargs.get('token')
# lets just do it via calling the terminal
# working = f"GIT_SSH_COMMAND='ssh -i {path_to_key}' git clone https://github.com/story-squad/GDE_Graph_IO.git {path}"
working = f"git clone {repo_url} {path}"
import os
import subprocess
# execute the command and retrieve the results
process = subprocess.Popen(working, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
out, err = process.communicate()
ld = os.listdir(path)
return (path, "\n".join(ld), ld, out.decode("utf-8") + err.decode("utf-8"),)
@ETK_git_base
class PullRepoNode:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"repo_path": ("STRING", {"multiline": False, "default": ""}),
"token": ("STRING", {"multiline": False, "default": ""}),
}
}
RETURN_TYPES = ()
FUNCTION = 'pull_repo'
CATEGORY = 'ETK/git'
def pull_repo(self, **kwargs):
kwargs = copy.deepcopy(kwargs)
repo_path = kwargs.get('repo_path')
token = kwargs.get('token')
repo = Repo(repo_path)
origin = repo.remotes.origin
# with repo.git(GIT_ASKPASS='echo', GIT_USERNAME=token):
origin.pull()
# with repo.git.custom_environment(GIT_ASKPASS='echo', GIT_USERNAME=token):
# origin.pull()
@ETK_git_base
class PushRepoNode:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"repo_path": ("STRING", {"multiline": False, "default": ""}),
"token": ("STRING", {"multiline": False, "default": ""}),
}
}
RETURN_TYPES = ()
FUNCTION = 'push_repo'
CATEGORY = 'ETK/git'
def push_repo(self, **kwargs):
kwargs = copy.deepcopy(kwargs)
repo_path = kwargs.get('repo_path')
token = kwargs.get('token')
repo = Repo(repo_path)
origin = repo.remotes.origin
# with repo.git.custom_environment(GIT_ASKPASS='echo', GIT_USERNAME=token):
origin.push()
@ETK_git_base
class SendCommitNode:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"repo_path": ("STRING", {"multiline": False, "default": ""}),
"commit_message": ("STRING", {"multiline": False, "default": ""}),
"token": ("STRING", {"multiline": False, "default": ""}),
}
}
RETURN_TYPES = ()
FUNCTION = 'send_commit'
CATEGORY = 'ETK/git'
def send_commit(self, **kwargs):
kwargs = copy.deepcopy(kwargs)
repo_path = kwargs.get('repo_path')
commit_message = kwargs.get('commit_message')
repo = Repo(repo_path)
repo.git.add(A=True)
repo.index.commit(commit_message)
@ETK_git_base
class CreateRepoNode:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"path": ("STRING", {"multiline": False, "default": ""}),
"token": ("STRING", {"multiline": False, "default": ""}),
}
}
RETURN_TYPES = ()
FUNCTION = 'create_repo'
CATEGORY = 'ETK/git'
def create_repo(self, **kwargs):
kwargs = copy.deepcopy(kwargs)
path = kwargs.get('path')
Repo.init(path)
hf_diffusers_nodes.py
-278
View File
@@ -1,278 +0,0 @@
import torch
import folder_paths
import torchvision
import torch.nn as nn
import functools
import torch
import json
try:
from .config import config_settings
except ImportError as e:
from config import config_settings
#NODE_CLASS_MAPPINGS = {}
#NODE_DISPLAY_NAME_MAPPINGS = {}
from . import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
def ETK_HF_Diffusers_base(cls):
cls.CATEGORY = "ETK/HF/Diffusers"
# Add spaces to the camel case class name
pretty_name = cls.__name__
for i in range(1, len(pretty_name)):
if pretty_name[i].isupper():
pretty_name = pretty_name[:i] + " " + pretty_name[i:]
cls.DISPLAY_NAME = pretty_name
NODE_CLASS_MAPPINGS[cls.__name__] = cls
NODE_DISPLAY_NAME_MAPPINGS[cls.DISPLAY_NAME] = pretty_name
# Wrap the function defined in the FUNCTION attribute
func_name = getattr(cls, "FUNCTION", None)
if func_name and hasattr(cls, func_name):
original_func = getattr(cls, func_name)
@functools.wraps(original_func)
def wrapped_func(*args, **kwargs):
with torch.inference_mode(False):
return original_func(*args, **kwargs)
setattr(cls, func_name, wrapped_func)
return cls
@ETK_HF_Diffusers_base
class DDPMPipline:
"""A node that loads the specified DDPMPipeline, supports only from_pretrained for now"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model_name": ("STRING", {"default": "johnowhitaker/ddpm-butterflies-32px"}),
"device": (["cpu", "cuda"], {"default": "cuda"}),
},
}
RETURN_TYPES = ("DDPM_PL",)
FUNCTION = "load_from_pretrained"
def load_from_pretrained(self, model_name, device):
from diffusers import DDPMPipeline
ret = DDPMPipeline.from_pretrained(model_name, device=device)
return (ret,)
@ETK_HF_Diffusers_base
class SampleDDPMPipline:
"""A node that samples from the specified DDPMPipeline"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": ("DDPM_PL", {}),
"batch_size": ("INT", {"default": 1}),
"num_timesteps": ("INT", {"default": 1000}),
},
}
RETURN_TYPES = ("TORCH_TENSOR",)
FUNCTION = "sample"
def sample(
self,
model,
batch_size,
num_timesteps,
):
ret = model(
batch_size=batch_size,
generator=None,
num_inference_steps=num_timesteps,
output_type="tensor",
)
ret = torch.tensor(ret.images)
return (ret,)
@ETK_HF_Diffusers_base
class TensorToImage:
"""A node that converts a tensor to an image"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"tensor": ("TORCH_TENSOR", {}),
"incoming_format": (["BHWC", "BCHW", "HWC"], {"default": "BHWC"}
),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "tensor_to_image"
def tensor_to_image(self, tensor, incoming_format):
# outgoing format is BHWC torch tensor float 32 0-1
if incoming_format == "BCHW":
tensor = tensor.permute(0, 2, 3, 1)
elif incoming_format == "HWC":
tensor = tensor.unsqueeze(0)
tensor = tensor.clamp(0, 1)
return (tensor,)
@ETK_HF_Diffusers_base
class LoadHFDataset:
"""A node that loads the specified HFDataset"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"dataset_name": ("STRING", {"default": "huggan/smithsonian_butterflies_subset"}),
},
}
RETURN_TYPES = ("HF_DATASET",)
FUNCTION = "load_from_hf"
def load_from_hf(self, dataset_name):
from datasets import load_dataset
ret = load_dataset(dataset_name,
split="train",
)
return (ret,)
@ETK_HF_Diffusers_base
class CreateUnet2dModel:
"""A node that creates a unet2d model"""
@classmethod
def INPUT_TYPES(cls):
down_blocks = ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D")
mid_blocks = ('UNetMidBlock2D', 'UnCLIPUNetMidBlock2D')
up_blocks = ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")
return {
"required": {
"name": ("STRING", {"default": "UNet2D"}),
},
"optional": {
"sample_size": ("TUPLE", {"default": (32, 32, 3,)}),
"in_channels": ("INT", {"default": 3}),
"out_channels": ("INT", {"default": 3}),
"norm_num_groups": ("INT", {"default": 32}),
"layers_per_block": ("INT", {"default": 2}),
"down_block_types": ("TUPLE", {"default": (down_blocks[0],)}),
"mid_block_scale_factor": ("FLOAT", {"default": 1.0}),
"up_block_types": ("TUPLE", {"default": (up_blocks[0],)}),
"block_out_channels": ("TUPLE", {"default": (64, 128, 256, 512,)}),
},
}
RETURN_TYPES = ("TORCH_MODEL",)
FUNCTION = "create_unet2d_model"
def create_unet2d_model(self, name, sample_size, in_channels, out_channels, norm_num_groups, layers_per_block,
down_block_types=("AttnDownBlock2D",),
mid_block_scale_factor=("UNetMidBlock2D",),
up_block_types=("AttnUpBlock2D",),
block_out_channels=(64, 128, 256, 512,)
):
from diffusers import UNet2DModel
with torch.inference_mode(False):
ret = UNet2DModel(
sample_size=sample_size,
in_channels=in_channels,
out_channels=out_channels,
norm_num_groups=norm_num_groups,
layers_per_block=layers_per_block,
down_block_types=down_block_types,
mid_block_scale_factor=mid_block_scale_factor,
up_block_types=up_block_types,
block_out_channels=block_out_channels,
)
for param in ret.parameters():
param.requires_grad = True
#param.retains_grad = True
return (ret,)
@ETK_HF_Diffusers_base
class TrainUnet2dModel:
"""A node that trains a unet2d model"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": ("TORCH_MODEL", {}),
"dataset": ("HF_DATASET", {}),
"batch_size": ("INT", {"default": 1}),
"num_epochs": ("INT", {"default": 1}),
"learning_rate": ("FLOAT", {"default": 1e-4}),
},
}
RETURN_TYPES = ("TORCH_MODEL",)
FUNCTION = "train_unet2d_model"
def train_unet2d_model(self, model, dataset, batch_size, num_epochs, learning_rate):
from torch.utils.data import DataLoader
# import mse loss
from torch.nn import MSELoss
from torch.optim import Adam
from diffusers import UNet2DModel
model = model.train()
model = model.cuda()
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
optimizer = Adam(model.parameters(), lr=learning_rate)
loss_fn = MSELoss()
for epoch in range(num_epochs):
for batch in dataloader:
batch = batch["image"]
batch = batch.cuda()
optimizer.zero_grad()
loss = loss_fn(model, batch)
loss.backward()
optimizer.step()
print(loss.item())
return (model,)
# @ETK_HF_Diffusers_base
# class LoadModelFromHF:
# """A node that loads the specified model from huggingface"""
#
# @classmethod
# def INPUT_TYPES(cls):
# return {
# "required": {
# "model_name": ("STRING", {"default": "johnowhitaker/ddpm-butterflies-32px"}),
# },
# }
#
# RETURN_TYPES = ("TORCH_MODEL",)
# FUNCTION = "load_model_from_hf"
#
# def load_model_from_hf(self, model_name):
# from huggingface_hub import hf_hub_url
# from diffusers import DDPMPipeline
# ret = DDPMPipeline.from_pretrained(model_name, device="cuda")
# return (ret,)
image.py
-3763
View File
File diff suppressed because it is too large Load Diff
issues/runtime_errors.md
-13
View File
@@ -1,13 +0,0 @@
# Runtime errors in PyTorch nodes
The following problems were discovered:
1. **Undefined helper method and variable** (#2)
- `TrainModel.train` calls `self.make_image_tensor()` which is not defined.
- Variable `best_model` may be used before assignment.
2. **Dataset handling in inference** (#3)
- `PyTorchInferenceNode.inference` assumes input dataset has a `.dataset` attribute, raising `AttributeError` for regular datasets.
3. **FlattenDataset assumes `.data`** (#4)
- `FlattenDataset.flatten` uses `dataset.data` which is not present in generic datasets.
These should be addressed to prevent runtime failures.
latent.py
-77
View File
@@ -1,77 +0,0 @@
#from nodes import MAX_RESOLUTION
import torch
#from custom_nodes.EternalKernelLiteGraphNodes.components import fields as field
class LatentInterpolation:
def __init__(self) -> None:
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"latent1": ("LATENT",),
"latent2": ("LATENT",),
"t": ("FLOAT", {"default": 0.5}),
}
}
RETURN_TYPES = ("LATENT",)
FUNCTION = "interpolate"
CATEGORY = "ETK/Latent"
def interpolate(self, latent1, latent2, t):
# Ensure that the latents are dictionaries
if not (isinstance(latent1, dict) and isinstance(latent2, dict)):
raise ValueError("Latent values must be dictionaries.")
# Ensure that both latents have the same keys
if set(latent1.keys()) != set(latent2.keys()):
raise ValueError("Latent dictionaries must have the same keys.")
inter_latent = {}
for key in latent1:
if isinstance(latent1[key], (list, tuple)):
inter_latent[key] = [(1 - t) * l1_val + t * l2_val for l1_val, l2_val in
zip(latent1[key], latent2[key])]
elif isinstance(latent1[key], torch.Tensor):
inter_latent[key] = (1 - t) * latent1[key] + t * latent2[key]
else:
raise ValueError(f"Unsupported value type for key '{key}' in latent dictionaries.")
return (inter_latent,)
def test_latent_interpolation():
import torch
# Create an instance of the LatentInterpolation class
li = LatentInterpolation()
# Define test cases for the interpolate method
# latents are at least 4 diminsional
# latents are of the same length
test_cases = [
{
"input": {
"latent1": torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]]]),
"latent2": torch.tensor([[[[5.0, 6.0], [7.0, 8.0]]]]),
"t": 0.5,
},
"expected": torch.tensor([[[[3.0, 4.0], [5.0, 6.0]]]]),
},
]
# Run the test cases
for test_case in test_cases:
result = li.interpolate(**test_case["input"])
if not (result[0][0] == test_case["expected"]).all().item():
print(f"Expected {test_case['expected']} but got {result}")
if __name__ == "__main__":
test_latent_interpolation()
local_shared.py
-7
View File
@@ -1,7 +0,0 @@
import os
# state that is shared, available to all modules
ETK_PATH = os.path.dirname(os.path.abspath(__file__))
# this is two levels down
GDE_PATH = os.path.normpath(os.path.join(ETK_PATH, "../../"))
modules/config.json
-158
View File
@@ -1,158 +0,0 @@
{
"output_path": "output",
"logger_uri": null,
"run_name": "run",
"project_name": null,
"run_description": "\ud83d\udc38Coqui trainer run.",
"print_step": 25,
"plot_step": 100,
"model_param_stats": false,
"wandb_entity": null,
"dashboard_logger": "tensorboard",
"save_on_interrupt": true,
"log_model_step": null,
"save_step": 10000,
"save_n_checkpoints": 5,
"save_checkpoints": true,
"save_all_best": false,
"save_best_after": 10000,
"target_loss": null,
"print_eval": false,
"test_delay_epochs": 0,
"run_eval": true,
"run_eval_steps": null,
"distributed_backend": "nccl",
"distributed_url": "tcp://localhost:54321",
"mixed_precision": false,
"precision": "fp16",
"epochs": 1000,
"batch_size": 32,
"eval_batch_size": 16,
"grad_clip": 0.0,
"scheduler_after_epoch": true,
"lr": 0.001,
"optimizer": "radam",
"optimizer_params": null,
"lr_scheduler": null,
"lr_scheduler_params": {},
"use_grad_scaler": false,
"allow_tf32": false,
"cudnn_enable": true,
"cudnn_deterministic": false,
"cudnn_benchmark": false,
"training_seed": 54321,
"model": "xtts",
"num_loader_workers": 0,
"num_eval_loader_workers": 0,
"use_noise_augment": false,
"audio": {
"sample_rate": 22050,
"output_sample_rate": 24000
},
"use_phonemes": false,
"phonemizer": null,
"phoneme_language": null,
"compute_input_seq_cache": false,
"text_cleaner": null,
"enable_eos_bos_chars": false,
"test_sentences_file": "",
"phoneme_cache_path": null,
"characters": null,
"add_blank": false,
"batch_group_size": 0,
"loss_masking": null,
"min_audio_len": 1,
"max_audio_len": Infinity,
"min_text_len": 1,
"max_text_len": Infinity,
"compute_f0": false,
"compute_energy": false,
"compute_linear_spec": false,
"precompute_num_workers": 0,
"start_by_longest": false,
"shuffle": false,
"drop_last": false,
"datasets": [
{
"formatter": "",
"dataset_name": "",
"path": "",
"meta_file_train": "",
"ignored_speakers": null,
"language": "",
"phonemizer": "",
"meta_file_val": "",
"meta_file_attn_mask": ""
}
],
"test_sentences": [],
"eval_split_max_size": null,
"eval_split_size": 0.01,
"use_speaker_weighted_sampler": false,
"speaker_weighted_sampler_alpha": 1.0,
"use_language_weighted_sampler": false,
"language_weighted_sampler_alpha": 1.0,
"use_length_weighted_sampler": false,
"length_weighted_sampler_alpha": 1.0,
"model_args": {
"gpt_batch_size": 1,
"enable_redaction": false,
"kv_cache": true,
"gpt_checkpoint": null,
"clvp_checkpoint": null,
"decoder_checkpoint": null,
"num_chars": 255,
"tokenizer_file": "",
"gpt_max_audio_tokens": 605,
"gpt_max_text_tokens": 402,
"gpt_max_prompt_tokens": 70,
"gpt_layers": 30,
"gpt_n_model_channels": 1024,
"gpt_n_heads": 16,
"gpt_number_text_tokens": 6153,
"gpt_start_text_token": null,
"gpt_stop_text_token": null,
"gpt_num_audio_tokens": 1026,
"gpt_start_audio_token": 1024,
"gpt_stop_audio_token": 1025,
"gpt_code_stride_len": 1024,
"gpt_use_masking_gt_prompt_approach": true,
"gpt_use_perceiver_resampler": true,
"input_sample_rate": 22050,
"output_sample_rate": 24000,
"output_hop_length": 256,
"decoder_input_dim": 1024,
"d_vector_dim": 512,
"cond_d_vector_in_each_upsampling_layer": true,
"duration_const": 102400
},
"model_dir": null,
"languages": [
"en",
"es",
"fr",
"de",
"it",
"pt",
"pl",
"tr",
"ru",
"nl",
"cs",
"ar",
"zh-cn",
"hu",
"ko",
"ja"
],
"temperature": 0.75,
"length_penalty": 1.0,
"repetition_penalty": 10.0,
"top_k": 50,
"top_p": 0.85,
"num_gpt_outputs": 1,
"gpt_cond_len": 30,
"gpt_cond_chunk_len": 4,
"max_ref_len": 30,
"sound_norm_refs": false
}
modules/fractal.py
View File
modules/image_utils.py
-441
View File
@@ -1,441 +0,0 @@
import torch
import torch.nn as nn
import cv2
import time
from PIL import Image
import numpy as np
import scipy.ndimage
class MaskExpansionModel(nn.Module):
"""
A PyTorch model that expands or erodes a mask using the provided function.
"""
def forward(self, mask, expand, tapered_corners):
assert mask.dim() == 4, "Input mask must have 3 dimensions (batch_size, height, width)."
bs = mask.size(0)
c = 0 if tapered_corners else 1
kernel = torch.tensor([[c, 1, c],
[1, 1, 1],
[c, 1, c]], dtype=mask.dtype, device=mask.device)
kernel = kernel.view(1, 1, 3, 3)
if expand < 0:
kernel = -kernel
expand = -expand
mask = 1 - mask
mask = mask.unsqueeze(1)
# print(mask.shape)
for _ in range(expand):
for i in range(bs):
# print(i)
mask[i] = torch.nn.functional.conv2d(mask[i], kernel, padding=1)
mask[i] = torch.clamp(mask[i], 0, 1) # Clip the mask values to the range [0, 1]
mask = mask.squeeze(1)
if expand < 0:
mask = 1 - mask
return mask
class LanczosInterpolationLayer(nn.Module):
"""
A PyTorch module that applies Lanczos interpolation to resize a batch of images.
Args:
target_size (tuple): The target size for resizing the images, in the format (width, height).
"""
def __init__(self, target_size):
super(LanczosInterpolationLayer, self).__init__()
self.target_size = target_size
def visualize_tensor(self,tensor, title):
from PIL import Image, ImageDraw, ImageFont
# Convert the tensor to a PIL image
tensor = tensor.squeeze().cpu().numpy()
tensor = (tensor * 255).astype('uint8')
image = Image.fromarray(tensor)
# Create a new image for the text
text_image = Image.new('RGB', (200, 100), color='white')
draw = ImageDraw.Draw(text_image)
font = ImageFont.truetype("arial.ttf", 16)
# Add the title, min, and max values to the text image
min_val = tensor.min()
max_val = tensor.max()
text = f"{title}\nMin: {min_val:.6f}\nMax: {max_val:.6f}"
draw.text((10, 10), text, fill='black', font=font)
# Combine the original image and the text image
combined_image = Image.new('RGB', (image.width + text_image.width, max(image.height, text_image.height)),
color='white')
combined_image.paste(image, (0, 0))
combined_image.paste(text_image, (image.width, 0))
# Display the combined image
combined_image.show()
def forward(self, x):
"""
Performs the forward pass of the Lanczos interpolation layer.
Args:
x (torch.Tensor): A batch of input images with shape (batch_size, channels, height, width).
Returns:
torch.Tensor: A batch of resized images with shape (batch_size, channels, target_height, target_width).
"""
device = torch.device("cpu")
interpolated_images = []
for image in x:
# Convert the PyTorch tensor to a NumPy array
image_np = image.permute(1, 2, 0).to("cpu", torch.float32).numpy()
# Resize the image using OpenCV's resize function with Lanczos interpolation
interpolated_image = cv2.resize(image_np, self.target_size, interpolation=cv2.INTER_LANCZOS4)
mn = interpolated_image.min()
mx = interpolated_image.max()
color_range = mx - mn
under_bottom_target_range = -10 // 255
over_top_target_range = 10 // 255
#self.visualize_tensor(torch.from_numpy(interpolated_image), "Original Image")
if color_range > 1:
# here we are going to compress the over the top values
over_top = interpolated_image > 1
not_over_top = interpolated_image <= 1
over_top_range = mx - 1
over_top_scale = over_top_target_range / over_top_range
# scale the over the top values to the target range
interpolated_image[over_top] = 1 + (interpolated_image[over_top] - 1) * over_top_scale
#self.visualize_tensor(torch.from_numpy(interpolated_image), "After Over-the-Top Compression")
# scale the non-over-the top to make room
interpolated_image[not_over_top] = interpolated_image[not_over_top] * (1 - over_top_target_range)
#self.visualize_tensor(torch.from_numpy(interpolated_image), "After Non-Over-the-Top Scaling")
# here we are going to compress the under the bottom values
under_bottom = interpolated_image < 0
not_under_bottom = interpolated_image >= 0
under_bottom_range = -mn
under_bottom_scale = under_bottom_target_range / under_bottom_range
# scale the under the bottom values to the target range
#print("Max value:", interpolated_image.max())
#print("Min value:", interpolated_image.min())
interpolated_image[under_bottom] = interpolated_image[under_bottom] * under_bottom_scale
#self.visualize_tensor(torch.from_numpy(interpolated_image), "After Under-the-Bottom Compression")
# scale and shift the non-under-the-bottom to make room
#interpolated_image[not_under_bottom] = interpolated_image[not_under_bottom] * (
# 1 - over_top_target_range) + under_bottom_target_range
#self.visualize_tensor(torch.from_numpy(interpolated_image),"After Non-Under-the-Bottom Scaling and Shifting")
#print("Max value:", interpolated_image.max())
#print("Min value:", interpolated_image.min())
# clamp
interpolated_image = np.clip(interpolated_image, 0, 1)
else:
pass
# Convert the resized image back to a PyTorch tensor
# interpolated_image = torch.from_numpy(interpolated_image).to(device, dtype=torch.float16).permute(2, 0, 1)
interpolated_image = torch.from_numpy(interpolated_image)
# convert fix range to 0-255
interpolated_image = interpolated_image * 255
# convert to uint8
interpolated_image = interpolated_image.to(device, dtype=torch.uint8).permute(2, 0, 1)
interpolated_images.append(interpolated_image)
del image
del image_np
del interpolated_image
# transform this list of tensors into a single tensor of shape (batch_size, channels, height, width)
# pre allocate a tensor of shape (batch_size, channels, target_height, target_width)
y = torch.zeros((x.shape[0], interpolated_images[0].shape[0], self.target_size[1], self.target_size[0]),
device=device,
dtype=interpolated_images[0].dtype)
y = torch.stack(interpolated_images, dim=0, out=y)
del interpolated_images
return y
class LanczosInterpolationModel(nn.Module):
"""
A PyTorch model that applies Lanczos interpolation to resize a batch of images.
Args:
target_size (tuple): The target size for resizing the images, in the format (width, height).
"""
def __init__(self, target_size):
super(LanczosInterpolationModel, self).__init__()
self.interpolation_layer = LanczosInterpolationLayer(target_size)
def forward(self, x):
"""
Performs the forward pass of the Lanczos interpolation model.
Args:
x (torch.Tensor): A batch of input images with shape (batch_size, channels, height, width).
Returns:
torch.Tensor: A batch of resized images with shape (batch_size, channels, target_height, target_width).
"""
return self.interpolation_layer(x)
def expand_mask(mask, expand, tapered_corners, device=None, dtype=None):
"""
Convenience function for expanding or eroding a mask using the MaskExpansionModel.
Args:
mask (torch.Tensor): The input mask tensor.
expand (int): The number of iterations to expand (positive) or erode (negative) the mask.
tapered_corners (bool): Whether to use tapered corners when expanding or eroding the mask.
device (torch.device, optional): The device to use for the model and input tensor.
If not provided, it will use the default device (CPU).
dtype (torch.dtype, optional): The data type to use for the input tensor.
If not provided, it will use the default dtype (float32).
Returns:
torch.Tensor: The expanded or eroded mask tensor.
"""
# Create an instance of the MaskExpansionModel
model = MaskExpansionModel()
# Move the model to the specified device (or default to CPU)
device = device or torch.device("cpu")
model = model.to(device)
# Move the input tensor to the specified device and convert to the specified dtype
dtype = dtype or torch.float32
mask = mask.to(device, dtype=dtype)
# Expand or erode the mask using the MaskExpansionModel
with torch.no_grad():
expanded_mask = model(mask, expand, tapered_corners)
# Clear the GPU cache if using a GPU
if device != torch.device("cpu"):
torch.cuda.empty_cache()
return expanded_mask
def lanczos_resize(image, target_size, device="cpu", dtype=torch.float32, output_type="tensor"):
"""
Convenience function for resizing an image using the Lanczos interpolation model.
Args:
image (PIL.Image or torch.Tensor): The input image to be resized.
target_size (tuple): The target size for resizing the image, in the format (width, height).
device (torch.device, optional): The device to use for the model and input tensor.
If not provided, it will use the default device (CPU).
dtype (torch.dtype, optional): The data type to use for the input tensor.
If not provided, it will use the default dtype (float32).
output_type (str, optional): The type of the output image.
If "tensor", the function will return a PyTorch tensor.
If "pil", the function will return a PIL Image.
If not provided, it will return a tensor by default.
Returns:
torch.Tensor or PIL.Image: The resized image.
"""
import gc
import os
import sys
if sys.platform == "win32":
import psutil
print(f"Free memory: {psutil.virtual_memory().available / 1024 / 1024} MB")
else:
# works on linux and mac
print(f"Free memory: {os.popen('free -m').readlines()[1].split()[3]} MB")
# Create an instance of the LanczosInterpolationModel
model = LanczosInterpolationModel(target_size)
# Move the model to the specified device (or default to CPU)
device = device or torch.device("cpu")
model = model.to(device)
# Convert the input image to a PyTorch tensor
if isinstance(image, list):
if isinstance(image[0], Image.Image):
try:
images = np.array(image)
except:
images = np.array([np.array(img) for img in image])
images = torch.from_numpy(images).to(device, dtype).permute(0, 3, 1, 2)
elif isinstance(image, torch.Tensor):
images = image.to(device, dtype)
del image
torch.cuda.empty_cache()
# Resize the image using the Lanczos interpolation model
with torch.no_grad():
resized_images = model(images)
# Convert the resized tensor back to a PIL Image or keep it as a tensor based on the output_type
if output_type.lower() == "pil":
resized_images = resized_images.squeeze(0).permute(0, 2, 3, 1).cpu().numpy()
resized_images = Image.fromarray(resized_images.astype(np.uint8))
elif output_type.lower() == "tensor":
pass # keep the resized_image as a tensor
else:
raise ValueError(f"Invalid output_type: {output_type}. Valid options are 'tensor' or 'pil'.")
# Clear the GPU cache if using a GPU
if device != torch.device("cpu"):
torch.cuda.empty_cache()
del images
del model
gc.collect()
# print current free memory on cpu, in a way that is cross platform 1 way that works on all
if sys.platform == "win32":
import psutil
print(f"Free memory: {psutil.virtual_memory().available / 1024 / 1024} MB")
else:
# works on linux and mac
print(f"Free memory: {os.popen('free -m').readlines()[1].split()[3]} MB")
return resized_images
def test_lanczos_interpolation_model():
"""
Tests the Lanczos interpolation model by resizing a batch of random images and benchmarking the performance.
"""
# Set the target sizes for resizing
hd_size = (1920, 1080)
sd_size = (1280, 720)
# Generate random PIL images
batch_size = 600
hd_images = [Image.fromarray(np.random.randint(0, 256, size=(hd_size[1], hd_size[0], 3), dtype=np.uint8))
for _ in range(batch_size)]
sd_images = [Image.fromarray(np.random.randint(0, 256, size=(sd_size[1], sd_size[0], 3), dtype=np.uint8))
for _ in range(batch_size)]
# Perform the benchmarks
num_iterations = 3
# HD to SD benchmark
start_time = time.time()
for _ in range(num_iterations):
sd_output_images = lanczos_resize(hd_images, sd_size)
end_time = time.time()
hd_to_sd_elapsed_time = end_time - start_time
hd_to_sd_avg_time_per_iteration = hd_to_sd_elapsed_time / num_iterations
# SD to HD benchmark
start_time = time.time()
for _ in range(num_iterations):
hd_output_images = lanczos_resize(sd_images, hd_size)
end_time = time.time()
sd_to_hd_elapsed_time = end_time - start_time
sd_to_hd_avg_time_per_iteration = sd_to_hd_elapsed_time / num_iterations
# Check the sizes of the output images
assert sd_output_images.shape[0] == batch_size
assert hd_output_images.shape[0] == batch_size
print(f"Lanczos interpolation test passed!")
print(f"Benchmark results:")
print(f"Number of iterations: {num_iterations}")
print(f"HD to SD:")
print(f" Total elapsed time: {hd_to_sd_elapsed_time:.4f} seconds")
print(f" Average time per iteration: {hd_to_sd_avg_time_per_iteration:.4f} seconds")
print(f" Average time per HD image -> SD image: {hd_to_sd_avg_time_per_iteration / batch_size:.4f} seconds")
print(f"SD to HD:")
print(f" Total elapsed time: {sd_to_hd_elapsed_time:.4f} seconds")
print(f" Average time per iteration: {sd_to_hd_avg_time_per_iteration:.4f} seconds")
print(f" Average time per SD image -> HD image: {sd_to_hd_avg_time_per_iteration / batch_size:.4f} seconds")
def test_mask_expansion_model():
"""
Tests the MaskExpansionModel by expanding and eroding random masks and benchmarking the performance.
"""
# Generate random masks
batch_size = 1000
mask_size = (512, 512)
# masks are 0-1 tensors float
masks = torch.rand((batch_size, mask_size[0], mask_size[1])).to("cuda")
print(masks.min(), masks.max())
# Perform the benchmarks
num_iterations = 1
expand_steps = 5
erode_steps = -5
# Mask expansion benchmark
start_time = time.time()
for _ in range(num_iterations):
print(masks.mean())
expanded_masks = expand_mask(masks, expand_steps, tapered_corners=True, device="cuda")
print(masks.mean())
end_time = time.time()
expansion_elapsed_time = end_time - start_time
expansion_avg_time_per_iteration = expansion_elapsed_time / num_iterations
# Mask erosion benchmark
start_time = time.time()
for _ in range(num_iterations):
print(masks.mean())
eroded_masks = expand_mask(masks, erode_steps, tapered_corners=True, device="cuda")
print(eroded_masks.mean())
end_time = time.time()
erosion_elapsed_time = end_time - start_time
erosion_avg_time_per_iteration = erosion_elapsed_time / num_iterations
# Check the sizes of the output masks
expected_mask_size = (batch_size, mask_size[0], mask_size[1])
assert expanded_masks.shape == expected_mask_size, f"Expected expanded mask size: {expected_mask_size}, but got: {expanded_masks.shape}"
assert eroded_masks.shape == expected_mask_size, f"Expected eroded mask size: {expected_mask_size}, but got: {eroded_masks.shape}"
print(f"Mask expansion/erosion test passed!")
print(f"Benchmark results:")
print(f"Number of iterations: {num_iterations}")
print(f"Mask expansion:")
print(f" Total elapsed time: {expansion_elapsed_time:.4f} seconds")
print(f" Average time per iteration: {expansion_avg_time_per_iteration:.4f} seconds")
print(f" Average time per mask expansion: {expansion_avg_time_per_iteration / batch_size:.4f} seconds")
print(f"Mask erosion:")
print(f" Total elapsed time: {erosion_elapsed_time:.4f} seconds")
print(f" Average time per iteration: {erosion_avg_time_per_iteration:.4f} seconds")
print(f" Average time per mask erosion: {erosion_avg_time_per_iteration / batch_size:.4f} seconds")
if __name__ == "image_utils":
print("image_utils module loaded.")
print("Testing Lanczos interpolation model...")
test_lanczos_interpolation_model()
print("Testing Mask expansion model...")
# test_mask_expansion_model()
exit(0)
if __name__ == "__main__":
test_lanczos_interpolation_model()
# test_mask_expansion_model()
modules/musicgencli.py
-224
View File
@@ -1,224 +0,0 @@
def set_all_seeds(seed):
import random
import os
import numpy as np
import torch
old_values = {}
old_values['random'] = random.getstate()
old_values['numpy'] = np.random.get_state()
old_values['torch'] = torch.get_rng_state()
old_values['torch_cuda'] = torch.cuda.get_rng_state()
old_values['torch_cuda_all'] = torch.cuda.get_rng_state_all()
old_values['torch_deterministic'] = torch.backends.cudnn.deterministic
# From https://gist.github.com/gatheluck/c57e2a40e3122028ceaecc3cb0d152ac
random.seed(seed)
os.environ["PYTHONHASHSEED"] = str(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.backends.cudnn.deterministic = True
return old_values
def reset_seeds(old_values):
import random
import numpy as np
import torch
random.setstate(old_values['random'])
np.random.set_state(old_values['numpy'])
torch.set_rng_state(old_values['torch'])
torch.cuda.set_rng_state(old_values['torch_cuda'])
torch.cuda.set_rng_state_all(old_values['torch_cuda_all'])
torch.backends.cudnn.deterministic = old_values['torch_deterministic']
import time # Add this at the top of your script
def generate_music_for_description(description, duration, file, model_name, position):
time.sleep(position * 5) # Wait for position * 5 seconds
from audiocraft.models import MusicGen
from audiocraft.data.audio import audio_write
import torch
model = MusicGen.get_pretrained(model_name)
print(f"Generating music for description: {description}")
print(f"Using model: {model_name}")
print(f"Using duration: {duration}")
model.set_generation_params(duration=duration)
wav = model.generate([description])
print(f"Writing to {file}")
audio_write(file.replace(".wav", ""), wav[0].cpu(), model.sample_rate, strategy="loudness")
del model
torch.cuda.empty_cache()
def generate_music(**kwargs):
from audiocraft.models import MusicGen
from audiocraft.data.audio import audio_write
import torch
import ast
import copy
import random
from concurrent.futures import ProcessPoolExecutor, as_completed
torch.cuda.empty_cache()
# Set seed, -1 means random seed
seed = kwargs.get('seed', -1)
if seed == -1:
seed = random.randint(0, 100000)
old_seed_values = set_all_seeds(seed)
kwargs = copy.deepcopy(kwargs)
description = kwargs.get('description')
file = kwargs.get('file')
# Add .wav if not there
if not file.endswith('.wav'):
file = file + ".wav"
model_name = kwargs.get('model_name', 'melody')
duration = kwargs.get('duration', 10)
descriptions = ast.literal_eval(description)
durations = ast.literal_eval(duration)
o_files = []
batch_size = 3
with ProcessPoolExecutor(max_workers=batch_size) as executor:
futures = []
for i, (description, duration) in enumerate(zip(descriptions, durations)):
ii = str(i).zfill(1 + len(str(1 + len(descriptions))))
file_to_use = file.replace(".wav", f"_{ii}.wav")
# Pass i (the position in the batch) as an additional argument
future = executor.submit(generate_music_for_description, description, duration, file_to_use, model_name, i//batch_size)
futures.append((future, file_to_use))
for future, file_to_use in futures:
# Wait for the future to complete, no need to get the result since the function handles file writing
future.result()
o_files.append(file_to_use)
# Reset seed
reset_seeds(old_seed_values)
torch.cuda.empty_cache()
return (o_files[0], o_files,)
import argparse
import copy
def main():
print("Running main function...")
# Define the model names as a list for the choices parameter
model_names = [
"facebook/musicgen-melody-large",
"facebook/musicgen-melody-small",
"facebook/musicgen-melody-medium",
"facebook/musicgen-stereo-melody-large",
"facebook/musicgen-stereo-melody-small",
"facebook/musicgen-stereo-melody-medium",
"facebook/musicgen-large",
"facebook/musicgen-small",
"facebook/musicgen-medium",
"facebook/musicgen-stereo-large",
"facebook/musicgen-stereo-small",
"facebook/musicgen-stereo-medium",
]
# Parse command line arguments
parser = argparse.ArgumentParser(description='Generate music based on a given description.')
parser.add_argument('--description', help='Description of the music to generate')
parser.add_argument('--file', type=str, default='output.wav', help='Output file name (default: output.wav)')
parser.add_argument('--model_name', type=str, default='facebook/musicgen-large', choices=model_names,
help='Model name to use for generation')
parser.add_argument('--duration', help='Duration of the generated music in seconds (default: 10)')
parser.add_argument('--seed', type=int, default=-1,
help='Seed for random number generation (default: -1 for random seed)')
parser.add_argument('--test', action='store_true', help='Run the test suite instead of generating music')
args = parser.parse_args()
# Convert argparse.Namespace to dictionary
kwargs = vars(args)
if args.test:
run_tests()
else:
print(f"Generating music with the following arguments: {kwargs}")
# Generate music
generated_files = generate_music(**kwargs)
print(f"Generated music files: {generated_files}")
def run_tests():
print("Running test suite...")
# Define your test cases here
test_model_name_validation()
test_file_generation()
# Add more tests as needed
print("All tests passed!")
def test_model_name_validation():
print("Testing model name validation...")
model_names = [
"facebook/musicgen-melody-large",
"facebook/musicgen-melody-small",
"facebook/musicgen-melody-medium",
"facebook/musicgen-stereo-melody-large",
"facebook/musicgen-stereo-melody-small",
"facebook/musicgen-stereo-melody-medium",
"facebook/musicgen-large",
"facebook/musicgen-small",
"facebook/musicgen-medium",
"facebook/musicgen-stereo-large",
"facebook/musicgen-stereo-small",
"facebook/musicgen-stereo-medium",
]
description = "A brief test piece" # Example description for testing
duration = 5 # 5 seconds long
file_template = "test_output_{}.wav" # Template for generated file names
for i, model_name in enumerate(model_names):
print(f"Testing {model_name}...")
kwargs = {
'description': description,
'duration': duration,
'model_name': model_name,
'file': file_template.format(i), # Ensure unique file names
'seed': 42 # Fixed seed for consistency, if needed
}
try:
# Assuming generate_music is adapted to work with these arguments
generate_music(**kwargs)
print(f"Successfully tested {model_name}")
except Exception as e:
print(f"Error testing {model_name}: {e}")
# Handle the error appropriately - for example, you might want to fail the test or log the error.
print("Model name validation test completed.")
def test_file_generation():
# Example test function
print("Testing file generation...")
# Implement the logic to test music file generation
# This could involve generating a small piece of music and checking if the file exists
# ...
if __name__ == "__main__":
main()
modules/noise.py
-87
View File
@@ -1,87 +0,0 @@
# TODO: replace pyfastnoisesimd with pyfastnoiselite (works with rpi)
import pyfastnoisesimd as fns
def get_perlin_2d(width, height, scale, octaves=6, lacunarity=3, gain=.5, seed=0, offset=0, workers=4) -> any:
"""
Returns a 2D fractal perlin noise array of size (width, height) with the given parameters.
"""
noise = fns.Noise(seed=seed, numWorkers=workers)
noise.frequency = scale
# noise.fractal.fractalType = fns.FractalType.RigidMulti
noise.fractal.octaves = octaves
noise.fractal.lacunarity = lacunarity
noise.fractal.gain = gain
noise.noiseType = fns.NoiseType.PerlinFractal
noise.fractal.offset = offset
return noise.genAsGrid([width, height])
def get_perlin_2d_as_torch_image_stack(width, height, scale, octaves=6, lacunarity=3, gain=.5, seed=0, offset=0,
workers=4, norm_zero_to_1=True) -> any:
"""
just uses get_perlin_2d and converts it to a torch image stack
(B,H,W,C=3) format
"""
import torch
import numpy as np
noise = get_perlin_2d(width, height, scale, octaves, lacunarity, gain, seed, offset, workers)
# noise = np.transpose(noise, (2, 0, 1))
noise = torch.from_numpy(noise)
# noise = noise.unsqueeze(0)
noise = torch.reshape(noise, (1, height, width, 1))
# now stack it 3 times
# noise = torch.normal(noise)
# use torch to put it from 0-1
# noise = torch.sigmoid(noise)
# normalize it using torch
if norm_zero_to_1:
# noise = noise.view(noise.size(0), -1)
noise -= noise.min()
noise /= noise.max()
# noise = noise + 0.5
# noise = noise / 1.0
return noise
if __name__ == '__main__':
import matplotlib.pyplot as plt
import time
for i in range(10):
start = time.time()
# noise = get_perlin_2d(1024, 1024, 0.001, i + 1, 3, .5, 1, 0)
# noise = get_perlin_2d(1024*10, 1024*10, i/5, i + 1, 3, .5, i, 0)
# print(type(noise))
# print(f"Time taken: {time.time() - start}")
# print the min and max
# print(f"Min: {noise.min()} Max: {noise.max()}")
# plt.imshow(noise, cmap='gray')
# plt.show()
start = time.time()
for i in range(10):
# now test the torch version
noise = get_perlin_2d_as_torch_image_stack(1024,
1024,
0.001,
6,
3,
.5,
1,
0,
4,
2)
print(f"Time taken: {time.time() - start}")
modules/speaker.mp3
BIN
View File
Binary file not shown.
modules/test.speech
-3
View File
@@ -1,3 +0,0 @@
This is a test.
Of the emergency AI TTS system.
This is only a test.
modules/vocab.json
-14470
View File
File diff suppressed because it is too large Load Diff
modules/xttscli.py
-219
View File
@@ -1,219 +0,0 @@
import sys
import os
import argparse
import logging
import traceback
import torch
import torchaudio
import tempfile
from pathlib import Path
from TTS.tts.configs.xtts_config import XttsConfig
from TTS.tts.models.xtts import Xtts
XTTS_MODEL = None
def prepare_speaker_audio(speaker):
"""
Prepares the speaker audio file and returns the speaker file path.
For .wav and .mp3 files, it directly returns the file path for model processing.
"""
if speaker.endswith(".wav") or speaker.endswith(".mp3"):
print(f"Using speaker audio file: {speaker}")
return speaker, None # No precomputed speaker embedding
elif speaker.endswith(".pth"):
print(f"Loading preprocessed speaker embedding from .pth file: {speaker}")
speaker_embedding = torch.load(speaker)
return None, speaker_embedding # Return precomputed speaker embedding
else:
raise ValueError("Unsupported speaker file format. Please provide a .wav, .mp3, or .pth file.")
def run_tts(lang, tts_text, speaker_audio_file, speaker_embedding, device):
"""
Generates TTS using the XTTS model, either from speaker audio or precomputed speaker embedding.
"""
if XTTS_MODEL is None:
raise ValueError("You need to load the model.")
# Generate conditioning latents from speaker audio if provided
if speaker_audio_file is not None:
print(f"Using speaker audio file: {speaker_audio_file}")
gpt_cond_latent, speaker_embedding = XTTS_MODEL.get_conditioning_latents(
audio_path=speaker_audio_file,
gpt_cond_len=XTTS_MODEL.config.gpt_cond_len,
max_ref_length=XTTS_MODEL.config.max_ref_len,
sound_norm_refs=XTTS_MODEL.config.sound_norm_refs
)
elif speaker_embedding is not None:
raise ValueError("Precomputed speaker embeddings are not directly usable without gpt_cond_latent.")
# Perform inference
out = XTTS_MODEL.inference(
text=tts_text,
language=lang,
gpt_cond_latent=gpt_cond_latent,
speaker_embedding=speaker_embedding,
temperature=XTTS_MODEL.config.temperature,
length_penalty=XTTS_MODEL.config.length_penalty,
repetition_penalty=XTTS_MODEL.config.repetition_penalty,
top_k=XTTS_MODEL.config.top_k,
top_p=XTTS_MODEL.config.top_p,
enable_text_splitting=True
)
# Move the output tensor to CPU before saving
out["wav"] = torch.tensor(out["wav"]).unsqueeze(0).to("cpu") # Ensure tensor is on CPU
# Save the generated speech to a temporary file
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as fp:
out_path = fp.name
torchaudio.save(out_path, out["wav"], 24000) # Save the tensor to disk
return out_path
def clear_gpu_cache():
if torch.cuda.is_available():
torch.cuda.empty_cache()
def load_model(xtts_checkpoint, xtts_config, xtts_vocab, device):
"""
Loads the XTTS model from checkpoint, config, and vocab files.
"""
global XTTS_MODEL
clear_gpu_cache()
if not xtts_checkpoint or not xtts_config or not xtts_vocab:
raise ValueError("You need to set the XTTS checkpoint path, XTTS config path, and XTTS vocab path!")
config = XttsConfig()
config.load_json(xtts_config)
XTTS_MODEL = Xtts.init_from_config(config)
print("Loading XTTS model!")
XTTS_MODEL.load_checkpoint(config, checkpoint_path=xtts_checkpoint, vocab_path=xtts_vocab, use_deepspeed=False)
if "cuda" in device and torch.cuda.is_available():
XTTS_MODEL.cuda(device)
print(f"Model moved to GPU: {device}")
else:
print(f"Using device: {device}")
def main():
# Set up logging
log_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xtts_cli.log')
logging.basicConfig(filename=log_file, level=logging.ERROR, format='%(asctime)s %(levelname)s %(message)s')
# Default paths for the config, checkpoint, and vocab
script_dir = os.path.dirname(os.path.abspath(__file__))
default_config_path = os.path.join(script_dir, 'config.json')
default_checkpoint_path = os.path.join(script_dir, 'model.pth')
default_vocab_path = os.path.join(script_dir, 'vocab.json')
default_speaker_path = os.path.join(script_dir, 'speaker.mp3') # Default speaker audio file
default_test_file_path = os.path.join(script_dir, 'test.speech') # Default text file for --test
parser = argparse.ArgumentParser(description="CLI for XTTS model loading and TTS generation.")
parser.add_argument('--checkpoint', default=default_checkpoint_path,
help=f"Path to XTTS checkpoint file. Default: {default_checkpoint_path}")
parser.add_argument('--config', default=default_config_path,
help=f"Path to XTTS config file. Default: {default_config_path}")
parser.add_argument('--vocab', default=default_vocab_path,
help=f"Path to XTTS vocab file. Default: {default_vocab_path}")
parser.add_argument('--speaker', default=default_speaker_path,
help=f"Path to speaker audio file (raw .wav, .mp3, or preprocessed .pth). Default: {default_speaker_path}")
parser.add_argument('--text', help="Text to synthesize.")
parser.add_argument('--file', help="File containing text to synthesize (one text per line).")
parser.add_argument('--lang', default='en', help="Language for the TTS.")
parser.add_argument('--folder', help="Path to save the output TTS file.")
parser.add_argument('--device', default='cuda:0', help="Device to run the model on. Default: cuda:0")
parser.add_argument('--test', action='store_true', help="Run a test TTS using the test.speech file.")
args = parser.parse_args()
try:
# Preprocess the speaker file if it's a .wav or .mp3 file
args.speaker, speaker_embedding = prepare_speaker_audio(args.speaker)
# Handle test case
if args.test:
print("Running test with default parameters...")
if not os.path.exists(default_test_file_path):
raise FileNotFoundError(f"The test file '{default_test_file_path}' does not exist. Please create it.")
args.file = default_test_file_path
args.folder = os.path.join(script_dir, 'test_output.wav')
print(f"Test file: {args.file}")
print(f"Output will be saved to: {args.folder}")
# Handle the --file argument
if args.file:
print(f"Reading text from file: {args.file}")
with open(args.file, 'r') as f:
args.text = f.read().splitlines()
elif not args.text:
raise ValueError("Text or a file with text must be provided unless running in test mode.")
# Load the XTTS model
load_model(args.checkpoint, args.config, args.vocab, args.device)
# Prepare to store generated files
all_output_files = []
# Process multiple lines of text (if from file)
if isinstance(args.text, list):
for i, tts_text in enumerate(args.text):
output_path = run_tts(
lang=args.lang,
tts_text=tts_text,
speaker_audio_file=args.speaker,
speaker_embedding=speaker_embedding,
device=args.device
)
output_file = f"{args.folder}_{i}.wav" if args.folder else f'./output_{i}.wav'
try:
os.rename(output_path, output_file)
except Exception as e:
import shutil
# try copy then delete
shutil.copy(output_path, output_file)
os.remove(output_path)
all_output_files.append(output_file)
else:
# Generate TTS for a single text input
output_path = run_tts(
lang=args.lang,
tts_text=args.text,
speaker_audio_file=args.speaker,
speaker_embedding=speaker_embedding,
device=args.device
)
output_file = args.folder if args.folder else './output.wav'
try:
os.rename(output_path, output_file)
except Exception as e:
import shutil
# try copy then delete
shutil.copy(output_path, output_file)
os.remove(output_path)
all_output_files.append(output_file)
# Final output format to match expected behavior
main_output_file = all_output_files[0]
print(f"Generated files: ({all_output_files})")
except Exception as e:
logging.error("An error occurred: %s", traceback.format_exc())
print(f"An error occurred. Check the log file at {log_file} for details.")
if __name__ == '__main__':
main()
modules/xttscli_old.py
-219
View File
@@ -1,219 +0,0 @@
import argparse
import torch.multiprocessing as mp
from concurrent.futures import ProcessPoolExecutor, as_completed
import unittest
from unittest.mock import patch, MagicMock
def generate_xtts_for_text(text, file, lang, speaker):
from TTS.api import TTS
import torch
tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to("cuda")
tts.tts_to_file(
text=text,
speaker_wav=speaker,
language=lang,
file_path=file,
)
def generate_xtts(**kwargs):
import torch
import ast
import uuid
import os
from concurrent.futures import ProcessPoolExecutor, as_completed
torch.cuda.empty_cache()
text = kwargs.get('text')
folder = kwargs.get('folder', '.')
folder = os.path.abspath(folder)
lang = kwargs.get('lang', 'en')
speaker = kwargs.get('speaker', '')
if isinstance(text, str):
if text != None and text.startswith("[") and text.endswith("]"):
try:
texts = ast.literal_eval(text)
except SyntaxError:
texts = [text]
elif isinstance(text, list):
texts = text
o_files = []
batch_size = 4
with ProcessPoolExecutor(max_workers=batch_size) as executor:
futures = []
for i, text in enumerate(texts):
if len(text) > 1:
ii = str(i).zfill(1 + len(str(1 + len(texts))))
file_to_use = f"{folder}/tts_{ii}_{str(uuid.uuid4())}.wav"
future = executor.submit(generate_xtts_for_text, text, file_to_use, lang, speaker)
futures.append((future, file_to_use))
for future, file_to_use in futures:
future.result()
o_files.append(file_to_use)
torch.cuda.empty_cache()
return (o_files[0], o_files,)
def cache_parameters(cache_file, params):
"""Cache the input parameters to a JSON file."""
import json
with open(cache_file, 'w') as f:
json.dump(params, f)
def load_cached_parameters(cache_file):
"""Load the cached parameters from the JSON file."""
import json
import os
if os.path.exists(cache_file):
with open(cache_file, 'r') as f:
return json.load(f)
return None
def run_test():
"""Run unit tests for the module."""
unittest.main(argv=[''], exit=False)
class TestGenerateXTTS(unittest.TestCase):
"""Unit tests for generate_xtts functions."""
@patch('os.path.exists', return_value=True)
@patch('TTS.api.TTS.tts_to_file')
def test_generate_xtts_single_text(self, mock_tts_to_file, mock_exists):
"""Test generating xtts with a single text input."""
mock_tts_to_file.return_value = None
kwargs = {
'text': 'Hello, this is a test.',
'folder': './output',
'lang': 'en',
'speaker': '/path/to/speaker.wav'
}
result = generate_xtts(**kwargs)
# Assertions
self.assertEqual(len(result[1]), 1)
mock_tts_to_file.assert_called_once()
@patch('os.path.exists', return_value=True)
@patch('TTS.api.TTS.tts_to_file')
def test_generate_xtts_multiple_texts(self, mock_tts_to_file, mock_exists):
"""Test generating xtts with multiple texts."""
mock_tts_to_file.return_value = None
kwargs = {
'text': ['Hello, this is a test.', 'This is another test.'],
'folder': './output',
'lang': 'en',
'speaker': '/path/to/speaker.wav'
}
result = generate_xtts(**kwargs)
# Assertions
self.assertEqual(len(result[1]), 2)
self.assertEqual(mock_tts_to_file.call_count, 2)
def test_cache_parameters(self):
"""Test caching parameters to a file."""
import os
import json
test_cache_file = './test_cache.json'
test_params = {'text': 'sample text', 'lang': 'en'}
cache_parameters(test_cache_file, test_params)
# Verify file content
with open(test_cache_file, 'r') as f:
data = json.load(f)
self.assertEqual(data, test_params)
# Clean up
os.remove(test_cache_file)
def test_load_cached_parameters(self):
"""Test loading cached parameters from a file."""
import os
import json
test_cache_file = './test_cache.json'
test_params = {'text': 'sample text', 'lang': 'en'}
# Write test data to the file
with open(test_cache_file, 'w') as f:
json.dump(test_params, f)
# Load the parameters
loaded_params = load_cached_parameters(test_cache_file)
self.assertEqual(loaded_params, test_params)
# Clean up
os.remove(test_cache_file)
def main():
import sys
import os
import argparse
import logging
import traceback
log_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'xttscli.log')
logging.basicConfig(filename=log_file, level=logging.ERROR,
format='%(asctime)s %(levelname)s %(message)s')
cache_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'params_cache.json')
try:
parser = argparse.ArgumentParser(description='Generate from given text.')
parser.add_argument('--text', help='Text to generate for')
parser.add_argument('--folder', type=str, default='.', help='Output folder')
parser.add_argument('--lang', type=str, default='en')
parser.add_argument('--speaker', type=str, default='', help='Speaker wav to use for generation')
parser.add_argument('--file', type=str, default="", help='Newline separated file.')
parser.add_argument('--rerun', action='store_true', help='Rerun with the last cached parameters.')
parser.add_argument('--test', action='store_true', help='Run the test.')
args = parser.parse_args()
kwargs = vars(args)
if args.test:
run_test()
sys.exit(0)
if args.rerun:
cached_params = load_cached_parameters(cache_file)
if cached_params:
print(f"Rerunning with cached parameters: {cached_params}")
kwargs.update(cached_params)
else:
print("No cached parameters found. Exiting.")
sys.exit(1)
else:
cache_parameters(cache_file, kwargs)
if kwargs['file']:
with open(kwargs['file']) as f:
kwargs['text'] = f.read().split('\n')
if kwargs['speaker']:
if not os.path.exists(kwargs['speaker']):
raise ValueError(f"Speaker file {kwargs['speaker']} does not exist.")
generated_files = generate_xtts(**kwargs)
print(f"Generated files: {generated_files}")
except Exception as e:
logging.error("An error occurred: %s", traceback.format_exc())
print(f"An error occurred. Check the log file at {log_file} for details.")
if __name__ == '__main__':
main()
modules/youtube.py
-184
View File
@@ -1,184 +0,0 @@
import httplib2
import os
import random
import sys
import time
from apiclient.discovery import build
from apiclient.errors import HttpError
from apiclient.http import MediaFileUpload
from oauth2client.client import flow_from_clientsecrets
from oauth2client.file import Storage
from oauth2client.tools import argparser, run_flow
from googleapiclient.http import MediaFileUpload
# Explicitly tell the underlying HTTP transport library not to retry, since
# we are handling retry logic ourselves.
httplib2.RETRIES = 1
# Maximum number of times to retry before giving up.
MAX_RETRIES = 10
# Always retry when these exceptions are raised.
RETRIABLE_EXCEPTIONS = (httplib2.HttpLib2Error, IOError)
# Always retry when an apiclient.errors.HttpError with one of these status
# codes is raised.
RETRIABLE_STATUS_CODES = [500, 502, 503, 504]
try:
import comfy
except ModuleNotFoundError:
print("comfy not found, using fallback")
comfy = lambda: None
comfy.__path__ = lambda: None
comfy.__path__._path = ["."]
comfy_dir = comfy.__path__._path[0]
base_dir = os.path.dirname(comfy_dir)
# this_dir
this_file_path = os.path.dirname(os.path.abspath(__file__))
CLIENT_SECRETS_FILE = os.path.join(base_dir, "google-credentials.json")
CLIENT_SECRETS_FILE = os.path.join(this_file_path, "youtube.py-oauth2-secrets.json")
# This OAuth 2.0 access scope allows an application to upload files to the
# authenticated user's YouTube channel, but doesn't allow other types of access.
YOUTUBE_UPLOAD_SCOPE = "https://www.googleapis.com/auth/youtube.upload"
YOUTUBE_API_SERVICE_NAME = "youtube"
YOUTUBE_API_VERSION = "v3"
# This variable defines a message to display if the CLIENT_SECRETS_FILE is
# missing.
MISSING_CLIENT_SECRETS_MESSAGE = """
WARNING: Please configure OAuth 2.0
To make this sample run you will need to populate the client_secrets.json file
found at:
%s
with information from the API Console
https://console.cloud.google.com/
For more information about the client_secrets.json file format, please visit:
https://developers.google.com/api-client-library/python/guide/aaa_client_secrets
""" % os.path.abspath(os.path.join(os.path.dirname(__file__),
CLIENT_SECRETS_FILE))
VALID_PRIVACY_STATUSES = ("public", "private", "unlisted")
def get_authenticated_service(args):
flow = flow_from_clientsecrets(CLIENT_SECRETS_FILE,
scope=YOUTUBE_UPLOAD_SCOPE,
message=MISSING_CLIENT_SECRETS_MESSAGE)
storage = Storage("%s-oauth2.json" % sys.argv[0])
credentials = storage.get()
if credentials is None or credentials.invalid:
credentials = run_flow(flow, storage, args)
return build(YOUTUBE_API_SERVICE_NAME, YOUTUBE_API_VERSION,
http=credentials.authorize(httplib2.Http()))
def initialize_upload(youtube_api, options):
tags = None
if options.keywords:
tags = options.keywords.split(",")
body = dict(
snippet=dict(
title=options.title,
description=options.description,
tags=tags,
categoryId=options.category
),
status=dict(
privacyStatus=options.privacyStatus
)
)
# Call the API's videos.insert method to create and upload the video.
insert_request = youtube_api.videos().insert(
part=",".join(body.keys()),
body=body,
media_body=MediaFileUpload(options.file, chunksize=-1, resumable=True)
)
video_id = resumable_upload(insert_request)
return video_id
def resumable_upload(insert_request):
response = None
error = None
retry = 0
while response is None:
try:
print("Uploading file...")
status, response = insert_request.next_chunk()
if response is not None:
if 'id' in response:
print("Video id '%s' was successfully uploaded." % response['id'])
return response['id']
else:
exit("The upload failed with an unexpected response: %s" % response)
except HttpError as e:
if e.resp.status in RETRIABLE_STATUS_CODES:
error = "A retriable HTTP error %d occurred:\n%s" % (e.resp.status, e.content)
else:
raise
except RETRIABLE_EXCEPTIONS as e:
error = "A retriable error occurred: %s" % e
if error is not None:
print(error)
retry += 1
if retry > MAX_RETRIES:
exit("No longer attempting to retry.")
max_sleep = 2 ** retry
sleep_seconds = random.random() * max_sleep
print("Sleeping %f seconds and then retrying..." % sleep_seconds)
time.sleep(sleep_seconds)
def upload_thumbnail(youtube, video_id, thumbnail_path):
print("Uploading thumbnail...")
request = youtube.thumbnails().set(
videoId=video_id,
media_body=MediaFileUpload(thumbnail_path)
)
response = request.execute()
print("Thumbnail uploaded successfully.")
if __name__ == '__main__':
argparser.add_argument("--file", required=True, help="Video file to upload")
argparser.add_argument("--title", help="Video title", default="Test Title")
argparser.add_argument("--description", help="Video description", default="Test Description")
argparser.add_argument("--category", default="22",
help="Numeric video category. See https://developers.google.com/youtube/v3/docs/videoCategories/list")
argparser.add_argument("--keywords", help="Video keywords, comma separated", default="")
argparser.add_argument("--privacyStatus", choices=VALID_PRIVACY_STATUSES,
default=VALID_PRIVACY_STATUSES[0], help="Video privacy status.")
argparser.add_argument("--thumbnail", help="Path to thumbnail image file")
args = argparser.parse_args()
if not os.path.exists(args.file):
exit("Please specify a valid file using the --file= parameter.")
youtube = get_authenticated_service(args)
try:
video_id = initialize_upload(youtube, args)
if args.thumbnail:
if not os.path.exists(args.thumbnail):
print("Thumbnail file does not exist. Skipping thumbnail upload.")
else:
upload_thumbnail(youtube, video_id, args.thumbnail)
except HttpError as e:
print("An HTTP error %d occurred:\n%s" % (e.resp.status, e.content))
noise_nodes.py
-95
View File
@@ -1,95 +0,0 @@
import torch
import folder_paths
import torchvision
import torch.nn as nn
import functools
import torch
NODE_CLASS_MAPPINGS = {}
NODE_DISPLAY_NAME_MAPPINGS = {}
def ETK_noise_base(cls):
cls.CATEGORY = "ETK/noise"
# Add spaces to the camel case class name
pretty_name = cls.__name__
for i in range(1, len(pretty_name)):
if pretty_name[i].isupper():
pretty_name = pretty_name[:i] + " " + pretty_name[i:]
cls.DISPLAY_NAME = pretty_name
NODE_CLASS_MAPPINGS[cls.__name__] = cls
NODE_DISPLAY_NAME_MAPPINGS[cls.DISPLAY_NAME] = pretty_name
# Wrap the function defined in the FUNCTION attribute
func_name = getattr(cls, "FUNCTION", None)
if func_name and hasattr(cls, func_name):
original_func = getattr(cls, func_name)
@functools.wraps(original_func)
def wrapped_func(*args, **kwargs):
with torch.inference_mode(False):
return original_func(*args, **kwargs)
setattr(cls, func_name, wrapped_func)
return cls
@ETK_noise_base
class PerlinFractalImage:
"""
use noise.perlin_fractal_image to generate a torch tensor (b,h,w,c=3)
"""
# noise = torch.cat((noise, noise, noise), dim=3)
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"width": ("INT", {"default": 256}),
"height": ("INT", {"default": 256}),
"scale": ("FLOAT", {"default": 0.01}),
"octaves": ("INT", {"default": 6}),
"lacunarity": ("FLOAT", {"default": 3}),
"gain": ("FLOAT", {"default": 0.5}),
"seed": ("INT", {"default": 1}),
"offset": ("FLOAT", {"default": 0}),
"workers": ("INT", {"default": 4}),
"norm_zero_to_1": ([True, False], {"default": True}),
}
}
# Define the return types
RETURN_TYPES = ("IMAGE", "IMAGE", "MASK",)
RETURN_NAMES = ("Image", "Grayscale Image", "Mask",)
FUNCTION = "handler"
# Method to download the dataset
def handler(self, **kwargs):
from .modules import noise
width = kwargs["width"]
height = kwargs["height"]
scale = kwargs["scale"]
octaves = kwargs["octaves"]
lacunarity = kwargs["lacunarity"]
gain = kwargs["gain"]
seed = kwargs["seed"]
offset = kwargs["offset"]
workers = kwargs["workers"]
norm_zero_to_1 = kwargs["norm_zero_to_1"]
noise = noise.get_perlin_2d_as_torch_image_stack(width, height, scale, octaves, lacunarity, gain, seed, offset,
workers, norm_zero_to_1)
noise_image = torch.cat((noise, noise, noise), dim=3)
# format for this according to comfyui is unknown atm
noise_grayscale = noise
# format for this according to comfyui is unknown atm
noise_mask = noise_grayscale > torch.median(noise_grayscale)
noise_mask = noise_mask.float()
return (noise_image, noise_grayscale, noise_mask,)
server_endpoints.py
-687
View File
@@ -1,687 +0,0 @@
from .config import config_settings
from .utils import FileBackedDict
import json
TEST= False
try:
from server import PromptServer
except ImportError:
TEST=True
# for now just exit
if not TEST:
from aiohttp import web
from importlib import reload
import logging
import os
from folder_paths import output_directory, input_directory, temp_directory
# from custom_nodes.EternalKernelLiteGraphNodes.local_shared import GDE_PATH, ETK_PATH
from .local_shared import GDE_PATH, ETK_PATH
from .config import config_settings
fake_git = None
# real_git_url = "https://github.com/story-squad/GDE_Graph_IO.git" # ssh url
# real_git_url = "https://github.com/Hopping-Mad-Games/COTRF.git" # ssh url
class GitIO:
def __init__(self, i_real_git_url: str, i_gde_path: str):
self.gde_path = i_gde_path
self.real_git_io_path = self.setup_git_io(i_real_git_url)
def setup_git_io(self, l_real_git_url: str):
git_name = l_real_git_url.split("/")[-1].split(".")[0]
real_git_path = os.path.join(self.gde_path, "data", "gitio", git_name)
# check to see if it exists and is a git repo, if not then clone it
if os.path.exists(real_git_path):
# check to see if it is a git repo
if os.path.exists(os.path.join(real_git_path, ".git")):
print("git io found")
else:
# not a git repo, delete it and clone it
print("git io error, not a git repo")
else:
# clone it
print("git io not found")
print("cloning git io repo")
os.system("git clone " + l_real_git_url + " " + real_git_path)
return real_git_path
def ensure_git_state(self, user: str):
import subprocess
"""works in windows or linux, ensures that the git repo is in the correct state"""
if not os.path.exists(self.real_git_io_path):
raise Exception("git io path not found")
if not os.path.exists(os.path.join(self.real_git_io_path, ".git")):
raise Exception("git io error, not a git repo")
def run_git_command(command: str):
"""Helper function to run a git command in the correct directory"""
return subprocess.check_output(command, cwd=self.real_git_io_path, text=True).strip()
# fetch first?
run_git_command(["git", "fetch"])
# get the current branch
current_branch = run_git_command(["git", "branch", "--show-current"])
##################
try:
origin_branch_exists = not run_git_command(["git", "rev-parse", "--verify", "--quiet", "origin/" + user])
origin_branch_exists = origin_branch_exists != ""
except subprocess.CalledProcessError:
origin_branch_exists = False
try:
local_branch_exists = run_git_command(["git", "rev-parse", "--verify", "--quiet", user])
local_branch_exists = local_branch_exists != ""
except subprocess.CalledProcessError:
local_branch_exists = False
if not origin_branch_exists and not local_branch_exists:
print(f"git branch {user} does not exist locally or on the remote")
# the branch does not exist locally or on the remote so create it here first
# checkout new_user
run_git_command(["git", "checkout", "new_user"])
# create new branch
run_git_command(["git", "branch", user]) # creates and checks out new branch
# git checkout feature-branch
run_git_command(["git", "checkout", user])
# set the upstream to the remote
# run_git_command(["git", "branch", "--set-upstream-to", "origin/" + user])
# push the branch to the remote
run_git_command(["git", "push", "-u", "origin", user])
# fetch and pull it back down for good measure
run_git_command(["git", "fetch"])
run_git_command(["git", "pull"])
elif not origin_branch_exists and local_branch_exists:
print(f"git branch {user} exists locally but not on the remote")
# the branch does not exist on the remote but does locally
# set the upstream
# run_git_command(["git", "branch", "--set-upstream-to", "origin/" + user])
# push the branch to the remote
run_git_command(["git", "push", "-u", "origin", user])
# fetch and pull it back down for good measure
run_git_command(["git", "fetch"])
run_git_command(["git", "pull"])
elif origin_branch_exists and not local_branch_exists:
print(f"git branch {user} exists on the remote but not locally")
# the branch does not exist locally but does on the remote
# fetch and pull it down
run_git_command(["git", "checkout", user])
# pull
run_git_command(["git", "pull"])
elif origin_branch_exists and local_branch_exists:
print(f"git branch {user} exists locally and on the remote")
# the branch exists locally and on the remote
# pull it down
run_git_command(["git", "checkout", user])
# set the upstream to the remote
run_git_command(["git", "branch", "--set-upstream-to", "origin/" + user])
current_branch = run_git_command(["git", "branch", "--show-current"])
# if the current branch is not the user's branch, then switch to it
if current_branch != user:
print(f"git branch {user} is not the current branch, switching to it")
# also check if the branch exists, if not then create it
branches = run_git_command(["git", "branch"])
if user not in branches:
run_git_command(["git", "branch", user])
# be verbose about what exactly is happening where
print(f"git branch {user} does not exist locally, creating it, and pushing it to the remote")
run_git_command(["git", "push", "--set-upstream", "origin", user])
# switch to the user's branch
run_git_command(["git", "checkout", user])
# pull the latest from the user's branch
run_git_command(["git", "pull"])
def git_save_version_1(self, user: str, file_name: str, file_data: str):
"""
for git save version 1 each user will have their own branch,
and the data will be saved to a file in the user's branch
version > 1 may include other strategies
"""
self.ensure_git_state(user)
# change to correct directory
os.chdir(self.real_git_io_path)
# switch to the users branch
os.system("git checkout " + user)
# save the file data as the file to the disk
with open(os.path.join(self.real_git_io_path, file_name), "w", encoding="utf-8") as f:
f.write(file_data)
# add the file to the git repo
os.system("git add " + file_name)
# commit the file to the git repo
os.system("git commit -m \"saving file " + file_name + "\"")
# push the file to the git repo
os.system("git push")
def git_load_version_1(self, user_name, file_name):
"""
for git load version 1 each user will have their own branch,
and the data will be loaded from a file in the user's branch
version > 1 may include other strategies
"""
self.ensure_git_state(user_name)
# change to correct directory
os.chdir(self.real_git_io_path)
# switch to the users branch
os.system("git checkout " + user_name)
# pull the latest from the user's branch
os.system("git pull")
# read the file from the disk
with open(os.path.join(self.real_git_io_path, file_name), "r", encoding="utf-8") as f:
file_data = f.read()
return file_data
git_io = None
user_data_base = {
"tasha_StSq": {
"password": "[REDACTED]",
"security": None,
"client_id": None,
"git_io": "StSq"
},
"tasha_COTRF": {
"password": "[REDACTED]",
"security": None,
"client_id": None,
"git_io": "COTRF"
},
"HMG": {
"password": "[REDACTED]",
"security": None,
"client_id": None,
"git_io": "HMG"
},
"API": {
"password": "[REDACTED]",
"security": None,
"client_id": None,
"git_io": "HMG"
},
}
tmp_dir = config_settings["tmp_dir"]
user_data_file = os.path.join(tmp_dir, 'user_data_tmp.json')
user_data = FileBackedDict(user_data_file)
if len(user_data.keys()) == 0:
user_data.update(user_data_base)
def get_user_by_client_id(client_id):
for user in user_data:
if user_data[user]["client_id"] == client_id:
return user
return None
async def get_user_if_client_id_logged_in(client_id):
try:
user = get_user_by_client_id(client_id)
except KeyError:
user = None
return user
def get_user_by_security(security):
for user in user_data:
if user_data[user]["security"] == security:
return user
return None
@PromptServer.instance.routes.post("/gde/logout")
async def logout(request):
"""
This function is called when the user logs out
"""
print("logout request")
# get the user and pass from the json
data = await request.json()
client_id = data["client_id"]
# find the user with the client id
user = get_user_by_client_id(client_id)
if user is None:
# use json to send the data back
return web.json_response({"error": "invalid client_id"}, status=401)
# clear the security and client id
user_data[user]["security"] = None
user_data[user]["client_id"] = None
global git_io
git_io = None
# use json to send the data back
return web.json_response({"success": True})
@PromptServer.instance.routes.post("/gde/login")
async def login(request):
"""
This function is called when the attemnpts to log in:
- validates login
- sets up security
- user and pass are in the json
"""
print("login request")
# get the user and pass from the json
data = await request.json()
try:
client_id = data["client_id"]
has_client_id = True
except:
# this means the client id is not set yet so return logged_in False
# use json to send the data back
has_client_id = False
return web.json_response({"logged_in": False})
try:
user = data["username"]
password = data["password"]
has_user_and_pass = True
except:
has_user_and_pass = False
# this means the client is just checking to see if it is already logged into the server
# so check
user = get_user_by_client_id(client_id)
if user is None:
print("user is none")
return web.json_response({"logged_in": False}, status=401)
else:
# send the client back its user so it can handle api calls correctly
print("user recconected " + user)
return web.json_response({"logged_in": True, "user": user}, status=200)
# validate the user and pass
# for now the users and passes are right here in the code
if user not in user_data:
# use json to send the data back
return web.json_response({"error": "invalid user"}, status=401)
if password != user_data[user]["password"]:
# use json to send the data back
return web.json_response({"error": "invalid password"}, status=401)
# if the user and pass are valid, then set up security
# for now the security is just a random string that will be sent back to the user
# in the future this will be a token that will be used to authenticate the user
import uuid
security = str(uuid.uuid4())
# there is no middleware installed
# so we will just store the security in memory
# this is not secure, but it is good enough for now
# store the security in memory
user_data[user]["security"] = security
user_data[user]["client_id"] = client_id
global git_io
# from . import config
# config_settings = config.config_settings
users_gitio = user_data[user]["git_io"]
# users_gitio_path = os.path.abspath(os.path.join(
# GDE_PATH,
# config_settings["gitio"]["base"],
# users_gitio)
# )
try:
users_gittio_url = config_settings["gitio"][users_gitio]
except KeyError:
return web.json_response({"error": "invalid gitio"}, status=401)
git_io = GitIO(users_gittio_url, GDE_PATH)
# return json with the security
return web.json_response({"security": security})
def hijack_prompt_server(route_path,
before_get_func: callable = None,
after_get_func: callable = None,
before_post_func: callable = None,
after_post_func: callable = None,
):
"""
This function is called by the GDE when it starts up, it hijacks the prompt server
for added security.
"""
from server import PromptServer
# find the existing /prompt route
rts = [i for i, s in enumerate(PromptServer.instance.routes) if (route_path == s.path)]
print(f"found {len(rts)} routes to remove: {rts}")
# remove the existing prompt route
for i, rt_num in enumerate(rts):
rt_to_pop = rt_num - i
route = PromptServer.instance.routes._items.pop(rt_to_pop)
# route = PromptServer.instance.routes._items[rt_to_pop]
# myr = web.RouteTableDef()
method = route.method
path = route.path
kwargs = route.kwargs
if method == "GET":
print(f"Hijacking route {method} {path} {kwargs}")
get_handler = route.handler
if not before_get_func:
if not after_get_func:
raise NotImplementedError
if before_get_func and after_get_func:
raise NotImplementedError
if before_get_func:
@PromptServer.instance.routes.get(path)
async def get_prompt(request):
part1 = await before_get_func(request)
if isinstance(part1, web.Response):
return part1
if isinstance(part1, web.Request):
modified_func_response = await get_handler(part1)
return modified_func_response
elif after_get_func:
@PromptServer.instance.routes.get(path)
async def get_prompt(request):
normal_handler_response = await get_handler(request)
modified_func_response = await after_get_func(normal_handler_response)
return modified_func_response
elif method == "POST":
print(f"Hijacking route {method} {path} {kwargs}")
post_handler = route.handler
if not before_post_func:
if not after_post_func:
raise NotImplementedError
if before_post_func and after_post_func:
raise NotImplementedError
if before_post_func:
@PromptServer.instance.routes.post(path)
async def post_prompt(request):
part1 = await before_post_func(request)
if isinstance(part1, web.Response):
return part1
if isinstance(part1, web.Request):
# modified_func_response = await post_handler(part1)
# if it has a json() method then it is a request, careful it might not have the attribute at all
if hasattr(part1, "json"):
modified_func_response = await post_handler(part1)
if isinstance(part1, dict):
a = await post_handler(part1)
b = part1
return a
return modified_func_response
elif after_post_func:
@PromptServer.instance.routes.post(path)
async def post_prompt(request):
normal_handler_response = await post_handler(request)
modified_func_response = await after_post_func(normal_handler_response)
return modified_func_response
from aiohttp import web
@PromptServer.instance.routes.post("/gde/git/user_graphs")
async def get_user_graphs(request):
"""
Returns a list of available user graphs (dropdown options).
"""
import os
import subprocess
if git_io == None:
return web.json_response({"error": "git_io not set"}, status=401)
def run_git_command(command: list, dir_path: str):
"""Helper function to run a git command in the specified directory"""
return subprocess.call(command, cwd=dir_path)
data = await request.json()
client_id = data.get("client_id", None)
user = data.get("user", None)
if user not in user_data:
user = get_user_by_client_id(client_id)
if await get_user_if_client_id_logged_in(client_id) is None:
return web.json_response({"error": "user not found"}, status=401)
# user = data.get("user", None)
print("getting user graphs for user:", user)
# Directory where the user graphs are stored
user_graphs_dir = git_io.real_git_io_path
# Ensure that the git repo is in the correct state
git_io.ensure_git_state(user)
run_git_command(["git", "pull", "origin", user], user_graphs_dir)
import os
user_graph_files = [
f for f in os.listdir(user_graphs_dir)
if os.path.isfile(os.path.join(user_graphs_dir, f)) and f.endswith(".json")
]
# Sort files by modification time, newest first
user_graph_files.sort(
key=lambda x: os.path.getmtime(os.path.join(user_graphs_dir, x)),
reverse=True
)
options_data = [{"value": f, "label": f} for f in user_graph_files if f != ""]
# options_data = []
return web.json_response(options_data)
@PromptServer.instance.routes.post("/gde/git/users")
async def get_users(request):
""" the list of users is just the list of branches in the git repo"""
import subprocess
print("getting users")
# Ensure that the git repo is in the correct state
git_io.ensure_git_state("main")
# Change to correct directory
os.chdir(git_io.real_git_io_path)
# Fetch updates from the remote
subprocess.call(["git", "fetch", "--prune"])
# List all remote branches
remote_branches_raw = subprocess.check_output(["git", "branch", "-r"]).decode("utf-8").split("\n")
remote_branches = [b.replace("origin/", "").strip() for b in remote_branches_raw]
# List all local branches
local_branches_raw = subprocess.check_output(["git", "branch"]).decode("utf-8").split("\n")
local_branches = [b.replace("*", "").strip() for b in local_branches_raw]
# Identify local branches not present in the remote branches
branches_to_delete = [b for b in local_branches if b not in remote_branches]
# Delete these local branches
for branch in branches_to_delete:
subprocess.call(["git", "branch", "-D", branch])
# Now, get the updated list of local branches
branches = subprocess.check_output(["git", "branch"]).decode("utf-8").split("\n")
branches = [b.replace("*", "").strip() for b in branches]
# Convert the list of branches to the desired format
options_data = [{"value": b, "label": b} for b in branches if b != ""]
return web.json_response(options_data)
@PromptServer.instance.routes.post("/gde/git/save")
async def save_to_git(request):
"""
for git save version 1 each user will have their own branch,
and the data will be saved to a file in the user's branch
version > 1 may include other strategies
"""
git_save_version = 1
req_json = await request.json()
print("saving to git")
req_json = await request.json()
data = req_json.get("data", None)
user_name = req_json.get("user_name", None)
file_name = req_json.get("file_name", None)
git_io.git_save_version_1(user_name, file_name, data)
# Implement your actual Git saving logic here
return web.Response(text="Saved to Git successfully.")
@PromptServer.instance.routes.post("/gde/git/load")
async def load_from_git(request):
"""use git_io to load the file from git"""
print("loading from git")
req_json = await request.json()
user_name = req_json.get("user_name", None)
file_name = req_json.get("file_name", None)
# Implement your actual Git loading logic here
file_data = git_io.git_load_version_1(user_name, file_name)
return web.json_response({"data": file_data})
# print(f"Hijacking route {method} {path} {kwargs}")
# post_handler = route.handler
#
async def my_post_hijack_before(request: web.Request):
json_data = await request.json()
client_id = json_data.get("client_id", None)
ret = request
print("client_id:", client_id)
if not isinstance(request, web.Request):
print("request is not a web.Request")
return request
user = await get_user_if_client_id_logged_in(client_id)
skp_lcl = config_settings["skip_login_for_local"]
is_local_request = request.transport.get_extra_info('peername')[0] in (
'127.0.0.1', '::1') if request.transport.get_extra_info('peername') else False
if user:
print("user:", user)
ret = request
elif skp_lcl and is_local_request:
ret = request
else:
ret = web.json_response({"error": "user not found"}, status=401)
return ret
async def my_get_hijack_after(response):
return response
# ------------------------------------
# WRAPPER: so we can override json()
# ------------------------------------
class ModifiedRequestWrapper:
"""
A wrapper around the original `web.Request` that overrides
the `json()` method to return our custom dictionary.
"""
def __init__(self, original_request: web.Request, new_json: dict):
self._original_request = original_request
self._new_json = new_json
def __getattr__(self, name):
# Delegate everything except .json() to the original request
return getattr(self._original_request, name)
async def json(self):
# Return our overridden data
return self._new_json
def fill_prompt(raw_json,**kwargs):
for node_num,node in raw_json.items():
inputs = node["inputs"]
if node["class_type"]=="RequestInput":
key = inputs["key"]
if key in kwargs:
raw_json[node_num]["inputs"]["overridden_value"] = kwargs[key]
return raw_json
# ------------------------------------------------------------------
# This function checks for "api_file" and if found, injects its JSON
# into the request by returning a ModifiedRequestWrapper.
# ------------------------------------------------------------------
async def my_post_hijack_before_api_file(request: web.Request):
import os
import json
# Convert the request body into a dict
data = await request.json()
api_file = data.get("api_file", None)
if api_file:
kwargs = data.get("kwargs", {})
# Build path to the file in ./api_graphs
current_dir = os.path.dirname(os.path.abspath(__file__))
api_graphs_dir = os.path.join(current_dir, "api_graphs")
file_path = os.path.join(api_graphs_dir, api_file)
if not os.path.isfile(file_path):
return web.json_response({"error": f"api_file '{api_file}' not found"}, status=400)
with open(file_path, "r", encoding="utf-8") as f:
loaded_json = json.load(f)
filled_json = fill_prompt(loaded_json,**kwargs)
# Insert loaded_json into the "prompt" key
data["prompt"] = filled_json
# Return a new request-like object that yields `data` from .json()
return ModifiedRequestWrapper(request, data)
# If "api_file" isn't there, just return the original request
return request
if __name__ == 'EternalKernelLiteGraphNodes.server_endpoints':
async def both_post_hijack_before(request: web.Request):
thing1 = await my_post_hijack_before(request)
thing2 = await my_post_hijack_before_api_file(thing1)
return thing2
hijack_prompt_server("/prompt", after_get_func=my_get_hijack_after, before_post_func=both_post_hijack_before)
pass
tests/test_descript.py
-16
View File
@@ -1,16 +0,0 @@
from unittest import TestCase
class TestDecodeTranscriptNode(TestCase):
def test_func(self):
from custom_nodes.EternalKernelLiteGraphNodes.descript import DecodeTranscriptNode
node = DecodeTranscriptNode()
string = """[00:00:00] jackie: Let's talk about a
[00:00:01] jackie: thing\n\n[00:00:05] jackie: Let's talk about an apple
"""
speakers, durations, paragraphs = node.func(string)
self.assertEqual(speakers, ["jackie", "jackie", "jackie"])
self.assertEqual(durations, [1, 4, 2])
self.assertEqual(paragraphs, ["Let's talk about a", 'thing', "Let's talk about an apple"])
tests/test_git.py
-120
View File
@@ -1,120 +0,0 @@
import unittest
from unittest.mock import patch, Mock
from git import Repo
from git.exc import GitCommandError
import os
os.putenv('UNIT_TEST', 'True')
from ..git import PublicGitRepo, CloneRepoNode, PullRepoNode, PushRepoNode, SendCommitNode, CreateRepoNode
repo_path = './tmp/test_repo'
commit_message = 'message'
token = 'token'
# use tashaskyups's repo on github for testing
repo_url = "https://github.com/TashaSkyUp/public_repo.git"
class BaseTestCase(unittest.TestCase):
pass
class TestPublicGitRepo(BaseTestCase):
def test_validate_git_repo_url(self):
repo = PublicGitRepo()
# Test with a valid URL
result = repo.validate_git_repo_url(repo_url=repo_url, token=token)
self.assertTrue(result)
# Test with an invalid URL
result = repo.validate_git_repo_url(repo_url='https://github.com', token=token)
self.assertFalse(result)
class TestCloneRepoNode(BaseTestCase):
# @patch.object(Repo, 'clone_from')
def test_clone_repo(self):
import shutil
# delete the repo if it exists first
if os.path.exists(repo_path):
# this fails so
# shutil.rmtree(repo_path)
# use this instead
os.system(f'rmdir "{repo_path}" /S /Q ')
node = CloneRepoNode()
node.clone_repo(repo_url=repo_url, path=repo_path, token=token)
# shutil.rmtree(repo_path)
# os.system(f'rmdir "{repo_path}" /S /Q ')
# mock_clone_from.assert_called_once_with(repo_url, repo_path,
# env={'GIT_ASKPASS': 'echo', 'GIT_USERNAME': 'token'})
class TestPullRepoNode(BaseTestCase):
# @patch.object(Repo, 'remotes')
# @patch.object(Repo, '__init__', return_value=None)
def test_pull_repo(self):
# mock_origin = Mock()
# mock_remotes.origin = mock_origin
node = PullRepoNode()
node.pull_repo(repo_path=repo_path, token=token)
# mock_origin.pull.assert_called_once()
class TestSendCommitNode(BaseTestCase):
# @patch.object(Repo, 'git')
# @patch.object(Repo, 'index')
# @patch.object(Repo, '__init__', return_value=None)
def test_send_commit(self):
import hashlib
data = hashlib.md5().hexdigest()
# write to a file
with open(f'{repo_path}/test.txt', 'w') as f:
f.write(data)
# now commit the file
node = SendCommitNode()
node.send_commit(repo_path=repo_path, commit_message=commit_message, token=token)
# mock_git.add.assert_called_once_with(A=True)
# mock_index.commit.assert_called_once_with(commit_message)
class TestPushRepoNode(BaseTestCase):
# @patch.object(Repo, 'remotes')
# @patch.object(Repo, '__init__', return_value=None)
def test_push_repo(self):
# mock_origin = Mock()
# mock_remotes.origin = mock_origin
node = PushRepoNode()
node.push_repo(repo_path=repo_path, token=token)
# mock_origin.push.assert_called_once()
class TestCreateRepoNode(BaseTestCase):
@patch.object(Repo, 'init')
def test_create_repo(self, mock_init):
node = CreateRepoNode()
node.create_repo(path=repo_path, token=token)
mock_init.assert_called_once_with(repo_path)
if __name__ == '__main__':
# run the tests one at a time in a n order that makes sense
runner = unittest.TextTestRunner()
suite = unittest.TestSuite()
suite.addTest(TestCloneRepoNode('test_clone_repo'))
runner.run(suite)
suite = unittest.TestSuite()
suite.addTest(TestSendCommitNode('test_send_commit'))
runner.run(suite)
suite = unittest.TestSuite()
suite.addTest(TestPushRepoNode('test_push_repo'))
runner.run(suite)
# suite.addTest(TestPublicGitRepo('test_validate_git_repo_url'))
# suite.addTest(TestCreateRepoNode('test_create_repo'))
# suite.addTest(TestPullRepoNode('test_pull_repo'))
# run all the tests
# unittest.main()
tests/test_hf_diffusers_nodes.py
-84
View File
@@ -1,84 +0,0 @@
import pytest
from ..hf_diffusers_nodes import DDPMPipline, SampleDDPMPipline, TensorToImage
from ..hf_diffusers_nodes import LoadHFDataset
# when testing set "UNIT_TEST" to true in the environment to avoid loading all of ComfyUI
def test_DDPMPipline():
p_node = DDPMPipline()
ret = p_node.load_from_pretrained("johnowhitaker/ddpm-butterflies-32px",
"cuda")
assert ret is not None
def test_SampleDDPMPipline():
p_node = DDPMPipline()
s_node = SampleDDPMPipline()
ret = p_node.load_from_pretrained("johnowhitaker/ddpm-butterflies-32px",
"cuda")
ret = ret[0]
assert ret is not None
ret = s_node.sample(ret,
1,
25
)
assert ret is not None
def test_TensorToImage():
p_node = DDPMPipline()
s_node = SampleDDPMPipline()
t_node = TensorToImage()
ret = p_node.load_from_pretrained("johnowhitaker/ddpm-butterflies-32px",
"cuda")
ret = ret[0]
assert ret is not None
ret = s_node.sample(ret,
1,
25
)
ret = ret[0]
assert ret is not None
ret = t_node.tensor_to_image(ret, "HWC")
ret = ret[0]
assert ret is not None
# now convert to pillow image
from PIL import Image
import numpy as np
# ret is float32 0-1 torch tensor we need to convert it to numpy and then to a pillow image
ret = ret.squeeze(0)
ret = ret.squeeze(0)
ret = ret.numpy()
ret = ret * 255
ret = np.uint8(ret)
ret = Image.fromarray(ret)
ret.show()
def test_LoadHFDataset():
p_node = LoadHFDataset()
ret = p_node.load_from_hf("huggan/smithsonian_butterflies_subset")
assert ret is not None
if __name__ == "__main__":
pytest.main()
def test_create_unet2d_model():
from ..hf_diffusers_nodes import CreateUnet2dModel
p_node = CreateUnet2dModel()
ret = p_node.create_unet2d_model(3,
3,
1,
1,
1,
1,
)
assert ret is not None
tests/test_image.py
-48
View File
@@ -1,48 +0,0 @@
import unittest
from PIL import Image
from custom_nodes.EternalKernelLiteGraphNodes.image import TextRender as Renderer
import os
hard_test = """Chuck Norris needs no mouse;
He stares down the computer
THE END.
SchwaaWeb, trained by Chuck!"""
hard_test2 ="""Chuck Norris can catch a data breach before it even happens.
SchwaaWeb, vigilant
like Chuck!"""
# set an env var to tell so we know we are in test mode
os.environ["ETERNAL_KERNEL_LITEGRAPH_NODES_TEST"] = "True"
class TestRenderText(unittest.TestCase):
def test_render_text(self):
# Initialize object that includes _render_text and _wrap_text methods
renderer = Renderer()
# Create a blank image
image = Image.new('RGB', (500, 500), color=(73, 109, 137))
# Define parameters for _render_text
font_name = 'arial'
size = 24
text = hard_test2
x, y = 0, 0
try:
# Call the function under test
renderer._render_text(font_name, image, size, text, x, y, allow_wrap=True, allow_shrink=True)
# Check that the image is not empty (i.e., it has been changed by the drawing operation)
self.assertNotEqual(image.tobytes(), Image.new('RGB', (500, 500), color=(73, 109, 137)).tobytes())
finally:
# Show the image whether the test passed or failed
image.show()
if __name__ == '__main__':
unittest.main()
if __name__ == '__main__':
unittest.main()
tests/test_torchvision_nodes.py
-62
View File
@@ -1,62 +0,0 @@
import pytest
from ..torchvision_nodes import (
TorchVisionTransformCompositionList,
TorchVisionTransformStringParser,
TorchVisionTransformNode,
TorchVisionCallComposition,
)
# when testing set "UNIT_TEST" to true in the environment to avoid loading all of ComfyUI
def test_TorchVisionTransformCompositionList():
tvtcl_node = TorchVisionTransformCompositionList()
ret = tvtcl_node.compose("test")
ret = ret[0]
assert ret is not None
def test_TorchVisionTransformStringParser():
node = TorchVisionTransformStringParser()
ret = node.compose("Resize(64,64)\nGaussianBlur(1)", "test")
transforms, name = ret
assert name == "test"
assert transforms[0][0] == "Resize"
assert transforms[0][1] == ["64", "64"]
assert transforms[1][0] == "GaussianBlur"
assert transforms[1][1] == ["1"]
def test_TorchVisionTransformNode():
node1 = TorchVisionTransformCompositionList()
node2 = TorchVisionTransformNode()
ret = node1.compose("test")
ret = ret[0]
assert ret is not None
ret = node2.compose(ret,"GaussianBlur","1")
ret = ret[0]
assert ret is not None
def test_TorchVisionCallComposition():
node1 = TorchVisionTransformCompositionList()
node2 = TorchVisionTransformNode()
node3 = TorchVisionCallComposition()
ret = node1.compose("test")
ret = ret[0]
assert ret is not None
ret = node2.compose(ret,"GaussianBlur","1")
ret = ret[0]
assert ret is not None
import torch
test_tensor = torch.rand(1,256,256,3)
ret = node3.compose(ret,test_tensor)
ret = ret[0]
assert ret is not None
if __name__ == "__main__":
pytest.main()
tests/test_utils.py
-41
View File
@@ -1,41 +0,0 @@
import torch
import numpy as np
from custom_nodes.EternalKernelLiteGraphNodes.utils import etk_deep_copy
def test_etk_deep_copy():
original = {
"list": [1, 2, 3, [4, 5, 6], {"inner_list": [7, 8, 9]}],
"dict": {"inner_dict": {"key": "value"}}
}
torch_test = {"torch_tensor": torch.tensor([1, 2, 3])}
np_test = {"np_array": np.array([1, 2, 3])}
test(original)
test(torch_test)
test(np_test)
def test(original):
# Creating a deep copy
copy = etk_deep_copy(original)
# Modifying the original structure
if "list" in original:
original["list"][0] = 100
if isinstance(original["list"][3], list):
original["list"][3][0] = 200
if "inner_list" in original["list"][4]:
original["list"][4]["inner_list"][0] = 300
if "dict" in original and "inner_dict" in original["dict"]:
original["dict"]["inner_dict"]["key"] = "modified"
if "torch_tensor" in original:
original["torch_tensor"][0] = 100
if "np_array" in original:
original["np_array"][0] = 200
# Returning both the original and the copy for inspection
return original, copy
tests/test_video.py
-277
View File
@@ -1,277 +0,0 @@
from unittest import TestCase
import os
import torch
import os
import shutil
import tempfile
import pytest
import torch
from PIL import Image
from ..audio_video_folder import AddFramesToVideoFolder, VideoFolderProvider, VideoDefinitionProvider
os.environ["ETERNAL_KERNEL_LITEGRAPH_NODES_TEST"] = "True"
class TestImageStackToVideoFile(TestCase):
def test_handler(self):
# test with random frames of noise
import numpy as np
from custom_nodes.EternalKernelLiteGraphNodes.video import ImageStackToVideoFile
# create a random stack of images
stack = np.random.randint(0, 255, (100, 100, 100, 3), dtype=np.uint8)
# it expects a tensor
stack = torch.from_numpy(stack)
# create the node
node = ImageStackToVideoFile()
# run the node
node.handler(image_stack=stack, video_out="test.mp4")
# check that the file exists
# find the actual filepath for the generated file, will be ../../../video/test.mp4
test_video_file_fullpath = os.path.join(
os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))), "video",
"test.mp4")
self.assertTrue(os.path.exists(test_video_file_fullpath))
# remove the file
os.remove(test_video_file_fullpath)
class TestTemporalSpatialSmoothing(TestCase):
def test_handler(self):
# test with random frames of noise
from custom_nodes.EternalKernelLiteGraphNodes.video import TemporalSpatialSmoothing, ImageStackToVideoFile
from custom_nodes.EternalKernelLiteGraphNodes.image import torch_image_show
# use only torch
# create a random stack of images
# use floats
stack = torch.rand((48, 512, 512, 3))
# _,path= ImageStackToVideoFile().handler(image_stack=stack, video_out="test_start.mp4",fps=24)
# launch the default video player
# os.system(f"start {path}")
# create the node
node = TemporalSpatialSmoothing()
# run the node
result = node.handler(image_stack=stack, spatial_kernel_size=3, temporal_kernel_size=3, spatial_sigma=.5,
temporal_sigma=.5)
# check that the file exists
_, path = ImageStackToVideoFile().handler(image_stack=result[0], video_out="test_result.mp4", fps=24)
# launch the default video player
os.system(f"start {path}")
# torch_image_show(result[0])
class TestTemporalOpticalFlowSmoothing(TestCase):
def test_handler(self):
from custom_nodes.EternalKernelLiteGraphNodes.video import TemporalOpticalFlowSmoothing, ImageStackToVideoFile
import time
import numpy as np
# create a random stack of images
# stack = torch.rand((24*4, 512, 512, 3))
h = 256
w = 512
fps = 24
seconds = 2
stack = torch.rand((fps * seconds, h, w, 3))
# now add a box that moves from left to right with a constant speed
# the box will be 20% of the width and 20% of the height
# the box will move from left to right
# create the box
box = torch.zeros((fps * seconds, int(h), int(w), 3))
# fill the box with ones
box = box + 1
# the box width
bw = min(int(w * .2), int(h * .2))
# the box height
bh = int(h * .2)
# the box x position
t = stack.shape[0]
for i in range(t):
pt = i / t
# the box x position
bx = int((w - bw) * pt)
# now add a jitter
bx += int(np.random.randint(-10, 10, 1))
# the box y position
by = int(h * .333)
# fill the box with ones
box[i, :, :, :] = 0
box[i, by:by + bh, bx:bx + bw, :] = 1
# create a mask for the box region
# box_mask = torch.zeros_like(stack)
# box_mask[:, by:by + bh, bx:bx + bw, :] = 1
# set the color of the stack to blue where the box mask is 1
stack = torch.where(box == 1, torch.tensor([.24, .5, .7], dtype=torch.float32), stack)
nm = "box1"
_, path = ImageStackToVideoFile().handler(image_stack=stack, video_out=nm, fps=fps)
assert os.path.exists(path), f"Output video file doesn't exist: {path}"
os.system(f"start {path}")
time.sleep(1)
# create the node
node = TemporalOpticalFlowSmoothing()
# run the node
result = node.handler(image_stack=stack, device="cuda")
result = result[0]
# result = torch.cat(result[0], dim=0)
# print(result.shape)
# now we can permute the dimensions
# result = result.permute(0, 2,3, 1)
_, path = ImageStackToVideoFile().handler(image_stack=result, video_out="test_result_optical_flow.mp4", fps=fps)
assert os.path.exists(path), f"Output video file doesn't exist: {path}"
os.system(f"start {path}")
## Check the shape of the output
# assert len(result) == stack.shape[0] - 1
# Each element in the result should have the shape of [1, 2, H, W]
# for flow in result:
# assert flow.shape == torch.Size([1, 2, 512, 512])
# check the file generation and video play
from ..audio_video_folder import AddFramesToVideoFolder, VideoFolderProvider, VideoDefinitionProvider
@pytest.fixture(scope="function")
def video_folder_def():
this_file_dir = os.path.dirname(os.path.abspath(__file__))
temp_dir = os.path.join(this_file_dir, "temp")
os.makedirs(temp_dir, exist_ok=True)
vfp = VideoFolderProvider().video_folder_provider(folder_path=temp_dir)[0]
d = {
"video_folder": vfp,
"fps": 24,
"width": 256,
"height": 256,
"frame format": "jpg",
}
VFD = VideoDefinitionProvider().video_definition_provider(**d)[0]
yield VFD
# shutil.rmtree(temp_dir)
def test_add_frames_to_video_folder(video_folder_def):
# Generate synthetic inputs
num_frames = 10
image_size = (256, 256)
frames = torch.randint(0, 256, (num_frames, *image_size, 3), dtype=torch.uint8)
rescale_method = "LANCZOS"
release_vram = True
# Call the function
result = AddFramesToVideoFolder().add_frames_to_video_folder(video_folder_def=video_folder_def,
IMAGE=frames,
Rescale_method=rescale_method,
Release_vram=release_vram
)[0]
# Check the output
assert len(os.listdir(result["video_folder"])) == num_frames
# Check the saved frames
for i in range(num_frames):
frame_path = video_folder_def["get_last_file_name"](i + 1)
saved_frame = Image.open(frame_path)
assert saved_frame.size == image_size
def test_save_or_copy_frame(video_folder_def):
# Test case 1: Different frames
# Generate synthetic inputs
current_frame = torch.randint(0, 256, (256, 256, 3), dtype=torch.uint8)
previous_frame = torch.randint(0, 256, (256, 256, 3), dtype=torch.uint8)
frame_fp = video_folder_def.get_next_file_name()
previous_frame_fp = video_folder_def.get_next_file_name()
# Save the previous frame
Image.fromarray(previous_frame.permute(2, 0, 1).permute(1, 2, 0).numpy()).save(previous_frame_fp)
# Call the function
result = AddFramesToVideoFolder.save_or_copy_frame(current_frame, previous_frame, frame_fp, previous_frame_fp)
# Check the output
assert torch.equal(result, current_frame)
assert os.path.exists(frame_fp)
# Check the saved frame
saved_frame = Image.open(frame_fp)
assert saved_frame.size == (256, 256)
saved_frame.close()
# clean up
os.remove(frame_fp)
# Test case 2: Duplicate frames
# Generate synthetic inputs
duplicate_frames = [torch.randint(0, 256, (256, 256, 3), dtype=torch.uint8)]
duplicate_frame = duplicate_frames * 2
previous_duplicate_frame_fp = video_folder_def.get_next_file_name(0)
# need to check the file names for 0000000000 and 0000000001
assert "0000000000" in previous_duplicate_frame_fp
#assert "0000000001" in duplicate_frame_fp
# Save the previous duplicate frame
Image.fromarray(duplicate_frames[0].permute(2, 0, 1).permute(1, 2, 0).numpy()).save(previous_duplicate_frame_fp)
# Call the function 100 times
for i in range(100):
duplicate_frame_fp = video_folder_def.get_next_file_name(0)
result = AddFramesToVideoFolder.save_or_copy_frame(duplicate_frames[0], duplicate_frame[1], duplicate_frame_fp,
previous_duplicate_frame_fp)
# Check the output
assert torch.equal(result, duplicate_frame[0])
assert os.path.exists(duplicate_frame_fp)
# check that here are 101 files
assert len(os.listdir(video_folder_def["video_folder"])) == 101
# clean up
for i in range(101):
os.remove(video_folder_def.get_last_file_name(0))
def test_process_frames(video_folder_def):
# Generate synthetic inputs
num_frames = 10
image_size = (256, 256)
frames = torch.randint(0, 256, (num_frames, *image_size, 3), dtype=torch.uint8)
rescale_method = "LANCZOS"
# Call the function
result = AddFramesToVideoFolder.process_frames(frames, rescale_method, video_folder_def)
# Check the output
# import pdb;
# pdb.set_trace()
assert result is None
test = video_folder_def["video_folder"]
print("path: ", test)
assert len(os.listdir(test)) == num_frames
# Check the saved frames
assert video_folder_def["get_frame_count"]() == num_frames
for i in range(num_frames):
frame_path = video_folder_def["get_last_file_name"](i + 1)
saved_frame = Image.open(frame_path)
assert saved_frame.size == image_size
tests/test_youtube_nodes.py
-17
View File
@@ -1,17 +0,0 @@
import os
os.environ["UNIT_TEST"] = "True"
from custom_nodes.EternalKernelLiteGraphNodes.youtube_nodes import YoutubeUploadNode
def test_upload():
node = YoutubeUploadNode()
output = node.upload_video(
file_path=os.path.join("tests", "example.mp4"),
category=22,
title="Test Video",
description="This is a test video for the Youtube Upload Node",
keywords="test, youtube, node",
privacy_status="unlisted"
)
print (output)
torchvision_nodes.py
-235
View File
@@ -1,235 +0,0 @@
try:
from .config import config_settings
except ImportError as e:
from config import config_settings
import functools
import torch
import torchvision
#NODE_CLASS_MAPPINGS = {}
#NODE_DISPLAY_NAME_MAPPINGS = {}
from . import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
def ETK_torchvision_base(cls):
cls.CATEGORY = "ETK/torchvision"
# Add spaces to the camel case class name
pretty_name = cls.__name__
for i in range(1, len(pretty_name)):
if pretty_name[i].isupper():
pretty_name = pretty_name[:i] + " " + pretty_name[i:]
cls.DISPLAY_NAME = pretty_name
NODE_CLASS_MAPPINGS[cls.__name__] = cls
NODE_DISPLAY_NAME_MAPPINGS[cls.DISPLAY_NAME] = pretty_name
# Wrap the function defined in the FUNCTION attribute
func_name = getattr(cls, "FUNCTION", None)
if func_name and hasattr(cls, func_name):
original_func = getattr(cls, func_name)
@functools.wraps(original_func)
def wrapped_func(*args, **kwargs):
with torch.inference_mode(False):
return original_func(*args, **kwargs)
setattr(cls, func_name, wrapped_func)
return cls
@ETK_torchvision_base
class TorchVisionTransformCompositionList:
"""Provides the base list that helps define a composition pipeline for torchvision transforms"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"name": ("STRING", {"default": "t_comp_00"}),
},
}
RETURN_TYPES = ("TV_COMPOSITION_LIST", "STRING")
RETURN_NAMES = ("transforms_list", "name")
FUNCTION = "compose"
def compose(self, name):
transforms = []
return (transforms, name,)
@ETK_torchvision_base
class TorchVisionTransformStringParser:
"""Parses a newline separated list of torchvision transform calls.
Each line should be either ``TransformName`` or ``TransformName(arg1, arg2)``.
This node helps users quickly define a composition pipeline without having
to add a separate node for each transform.
"""
@classmethod
def INPUT_TYPES(cls):
import inspect
import torchvision.transforms as transforms
import torch.nn as nn
# build default definition text listing all available transform
# classes and their constructor parameters for quick reference
lines = []
for name, obj in inspect.getmembers(transforms, inspect.isclass):
if obj.__module__.startswith("torchvision.transforms") and issubclass(obj, nn.Module):
sig = inspect.signature(obj.__init__)
params = [p for p in sig.parameters.values() if p.name != "self"]
p_names = ", ".join(p.name for p in params)
lines.append(f"{name}({p_names})")
default_def = "\n".join(lines)
return {
"required": {
"definition": ("STRING", {"default": default_def, "multiline": True}),
"name": ("STRING", {"default": "t_comp_00"}),
},
}
RETURN_TYPES = ("TV_COMPOSITION_LIST", "STRING")
RETURN_NAMES = ("transforms_list", "name")
FUNCTION = "compose"
def compose(self, definition, name):
import re
lines = [l.strip() for l in definition.splitlines() if l.strip()]
transforms = []
for line in lines:
m = re.match(r"(\w+)(?:\((.*)\))?", line)
if not m:
continue
t_name = m.group(1)
args_str = m.group(2)
if args_str:
args = [a.strip() for a in args_str.split(',') if a.strip()]
else:
args = None
transforms.append([t_name, args])
return (transforms, name,)
@ETK_torchvision_base
class TorchVisionTransformNode:
"""Provides the base class for torchvision transforms"""
@classmethod
def INPUT_TYPES(cls):
import torchvision.transforms as transforms
import inspect
import torch.nn as nn
# gather torchvision transform classes for dropdown selection
t_list = []
c_list = []
for name, obj in inspect.getmembers(transforms, inspect.isclass):
if obj.__module__.startswith("torchvision.transforms") and issubclass(obj, nn.Module):
t_list.append([name, obj])
c_list.append(name)
return {
"required": {
"composition": ("TV_COMPOSITION_LIST", {"default": None}),
"transform": (c_list,),
}, "optional": {
"args": ("STRING", {"default": ""}),
},
}
RETURN_TYPES = ("TV_COMPOSITION_LIST",)
FUNCTION = "compose"
def compose(self, composition: list, transform, args=None):
composition = composition.copy()
if args:
if not "(" in args[0] and not ")" in args[-1]:
args = args.split(",")
else:
args = [args]
composition.append([transform, args])
return (composition,)
@ETK_torchvision_base
class TorchVisionCallComposition:
"""
creates the composition object and calls it on the input TORCH_TENSOR or IMAGE
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"composition": ("TV_COMPOSITION_LIST", {"default": None}),
}, "optional": {
"tensor": ("TORCH_TENSOR",),
"image": ("IMAGE",),
},
}
RETURN_TYPES = ("TORCH_TENSOR", "IMAGE", "TV_COMPOSITION")
FUNCTION = "compose"
def compose(self, composition, tensor=None, image=None):
# we assume that both tensor and image are fit for the transform and that the transform is valid
if tensor == None and image == None:
raise ValueError("Must provide either a tensor or an image")
if tensor != None and image != None:
raise ValueError("Must provide either a tensor or an image, not both")
if tensor == None and image != None:
tensor = image
# transform args that might be ints or floats
composition = composition.copy()
for c_name, c_args in composition:
if c_args:
if isinstance(c_args, list):
for i in range(len(c_args)):
try:
if "(" in c_args[i][0] and ")" in c_args[i][-1]:
c_args[i] = [eval(c_args[i])][0]
elif "." in c_args[i]:
c_args[i] = float(c_args[i])
else:
c_args[i] = int(c_args[i])
except:
pass
# create the transform composition list that will be used with transforms.Compose
transform_composition = []
for l_c_name, l_c_args in composition:
if l_c_name:
if l_c_args is None:
l_c_args = []
transform_composition.append(getattr(torchvision.transforms, l_c_name)(*l_c_args))
# now that the list is composed use it to create the compostion object
transform_composition = torchvision.transforms.Compose(transform_composition)
# apply the transform to each item in the batch
# many torchvision transforms expect a single image tensor of shape C,H,W
tensor = tensor.permute(0, 3, 1, 2) # B,C,H,W
out_list = []
for img in tensor:
out_list.append(transform_composition(img))
tensor = torch.stack(out_list, dim=0)
tensor = tensor.permute(0, 2, 3, 1)
return (tensor, tensor, transform_composition,)
utils.py
-155
View File
@@ -1,155 +0,0 @@
import copy
from collections import OrderedDict
def etk_deep_copy(x, verbose=False):
if type(x).__module__ == "numpy":
return x.copy()
elif type(x).__module__ == "torch":
return x.clone()
elif isinstance(x, list):
return [etk_deep_copy(item) for item in x]
elif isinstance(x, tuple):
return tuple(etk_deep_copy(item) for item in x)
elif isinstance(x, dict):
return {k: etk_deep_copy(v) for k, v in x.items()}
elif isinstance(x, set):
return set(etk_deep_copy(item) for item in x)
elif isinstance(x, OrderedDict):
return OrderedDict((k, etk_deep_copy(v)) for k, v in x.items())
else:
if verbose:
print(f"copying {x}, of type {type(x)}")
return copy.deepcopy(x)
import json
import os
from collections.abc import MutableMapping
class FileBackedDict(MutableMapping):
FBD_MARKER = "__is_fbd__" # Marker to indicate nested FileBackedDicts
def __init__(self, file_path):
# Ensure the provided path is absolute, raise an error if not
if not os.path.isabs(file_path):
raise ValueError(f"The path '{file_path}' is not an absolute path.")
# Convert file_path to an absolute path (if needed)
self.file_path = os.path.abspath(file_path)
# check if the file exists, if it doesnt create it with the current data
if not os.path.exists(self.file_path):
self._data = {}
self._save()
else:
self._data = {}
self._load()
def _load(self):
"""Loads the data from the file into memory."""
if os.path.exists(self.file_path):
with open(self.file_path, 'r', encoding='utf-8') as file:
try:
raw_data = json.load(file)
# Convert any sub-dictionaries marked as FileBackedDicts
for key, value in raw_data.items():
if isinstance(value, dict) and value.get(self.FBD_MARKER, False):
# Create a nested FileBackedDict for each sub-dictionary
sub_dict_file = f"{self.file_path}_{key}.json"
self._data[key] = FileBackedDict(sub_dict_file)
else:
self._data[key] = value
except json.JSONDecodeError:
self._data = {}
else:
self._data = {}
def _save(self, debug=False):
"""Saves the current data to the file."""
# Debug print to check if the directory and file path exist
if debug:
print(f"Saving file to: {self.file_path}")
directory = os.path.dirname(self.file_path)
# Ensure the directory exists, create if it doesn't
if not os.path.exists(directory):
if debug:
print(f"Directory does not exist, creating: {directory}")
os.makedirs(directory, exist_ok=True)
else:
if debug:
print(f"Directory exists: {directory}")
# Check if the file exists
if not os.path.exists(self.file_path):
if debug:
print(f"File does not exist, it will be created: {self.file_path}")
else:
if debug:
print(f"File already exists: {self.file_path}")
# Save only the top-level keys that are not FileBackedDicts
data_to_save = {}
for key, value in self._data.items():
if isinstance(value, FileBackedDict):
# Mark this key as being backed by a FileBackedDict
data_to_save[key] = {self.FBD_MARKER: True}
value._save() # Trigger save for the sub-dictionaries
else:
data_to_save[key] = value
# Save the metadata about the keys, creating the file if it does not exist using utf-8 encoding
with open(self.file_path, 'w', encoding='utf-8') as file:
if debug:
print(f"Writing data to file: {self.file_path}")
json.dump(data_to_save, file, ensure_ascii=False, indent=4)
if debug:
print("File written successfully.")
def __getitem__(self, key):
self._load() # Ensure the data is up-to-date
return self._data[key]
def __setitem__(self, key, value):
self._load() # Ensure the data is up-to-date
if isinstance(value, dict):
# Create a nested FileBackedDict for each sub-dictionary
sub_dict_file = f"{self.file_path}_{key}.json"
self._data[key] = FileBackedDict(sub_dict_file)
self._data[key].update(value)
else:
self._data[key] = value
self._save()
def __delitem__(self, key):
self._load() # Ensure the data is up-to-date
del self._data[key]
self._save()
def __iter__(self):
self._load() # Ensure the data is up-to-date
return iter(self._data)
def __len__(self):
self._load() # Ensure the data is up-to-date
return len(self._data)
def __repr__(self):
self._load() # Ensure the data is up-to-date
return repr(self._data)
def clear(self):
"""Clear the dictionary and save the changes to the file."""
self._data.clear()
self._save()
def update(self, *args, **kwargs):
self._load() # Ensure the data is up-to-date
"""Update the dictionary with new data and save the changes to the file."""
for key, value in dict(*args, **kwargs).items():
self.__setitem__(key, value)
self._save()
video.py
-2282
View File
File diff suppressed because it is too large Load Diff
video_v2.py
View File
web/marked.js
-1191
View File
File diff suppressed because it is too large Load Diff
web/moviemain.js
-107
View File
@@ -1,107 +0,0 @@
import {app} from "../../scripts/app.js";
import {api} from "../../scripts/api.js";
app.registerExtension({
name: "movienodes.MovieMain",
async init(app) {
const canvas = document.querySelector("#graph-canvas");
// Add event listener for the drop event
canvas.addEventListener("drop", this.handleDrop.bind(this), false);
//canvas.addEventListener("dragover", (e) => {
// e.preventDefault()
// console.log(e), false
//});
// Allow drop events on the app canvas
canvas.addEventListener("dragover", this.handleDragOver.bind(this), false);
canvas.addEventListener("dragenter", this.handleDragEnter.bind(this), false);
canvas.addEventListener("dragleave", this.handleDragLeave.bind(this), false);
},
async handleDrop(event) {
// Prevent default behavior to prevent the browser from opening the dropped file
event.preventDefault();
const supportsFileSystemAccessAPI =
'getAsFileSystemHandle' in DataTransferItem.prototype;
// Get the file(s) from the dropped event
const files = event.dataTransfer.files;
// Check if there are any files dropped
if (files.length > 0) {
// Iterate through the dropped files
for (let i = 0; i < files.length; i++) {
const fileHandlesPromises = [...event.dataTransfer.items]
// …by including only files (where file misleadingly means actual file _or_
// directory)…
//.filter((item) => item.kind === 'file')
// …and, depending on previous feature detection…
.map((item) =>
supportsFileSystemAccessAPI
// …either get a modern `FileSystemHandle`…
? item.getAsFileSystemHandle()
// …or a classic `File`.
: item.getAsFile(),
);
for await (const handle of fileHandlesPromises) {
// This is where we can actually exclusively act on the files.
console.log(`File: ${handle.name}`);
}
const file = fileHandlesPromises[i];
const formData = new FormData();
formData.append("name", file.name);
formData.append("video", file);
try {
const resp = await api.fetchApi("/SWAIN/SaveFile", {
method: "POST",
body: formData
});
if (resp.status === 200) {
// File uploaded successfully
// Get uploaded filename
const filename = await resp.json();
// Create node
const node = LiteGraph.createNode("VideoFileToImageStack");
app.graph.add(node);
// Set widget value to filename
// filter to the widget that has the name "vide_in"
const widget = node.widgets.find(w => w.name === "video_in");
widget.value = file.name;
} else {
console.error("Upload failed");
}
} catch (err) {
console.error("Error uploading file", err);
}
}
}
},
handleDragOver(event) {
// Allow dropping on the document
event.preventDefault();
}
,
handleDragEnter(event) {
// Add a class to the document when a file is dragged over it
document.body.classList.add("dragover");
}
,
handleDragLeave(event) {
// Remove the class from the document when a file is dragged out of it
document.body.classList.remove("dragover");
}
,
})
;
web/showMarkdown.js
-47
View File
@@ -1,47 +0,0 @@
import { app } from "../../scripts/app.js";
import { ComfyWidgets } from "../../scripts/widgets.js";
import Marked from "./marked.js";
// Displays input markdown on a node
app.registerExtension({
name: "pysssss.ShowMarkdown",
async beforeRegisterNodeDef(nodeType, nodeData, app) {
if (nodeData.internal_state_display || nodeData.name == "ShowMarkdown") {
// When the node is executed we will be sent the input markdown, display this in the widget
const onExecuted = nodeType.prototype.onExecuted;
nodeType.prototype.onExecuted = function (message) {
onExecuted?.apply(this, arguments);
if (this.widgets) {
const pos = this.widgets.findIndex((w) => w.name === "markdown");
if (pos !== -1) {
for (let i = pos; i < this.widgets.length; i++) {
this.widgets[i].onRemove?.();
}
this.widgets.length = pos;
}
}
for (const list of message.text) {
const w = ComfyWidgets["STRING"](this, "markdown", ["STRING", { multiline: true }], app).widget;
// Convert markdown to HTML
const html = Marked(list);
// Create a new div element to hold the rendered markdown
const div = document.createElement("div");
div.classList.add("markdown-widget");
div.style.opacity = 0.6;
div.innerHTML = html;
// Replace the input element with the div containing the rendered markdown
w.inputEl.parentNode.replaceChild(div, w.inputEl);
}
this.onResize?.(this.size);
};
}
},
});
web/showText.js
-118
View File
@@ -1,118 +0,0 @@
import {app} from "/scripts/app.js";
import {ComfyWidgets} from "/scripts/widgets.js";
//import {createAddNodeWidget} from "./AddNodeWidget.js";
app.registerExtension({
name: "etk.ShowText",
async beforeRegisterNodeDef(nodeType, nodeData, app) {
const use_name = "ETK Show Text";
const use_name_2 = "ShowList";
// create an array with these two
const use_names =
[
"ETK Show Text",
"ShowList",
"ShowDict"
];
// if nodeData.name in use_names
if (use_names.includes(nodeData.name)) {
const onNodeCreated = nodeType.prototype.onNodeCreated;
nodeType.prototype.onNodeCreated = function () {
onNodeCreated?.apply(this, arguments);
// const nodeParams = {
// buttonLabel: "Add Node",
// nodeName: "Show Text",
// outputType: "STRING",
// inputSlotName: "text"
// };
// createAddNodeWidget(this, app, nodeParams);
if (!nodeData.widgets) {
return;
} else {
const widget = this.widgets.find((w) => w.name === "text_display");
if (widget) {
widget.inputEl.readOnly = true;
widget.inputEl.style.opacity = 0.6;
widget.inputEl.style.userSelect = "text";
widget.inputEl.style.webkitUserSelect = "text";
widget.inputEl.style.msUserSelect = "text";
widget.inputEl.style.mozUserSelect = "text";
}
}
}
}
;
// Check if nodeData has a property 'internal_state_display' to identify nodes used for display.
if (use_names.includes(nodeData.name)) {
// Debugging: print the nodeData to the console.
console.log(nodeData);
// Function to populate the node with text widgets.
function populate(text) {
this.display_widget_idx = this.widgets.findIndex((w) => w.name === "text_display");
// Remove the widget with the name of the text display widget
const widget = this.widgets.find((w) => w.name === "text_display");
if (widget) {
widget.onRemove?.();
this.widgets.splice(this.widgets.indexOf(widget), 1);
}
// Create a new widget with the name of the text display widget and the correct text
const w = ComfyWidgets["STRING"](this, "text_display", ["STRING", {multiline: true}], app).widget;
// Make the widget's input element read-only and semi-transparent.
w.inputEl.readOnly = true;
w.inputEl.style.opacity = 0.6;
w.inputEl.style.userSelect = "text";
w.inputEl.style.webkitUserSelect = "text";
w.inputEl.style.msUserSelect = "text";
w.inputEl.style.mozUserSelect = "text";
// If text has only one item, update the w.value directly without indexing
if (text.length === 1) {
w.value = text[0];
} else {
// Set the value of the widget to the current text item, use the correct item index
w.value = text[this.display_widget_idx];
}
// Adjust the node size after adding widgets.
requestAnimationFrame(() => {
const sz = this.computeSize();
// Ensure the node size is not smaller than its initial size.
if (sz[0] < this.size[0]) {
sz[0] = this.size[0];
}
if (sz[1] < this.size[1]) {
sz[1] = this.size[1];
}
// Trigger a resize event with the new size.
this.onResize?.(sz);
// Mark the canvas as dirty to update the UI.
app.graph.setDirtyCanvas(true, false);
});
}
// Override the 'onExecuted' method to display text when the node is executed.
const onExecuted = nodeType.prototype.onExecuted;
nodeType.prototype.onExecuted = function (message) {
// Call the original 'onExecuted' method.
onExecuted?.apply(this, arguments);
// Populate the node with the text from the execution message.
populate.call(this, message.text);
};
// Override the 'onConfigure' method to display stored widget values.
const onConfigure = nodeType.prototype.onConfigure;
nodeType.prototype.onConfigure = function () {
// If there are stored widget values, populate the node with these values.
if (this.widgets_values?.length) {
populate.call(this, this.widgets_values);
}
};
}
},
});
web/ui_helpers.js
-225
View File
@@ -1,225 +0,0 @@
import {app} from "../../../scripts/app.js";
import {api} from "../../../scripts/api.js";
export function setupButtons(gdeState, app, api) {
// Save To Git Button
const menu = document.querySelector(".comfy-menu");
const saveButton = document.createElement("button");
saveButton.textContent = "Save To Cloud";
saveButton.onclick = async () => {
try {
const json = JSON.stringify(app.graph.serialize(), null, 2);
const userNameValue = gdeState.user;
const graphSelected = gdeState.graph;
await api.fetchApi("/gde/git/save", {
method: "POST",
headers: {
"Content-Type": "application/json",
},
body: JSON.stringify({
data: json,
file_name: graphSelected,
user_name: userNameValue
}),
});
} catch (error) {
console.error(error);
}
};
menu.append(saveButton);
// Load From Git Button
const loadButton = document.createElement("button");
loadButton.textContent = "Load From Cloud";
loadButton.onclick = async () => {
try {
const userNameValue = gdeState.user;
const graphSelected = gdeState.graph;
const response = await api.fetchApi("/gde/git/load", {
// needs user_name file_name
method: "POST",
headers: {
"Content-Type": "application/json",
},
body: JSON.stringify({
user_name: userNameValue,
file_name: graphSelected
}
),
});
const data = await response.json();
app.loadGraphData(JSON.parse(data.data));
} catch (error) {
console.error(error);
}
};
menu.append(loadButton);
}
export function setupModal(gdeState) {
// Modal
const modal = document.createElement("div");
modal.id = "graph-name-modal";
modal.style.display = "none"; // Initially hidden
modal.style.position = "fixed";
modal.style.top = "50%";
modal.style.left = "50%";
modal.style.transform = "translate(-50%, -50%)";
modal.style.padding = "20px";
modal.style.backgroundColor = "#fff";
modal.style.border = "1px solid #000";
modal.style.zIndex = "1000"; // Ensure it's on top of other elements
document.body.append(modal);
// Input field inside the modal
const modalInput = document.createElement("input");
modalInput.type = "text";
modalInput.placeholder = "Enter graph name";
modal.append(modalInput);
// Submit button inside the modal
const submitButton = document.createElement("button");
submitButton.textContent = "Submit";
let targetDropdownId = ''; // This will store the ID of the dropdown we're targeting
modal.submit = () => {
// this is the function that will be called when the submit button is clicked
// it will add the new option to the dropdown and select it
const dropdown = document.querySelector(`#${targetDropdownId}`);
const newOption = document.createElement("option");
newOption.value = modalInput.value;
newOption.textContent = modalInput.value;
dropdown.append(newOption);
dropdown.value = modalInput.value;
console.log("Submitted value:", modalInput.value);
modal.style.display = "none"; // Hide the modal
return modalInput.value;
}
submitButton.onclick = () => {
modal.submit()
};
modal.append(submitButton);
// Function to show the modal (this can be used outside as well)
window.showModal = function (dropdownId) {
targetDropdownId = dropdownId; // Store the passed dropdown ID
modal.style.display = "block";
};
}
HTMLSelectElement.prototype.clearDropDown = function (placeHolder = "") {
// remove everything
while (this.options.length > 0) {
this.remove(0);
}
// only add if there is a placeholder
if (placeHolder != "") {
//add select+type
const defaultOption = document.createElement("option");
defaultOption.value = "";
defaultOption.textContent = placeHolder
this.append(defaultOption);
// Add "Create New" option
const createNewOption = document.createElement("option");
createNewOption.value = "create_new";
createNewOption.textContent = "Create New";
this.append(createNewOption);
}
};
export function setupDropdown(menu, api, placeholder, apiEndpoint, dropdownId,
user = null, onCreateNew = null,
onChanged = null) {
let dropdown = document.getElementById(dropdownId);
// If dropdown already exists, clear its options first
if (dropdown) {
dropdown.clearDropDown(placeholder)
} else {
// Create new dropdown if it doesn't exist
dropdown = document.createElement("select");
dropdown.style.marginTop = "2px";
dropdown.style.width = "100%";
dropdown.style.borderRadius = "8px";
dropdown.style.fontSize = "18px";
dropdown.style.backgroundColor = "#222222";
dropdown.style.color = "#EEEEEE";
dropdown.id = dropdownId;
//use comfy-list style
dropdown.clearDropDown(placeholder)
menu.append(dropdown);
// Add event listener to the dropdown
dropdown.addEventListener("change", (event) => {
// Check if "Create New" option was selected
if (event.target.value === "create_new" && onCreateNew) {
onCreateNew(event, dropdown);
}
if (onChanged) {
onChanged(event, dropdown);
}
}
);
}
// dropdown.addEventListener("change", (event) => {
// // Check if "Create New" option was selected
// if (event.target.value === "create_new" && onCreateNew) {
// onCreateNew(event, dropdown);
// }
//
// });
setTimeout(async () => {
dropdown.clearDropDown(placeholder)
await someFunction(dropdownId, api, apiEndpoint, placeholder, user)
}, 50)
return dropdown;
}
export async function someFunction(dropdownId, api, apiEndpoint, placeholder, user = null) {
// this is the function that will be called to fetch the options for the dropdown
try {
const dropdown = document.getElementById(dropdownId);
dropdown.disabled = true; // Disable the file dropdown
const clientId = api.socket.url.split("=")[1]
const queryParams = {
user: user || "",
client_id: clientId
};
const response = await api.fetchApi(apiEndpoint, {
method: "POST",
headers: {
"Content-Type": "application/json",
},
body: JSON.stringify(queryParams)
});
const optionsData = await response.json();
optionsData.forEach(optionItem => {
// if the option already exists, don't add it again
if (dropdown.querySelector(`option[value="${optionItem.value}"]`)) {
return;
}
const optionElement = document.createElement("option");
optionElement.value = optionItem.value;
optionElement.textContent = optionItem.label;
dropdown.append(optionElement);
});
dropdown.disabled = false; // Enable the file dropdown
} catch (error) {
console.error(`Error fetching options for ${placeholder}:`, error);
}
}
web/ui_modifications.js
-413
View File
@@ -1,413 +0,0 @@
import {app} from "../../../scripts/app.js";
import {api} from "../../../scripts/api.js";
import {setupButtons, setupModal, setupDropdown} from './ui_helpers.js';
let loggedInEvent = null;
let loggedOutEvent = new Event('loggedOutEvent');
export let selectedGraph = null;
//const graphCreateNewEvent = new Event('graphCreateNew');
// create a place to store some variables, including security_token
class GDELogin {
constructor(api) {
this.api = api;
this.security_token = null;
this.initElements();
}
initElements() {
this.login_form = this.createElement("form", {id: "gde-login"});
// make sure the form doesnt cause the page to reload
this.login_form.addEventListener("submit", (event) => {
event.preventDefault();
});
//this.login_div = this.createElement("div", {id: "gde-login"});
this.usernameField = this.createInputField("username-field", "Username");
this.passwordField = this.createInputField("password-field", "Password", "password");
this.loginButton = this.createButton("login-button", "Login", this.handleLogin.bind(this));
this.logoutButton = this.createButton("logout-button", "Logout", this.handleLogout.bind(this));
}
createInputField(id, placeholder, type = "text") {
return this.createElement("input", {
id: id,
placeholder: placeholder,
type: type,
});
}
createButton(id, text, clickHandler) {
const button = this.createElement("button", {
id: id,
innerText: text,
});
button.addEventListener("click", clickHandler);
return button;
}
createElement(tag, attributes) {
const element = document.createElement(tag);
for (let attr in attributes) {
element[attr] = attributes[attr];
}
return element;
}
async handleLogin() {
const username = this.usernameField.value;
const password = this.passwordField.value;
gdeState.user = username;
const res = await this.api.fetchApi("/gde/login", {
method: "POST",
body: JSON.stringify({
username: username,
password: password,
client_id: this.api.clientId,
}),
});
if (res.status !== 200) {
//if res.logged_in is false then the login failed but only need to alert if this.api.clientId is set
if (res.logged_in === false && this.api.clientId) {
alert("Login failed. Please check your username and password and try again.");
}
if (res.logged_in === true) {
this.username = res.user;
gdeState.user = username;
}
} else {
const security_json = await res.json();
this.security_token = security_json.security;
this.switchToLogoutView();
if (username) {
const event = new CustomEvent('loggedInEvent',
{detail: {'user': username}});
//loggedInEvent.target = {value:json.user};
document.dispatchEvent(event);
} else { // client was refreshed user was lost but client_id was not
this.usernameField.value = security_json.user;
loggedInEvent.target.value = security_json.user;
document.dispatchEvent(loggedInEvent);
}
const welcomeMessage = this.createElement("div", {
id: "welcome-message",
innerText: `Welcome ${this.usernameField.value}`
});
this.login_form.append(welcomeMessage);
}
}
async handleLogout() {
await this.api.fetchApi("/gde/logout", {
method: "POST",
body: JSON.stringify({
client_id: this.api.clientId,
}),
});
document.dispatchEvent(loggedOutEvent);
this.clearDropdowns();
this.switchToLoginView();
const welcomeMessage = document.getElementById("welcome-message");
if (welcomeMessage) welcomeMessage.remove();
}
switchToLogoutView() {
this.usernameField.remove();
this.passwordField.remove();
this.loginButton.remove();
this.login_form.append(this.logoutButton);
}
switchToLoginView() {
this.logoutButton.remove();
this.login_form.append(this.usernameField, this.passwordField, this.loginButton);
}
clearDropdowns() {
this.clearDropdown("graph-name-field");
//this.clearDropdown("branch-name-field");
}
clearDropdown(id) {
const dropdown = document.getElementById(id);
if (dropdown) {
while (dropdown.options.length > 2) {
dropdown.remove(2);
}
}
}
async getServerState(client_id) {
// check with the server to see if we have a valid client_id
const response = await api.fetchApi("/gde/login", {
method: "POST",
headers: {
"Content-Type": "application/json",
},
body: JSON.stringify({
client_id: client_id,
}),
});
const data = await response.json();
return data;
}
async setupLoginForm() {
const menu = document.querySelector(".comfy-menu");
// create a div in the menu for this extension
//need to check with server for the status of our current api.clientId
const cid = api.clientId || api.socket.url.split("=")[1];
const server_state = await this.getServerState(cid).then((server_state) => {
const logged_in = server_state.logged_in;
gdeState.user = server_state.user;
const separator = document.createElement("hr");
separator.style.margin = "20px 0";
separator.style.width = "100%";
menu.append(separator);
menu.append(this.login_form);
if (logged_in) {
this.login_form.append(this.logoutButton);
} else {
this.login_form.append(this.usernameField);
this.login_form.append(this.passwordField);
this.login_form.append(this.loginButton);
}
});
}
}
const gdeLogin = new GDELogin(api);
let gdeState = {
// this is for holding the state of the ui like what is selected etc
graph: null,
branch: null,
user: null,
}
// create a dummy function
// gde.setup_login_form = async () => {
// const menu = document.querySelector(".comfy-menu");
// // create a div in the menu for this extension
//
//
// //need to check with server for the status of our current api.clientId
// const cid = api.clientId || api.socket.url.split("=")[1];
// const server_state = await gde.get_server_state(cid).then((server_state) => {
// const logged_in = server_state.logged_in;
//
// const separator = document.createElement("hr");
// separator.style.margin = "20px 0";
// separator.style.width = "100%";
// menu.append(separator);
// menu.append(gde.login_div);
//
// if (logged_in) {
// gde.login_div.append(gde.logoutButton);
//
// } else {
// gde.login_div.append(gde.usernameField);
// gde.login_div.append(gde.passwordField);
// gde.login_div.append(gde.loginButton);
//
// }
// });
//
// }
// gde.hijack_queueprompt = async () => {
// // is this actually needed?
// }
app.registerExtension({
name: "GDE.UI.RemoveDefaultButtons",
async setup(app) {
document.getElementById("comfy-load-default-button").remove();
document.getElementById("comfy-clipspace-button").remove();
},
});
app.registerExtension({
name: "GDE.UI.Login",
async setup() {
setTimeout(async () => {
await gdeLogin.setupLoginForm();
//await gde.hijack_queueprompt();
}, 10);
}
});
function SetupGraphDropDown(graphDropdown,
event,
menu,
handleGraphCreateNew,
handleGraphSelect) {
// event may come in as null, so we need to populate it with default values
if (!event) {
event = {
detail: {
user: "main",
graph: "default.json",
branch: "main"
},
target: {
value: "main"
}
}
gdeState.user = "main"
}
// graphDropdown.clearDropDown();
graphDropdown.clearDropDown("Select Graph");
const selectedGraph = null
graphDropdown = setupDropdown(menu, api,
"Select Graph",
"/gde/git/user_graphs",
"graph-name-field",
gdeState.user,
handleGraphCreateNew,
handleGraphSelect);
// modify graphDropdown to change gdeState.graph when it is changed
graphDropdown.addEventListener("change", (event) => {
const selectedGraph = event.target.value;
if (selectedGraph && selectedGraph !== gdeState.graph) {
gdeState.graph = selectedGraph;
const graphChangeEvent = new CustomEvent('graphChangeEvent', {
detail: {
graph: selectedGraph
}
});
document.dispatchEvent(graphChangeEvent);
}
});
return graphDropdown;
}
app.registerExtension({
name: "GDE.UI.Git",
async setup() {
setupModal();
setupButtons(gdeState, app, api);
const separator = document.createElement("hr");
const menu = document.querySelector(".comfy-menu");
separator.style.margin = "20px 0";
separator.style.width = "100%";
menu.append(separator);
// const handleBranchCreateNew = (event, dropdown) => {
// // Custom logic for when "Create New" is selected in the branch dropdown
// showModal('branch-name-field'); // Show the modal for creating a new branch
// dropdown.value = ""; // Reset the dropdown value
// // Clear the file dropdown
// const fileDropdown = document.getElementById("graph-name-field");
//
//
// };
document.addEventListener('loggedInEvent', async (event) => {
gdeState.user = event.detail.user;
//gdeState.graph = event.detail.graph;
graphDropdown = SetupGraphDropDown(
graphDropdown,
event,
menu,
handleGraphCreateNew,
handleGraphSelect);
});
// no lets call clear dropdown when there is a loggout
document.addEventListener('loggedOutEvent', async (event) => {
graphDropdown.clearDropDown();
}
);
const handleGraphCreateNew = (event, dropdown) => {
// Custom logic for when "Create New" is selected in the branch dropdown
showModal('graph-name-field');
dropdown.disabled = true;
const modal = document.getElementById("graph-name-modal");
// Hijack its onsubmit function
const prev_submit = modal.submit;
modal.submit = () => {
let graphName = prev_submit();
if (graphName) {
// Add the new graph to the dropdown
const optionElement = document.createElement("option");
optionElement.value = graphName;
optionElement.textContent = graphName;
dropdown.append(optionElement);
// Set the new graph as the selected value
dropdown.value = graphName;
gdeState.graph = graphName;
// Dispatch a custom event to notify other parts of the application
const graphChangeEvent = new CustomEvent('graphChangeEvent', {
detail: {
graph: graphName
}
});
document.dispatchEvent(graphChangeEvent);
}
dropdown.disabled = false;
};
};
const handleGraphSelect = (event, dropdown) => {
// change gdeState.graph when it is changed
const selectedGraph = event.target.value;
gdeState.graph = selectedGraph;
};
//setupDropdown(menu, api, "Select Branch", "/gde/git/users", "branch-name-field", null, null, handleBranchCreateNew);
//const branchDropdown = document.getElementById("branch-name-field");
// branchDropdown.addEventListener("change", (event) => {
// const selectedBranch = event.target.value;
// if (selectedBranch && selectedBranch !== "Create New") {
// // Repopulate the file/graph dropdown based on the selected branch
// setupDropdown(menu, api, "Select Graph", "/gde/git/user_graphs", "graph-name-field", selectedBranch, null, handleGraphCreateNew);
// }
//
// });
let tmp = setupDropdown(menu,
api,
"Select Graph",
"/gde/git/user_graphs",
"graph-name-field",
gdeState.user,
handleGraphCreateNew,
handleGraphSelect);
let graphDropdown = SetupGraphDropDown(tmp,
null,
menu,
handleGraphCreateNew,
handleGraphSelect);
}
});
youtube_nodes.py
-76
View File
@@ -1,76 +0,0 @@
import subprocess
import os
NODE_CLASS_MAPPINGS = {}
NODE_DISPLAY_NAME_MAPPINGS = {}
from . import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
def ETK_youtube_base(cls):
cls.CATEGORY = "ETK/Youtube"
NODE_CLASS_MAPPINGS[cls.__name__] = cls
return cls
@ETK_youtube_base
class YoutubeUploadNode:
VALID_PRIVACY_STATUSES = ["public", "private", "unlisted"]
VALID_CATEGORIES = [1, 2, 10, 15, 17, 19, 20, 22, 23, 24, 25, 26, 27, 28, 29]
CATEGORY_NAMES = ['Film & Animation', 'Autos & Vehicles', 'Music', 'Pets & Animals', 'Sports', 'Travel & Events',
'Gaming', 'People & Blogs', 'Comedy', 'Entertainment', 'News & Politics', 'Howto & Style',
'Education', 'Science & Technology', 'Nonprofits & Activism']
OUTPUT_NODE = True
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"file_path": ("STRING", {"description": "Path to the video file to upload"}),
"category": (cls.CATEGORY_NAMES,),
"title": ("STRING", {"default": "My video", "description": "Title of the video"}),
"description": ("STRING", {"default": "My video description", "multiline": True,
"description": "Description of the video"}),
"keywords": ("STRING", {"default": "tag1, tag2, tag3", "description": "Keywords for the video"}),
"privacy_status": (cls.VALID_PRIVACY_STATUSES,)
},
"optional": {
"thumbnail_path": ("STRING", {"default": "", "description": "Path to the thumbnail image file"})
}
}
RETURN_TYPES = ('STRING',)
FUNCTION = "upload_video"
CATEGORY = "ETK/Youtube"
def upload_video(self, file_path, category, title, description, keywords, privacy_status, thumbnail_path=""):
import sys
my_python_exe = sys.executable
category_idx = self.CATEGORY_NAMES.index(category)
category = self.VALID_CATEGORIES[category_idx]
args = ["--file", file_path, "--category", str(category), "--title", title, "--description", description,
"--keywords", keywords, "--privacyStatus", privacy_status]
if thumbnail_path:
args.extend(["--thumbnail", thumbnail_path])
current_dir = os.getcwd()
os.chdir(os.path.dirname(os.path.realpath(__file__)))
result = subprocess.run([my_python_exe, "./modules/youtube.py"] + args,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True)
os.chdir(current_dir)
error = result.stderr
if error:
raise Exception(error)
return (result.stdout,)
if __name__ == "__main__":
pass