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:
@@ -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
@@ -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?
|
||||
|
||||
@@ -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}')
|
||||
|
||||
@@ -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
@@ -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 = {}
|
||||
|
||||
|
||||
@@ -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,)
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -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()
|
||||
@@ -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,3 +1,4 @@
|
||||
# TODO: replace pyfastnoisesimd with pyfastnoiselite (works with rpi)
|
||||
import pyfastnoisesimd as fns
|
||||
|
||||
|
||||
|
||||
@@ -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()
|
||||
@@ -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
@@ -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
|
||||
@@ -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
@@ -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
|
||||
|
||||
+1191
File diff suppressed because it is too large
Load Diff
@@ -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
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user