New Features: Added nodes for advanced video and audio processing

Key updates:

Added AddFramesToVideoFolder node to add frames to a video folder
Introduced VideoFolderProvider and VideoDefinitionProvider nodes to manage video folders and video definitions
Added FadeToInVideoFolder node for fading between frames in a video folder
Implemented ListedSVDConditioning node for applying SVD conditioning to a list of parameters
Added VideoFolderFFMPEGVideoInterpolate node to interpolate frames in a video folder using FFmpeg
Introduced AudioFolderProvider and AudioDefinitionProvider nodes for managing audio folders and definitions
Added AddAudioToAudioFolder and AddSilenceToAudioFolder nodes for adding audio files and silence to audio folders
Implemented GetAudioFilesLengthsFromAudioFolder node to retrieve lengths of audio files in a folder
Added AudioFolderRenderAudio node to join audio files in an audio folder
Introduced GenerateMusic node for generating music from descriptions using the AudioCraft library
Added XttsNode for text-to-speech conversion using the XTTS library
Implemented AdjustAudioVolumeForFiles node to adjust volume for a list of audio files
The changes also include the addition of LiteLLM nodes for text generation:

Added FuncAnysToText node that runs Python code with multiple input types and outputs a string
Introduced FuncAnysToLatent node for running code that outputs a latent representation
Added FuncAnysToString node for executing code with various inputs and producing a string output
Implemented FuncAnysToList node to run code that takes different input types and returns a list
Added FuncAnysToCond node for running code with multiple inputs and generating a conditioning output

These new nodes expand the capabilities for video and audio processing, allowing more advanced manipulations and generation of video and audio content within the graph.
This commit is contained in:
Tasha Upchurch
2024-04-15 13:29:31 -04:00
parent 447e9403ae
commit d9170f804b
20 changed files with 4832 additions and 175 deletions
+159
View File
@@ -0,0 +1,159 @@
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
+5 -2
View File
@@ -1,5 +1,6 @@
import os
import importlib
WEB_DIRECTORY = "./web"
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
OLD_NODE_CLASS_MAPPINGS = {
@@ -33,12 +34,14 @@ NODE_CLASS_MAPPINGS = {}
NODE_DISPLAY_NAME_MAPPINGS = {}
# Try importing optional dependencies
import sys
if os.getenv("UNIT_TEST", False):
if os.getenv("UNIT_TEST", False) or 'pytest' in sys.modules:
pass
else: # if not a unit test
try:
import comfy.samplers
use_generic_nodes = False
except ImportError:
use_generic_nodes = True
@@ -53,7 +56,7 @@ else: # if not a unit test
if file.endswith(".py") and file[0] != "_":
# import the NODE_CLASS_MAPPINGS from the file
try:
module = importlib.import_module(f".{file[:-3]}", __name__)
module = importlib.import_module(f".{file[:-3]}", __name__)
except ImportError:
try:
module = importlib.import_module(f"{file[:-3]}") # needed for testing?
+408 -5
View File
@@ -1,9 +1,18 @@
import copy
import shutil
import copy
import imageio_ffmpeg as ffmpeg
import joblib
from config import config_settings
temp_dir = config_settings["tmp_dir"]
NODE_CLASS_MAPPINGS = {}
NODE_DISPLAY_NAME_MAPPINGS = {}
# Create a memory object for caching
# memory = joblib.Memory(temp_dir, verbose=0)
def ETK_audio_base(cls):
cls.CATEGORY = "ETK/audio"
NODE_CLASS_MAPPINGS[cls.__name__] = cls
@@ -73,10 +82,9 @@ class AudioFileListFromPathPattern:
return (l_fpaths, l_bytes, l_metadata,)
import librosa
import soundfile as sf
import numpy as np
@ETK_audio_base
class JoinAudioFiles:
""" Joins a list of audio files into a single audio file """
@@ -95,6 +103,8 @@ class JoinAudioFiles:
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']
@@ -123,6 +133,7 @@ class JoinAudioFiles:
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 """
@@ -132,7 +143,7 @@ class ExtractAudioFromVideo:
return {
"required": {
"video_path": ("STRING", {"multiline": False, "default": ""}),
"format":(["wav", "mp3"], {"default": "wav"}),
"format": (["wav", "mp3"], {"default": "wav"}),
"output_path": ("STRING", {"multiline": False, "default": ""}),
},
"optional": {
@@ -151,7 +162,7 @@ class ExtractAudioFromVideo:
kwargs = copy.deepcopy(kwargs)
video_path = kwargs.get('video_path')
output_path = kwargs.get('output_path',"tmp.wav")
output_path = kwargs.get('output_path', "tmp.wav")
audio_format = kwargs.get('format', 'wav')
bitrate = kwargs.get('bitrate', None) # Default bitrate is None
@@ -174,3 +185,395 @@ class ExtractAudioFromVideo:
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
from .config import this_file_path
audio_folder_def = kwargs.get("AUDIO_FOLDER_DEF", None)
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 = str(kwargs['[text]'])
else:
texts = str([kwargs.get('text', "")])
# Create the command
command = [
python_exe, xttscli_path,
"--text", texts,
"--folder", kwargs.get('output_path', "."),
"--lang", kwargs.get('lang', 'en'),
"--speaker", kwargs.get('speaker_wav', ''),
]
# execute the command in its own full porcess and recieve the output
output = subprocess.check_output(command, shell=True)
generated = output.decode("utf-8")
generated_files = generated.split("Generated files: (")[1].split(")")[0]
generated_files = eval(generated_files)
generated_files = generated_files[1]
# get the output folder
output_folder = os.path.dirname(kwargs.get('file', "output.wav"))
# move the files to the correct files
out_files = []
for src_file in generated_files:
target_file = audio_folder_def.get_next_file_name()
shutil.copy(src_file, target_file)
out_files.append(target_file)
return (out_files[0], out_files, output_folder, audio_folder_def,)
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 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}')
+2
View File
@@ -3,4 +3,6 @@ tmp_dir: comfy_path/tmp
gitio:
COTRF: https://github.com/Hopping-Mad-Games/COTRF.git
StSq: https://github.com/story-squad/GDE_Graph_IO.git
HMG: https://github.com/TashaSkyUp/HMG_GDE_GITIO
base: ./data/gitio
tesseract_location: /usr/bin/tesseract
+79 -8
View File
@@ -22,6 +22,7 @@ def ETK_functional_base(cls):
NODE_CLASS_MAPPINGS[cls.__name__] = cls
return cls
def new_exec(c, g, l):
global old_exec
import os
@@ -36,6 +37,7 @@ def new_exec(c, g, l):
return locals
old_exec = exec
exec = new_exec
@@ -336,7 +338,7 @@ class FuncListToList:
exec(code, my_globals, my_locals)
return (my_locals["y_list"],)
return (my_locals["y"],)
@ETK_functional_base
@@ -390,7 +392,7 @@ class FuncStrToList:
exec(code, my_globals, my_locals)
return (my_locals["y_list"],)
return (my_locals["y"],)
@ETK_functional_base
@@ -483,14 +485,16 @@ class FuncAnysToImage:
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()
my_locals = locals()
my_globals = globals().copy()
my_locals = {}
my_locals["any1"] = any1
my_locals["any2"] = any2
@@ -499,11 +503,21 @@ class FuncAnysToImage:
exec(code, my_globals, my_locals)
if "y" in my_locals.keys():
return (my_locals["y"],)
y = my_locals["y"]
elif "y" in my_globals.keys():
return (my_globals["y"],)
y = my_globals["y"]
else:
return (None,)
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:
@@ -552,6 +566,7 @@ class FuncAnysToLatent:
else:
return (None,)
@ETK_functional_base
class FuncAnysToString:
"""runs eval on the given text, with 5 anys as input and outputs a string"""
@@ -599,6 +614,7 @@ class FuncAnysToString:
else:
return (None,)
@ETK_functional_base
class FuncAnysToList:
"""runs eval on the given text, with 5 anys as input and outputs a list"""
@@ -631,7 +647,7 @@ class FuncAnysToList:
any4 = kwargs.get("any4", None)
my_globals = globals()
my_locals = locals()
my_locals = {}
my_locals["any1"] = any1
my_locals["any2"] = any2
@@ -645,6 +661,61 @@ class FuncAnysToList:
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 = {}
+24
View File
@@ -251,3 +251,27 @@ class TrainUnet2dModel:
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,)
+783 -117
View File
File diff suppressed because it is too large Load Diff
+404
View File
@@ -0,0 +1,404 @@
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 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 = 1/255
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
# subtract the normal top value
interpolated_image[over_top] = interpolated_image[over_top] - 1
# scale the over the top values to the target range
interpolated_image[over_top] *= (over_top_target_range / over_top_range)
# add the target range to the over the top values
interpolated_image[over_top] += (1-over_top_target_range)
# scale the non-over-the top to make room
interpolated_image[not_over_top] *= (1 - over_top_target_range)
# 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
# subtract the normal bottom value
interpolated_image[under_bottom] = interpolated_image[under_bottom] - 0
# scale the under the bottom values to the target range
interpolated_image[under_bottom] *= (under_bottom_target_range / under_bottom_range)
# add the target range to the under the bottom values
interpolated_image[under_bottom] += under_bottom_target_range
# scale the non-under-the-bottom to make room
interpolated_image[not_under_bottom] *= (1 - under_bottom_target_range)
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()
+224
View File
@@ -0,0 +1,224 @@
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()
+1
View File
@@ -1,3 +1,4 @@
# TODO: replace pyfastnoisesimd with pyfastnoiselite (works with rpi)
import pyfastnoisesimd as fns
+94
View File
@@ -0,0 +1,94 @@
import argparse
import torch.multiprocessing as mp
from concurrent.futures import ProcessPoolExecutor, as_completed
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,
# emotion="angry", # doesnt do anything as far as i can tell with xtts2
)
def generate_xtts(**kwargs):
import torch
import ast
import uuid
from concurrent.futures import ProcessPoolExecutor, as_completed
torch.cuda.empty_cache()
text = kwargs.get('text')
folder = kwargs.get('folder', '.')
lang = kwargs.get('lang', 'en')
speaker = kwargs.get('speaker', '')
try:
texts = ast.literal_eval(text)
except SyntaxError:
texts = [text]
o_files = []
batch_size = 3
with ProcessPoolExecutor(max_workers=batch_size) as executor:
futures = []
for i, text in enumerate(texts):
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 generate_xtts_for_text_advanced(text, file, lang, speaker):
# from TTS.api import TTS
# import torch
#
# # Create a SharedMemoryManager
# manager = mp.Manager()
# # Create a shared list to store the result
# result = manager.list()
# model = manager.model()
#
#
# # also we want to SHARE the TTS model
# something = manager.
def main():
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 use for generation')
args = parser.parse_args()
kwargs = vars(args)
print(f"Generating with the following arguments: {kwargs}")
generated_files = generate_xtts(**kwargs)
print(f"Generated files: {generated_files}")
if __name__ == "__main__":
main()
+1
View File
@@ -1186,6 +1186,7 @@ class SaveTorchTensor:
RETURN_TYPES = ("STRING",)
FUNCTION = "save_torch_tensor"
OUTPUT_NODE = True
def save_torch_tensor(self, tensor, file):
import torch
+52 -1
View File
@@ -1,3 +1,4 @@
git+https://github.com/facebookresearch/audiocraft.git
google-api-python-client
google-auth-oauthlib
google-auth-httplib2
@@ -8,4 +9,54 @@ librosa
moviepy
openai-whisper
addict
diffusers
diffusers
# core deps for xtts
numpy==1.22.0;python_version<="3.10"
numpy>=1.24.3;python_version>"3.10"
cython>=0.29.30
scipy>=1.11.2
torch>=2.1
torchaudio
soundfile>=0.12.0
librosa>=0.10.0
scikit-learn>=1.3.0
numba==0.55.1;python_version<"3.9"
numba>=0.57.0;python_version>="3.9"
inflect>=5.6.0
tqdm>=4.64.1
anyascii>=0.3.0
pyyaml>=6.0
fsspec>=2023.6.0 # <= 2023.9.1 makes aux tests fail
aiohttp>=3.8.1
packaging>=23.1
mutagen==1.47.0
# deps for inference
pysbd>=0.3.4
# deps for training
matplotlib>=3.7.0
# coqui stack
trainer>=0.0.36
# config management
coqpit>=0.0.16
#deps for tortoise
einops>=0.6.0
transformers>=4.33.0
#deps for bark
encodec>=0.1.1
# deps for XTTS
unidecode>=1.3.2
num2words
spacy[ja]>=3
TTS
#deps for ocr
pytesseract
+13
View File
@@ -197,6 +197,12 @@ user_data = {
"client_id": None,
"git_io": "COTRF"
},
"HMG": {
"password": "[REDACTED]",
"security": None,
"client_id": None,
"git_io": "HMG"
},
}
@@ -409,6 +415,10 @@ def hijack_prompt_server(route_path,
return part1
if isinstance(part1, web.Request):
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)
@@ -545,6 +555,9 @@ async def my_post_hijack_before(request: web.Request):
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"]
+153 -14
View File
@@ -1,8 +1,17 @@
from unittest import TestCase
import os
import torch
import os
import shutil
import tempfile
import pytest
import torch
from PIL import Image
from ..video 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
@@ -62,9 +71,9 @@ class TestTemporalOpticalFlowSmoothing(TestCase):
import time
import numpy as np
# create a random stack of images
#stack = torch.rand((24*4, 512, 512, 3))
h=256
w=512
# stack = torch.rand((24*4, 512, 512, 3))
h = 256
w = 512
fps = 24
seconds = 2
stack = torch.rand((fps * seconds, h, w, 3))
@@ -77,7 +86,7 @@ class TestTemporalOpticalFlowSmoothing(TestCase):
# fill the box with ones
box = box + 1
# the box width
bw = min(int(w * .2),int(h*.2))
bw = min(int(w * .2), int(h * .2))
# the box height
bh = int(h * .2)
# the box x position
@@ -93,12 +102,12 @@ class TestTemporalOpticalFlowSmoothing(TestCase):
# the box y position
by = int(h * .333)
# fill the box with ones
box[i,:,:,:] = 0
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
# 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)
@@ -115,24 +124,154 @@ class TestTemporalOpticalFlowSmoothing(TestCase):
result = node.handler(image_stack=stack, device="cuda")
result = result[0]
#result = torch.cat(result[0], dim=0)
#print(result.shape)
# result = torch.cat(result[0], dim=0)
# print(result.shape)
# now we can permute the dimensions
#result = result.permute(0, 2,3, 1)
# 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
# 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:
# for flow in result:
# assert flow.shape == torch.Size([1, 2, 512, 512])
# check the file generation and video play
from ..video 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
+1189 -27
View File
File diff suppressed because it is too large Load Diff
View File
+1191
View File
File diff suppressed because it is too large Load Diff
+48
View File
@@ -0,0 +1,48 @@
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);
};
}
},
});
+2 -1
View File
@@ -43,7 +43,8 @@ class YoutubeUploadNode:
import sys
my_python_exe = sys.executable
category = self.CATEGORY_NAMES.index(category)
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]