diff --git a/3d.py b/3d.py
new file mode 100644
index 0000000..804d3b2
--- /dev/null
+++ b/3d.py
@@ -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
diff --git a/__init__.py b/__init__.py
index 1936540..6df4939 100644
--- a/__init__.py
+++ b/__init__.py
@@ -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?
diff --git a/audio.py b/audio.py
index 57cdb8e..c5cd55a 100644
--- a/audio.py
+++ b/audio.py
@@ -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}')
diff --git a/config.yaml b/config.yaml
index 94cc943..6314cee 100644
--- a/config.yaml
+++ b/config.yaml
@@ -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
\ No newline at end of file
diff --git a/functional.py b/functional.py
index 13d5717..4c1c62b 100644
--- a/functional.py
+++ b/functional.py
@@ -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 = {}
diff --git a/hf_diffusers_nodes.py b/hf_diffusers_nodes.py
index e00a205..8777c91 100644
--- a/hf_diffusers_nodes.py
+++ b/hf_diffusers_nodes.py
@@ -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,)
+
+
+
+
diff --git a/image.py b/image.py
index 399bf37..2ad9eec 100644
--- a/image.py
+++ b/image.py
@@ -1,4 +1,5 @@
import os
+
import comfy
try:
@@ -29,7 +30,8 @@ if testing:
pass
else:
try:
- from nodes import VAEDecode, KSampler, CheckpointLoaderSimple, EmptyLatentImage
+ from nodes import VAEDecode, KSampler, CheckpointLoaderSimple, EmptyLatentImage, CLIPTextEncode, VAEDecode, \
+ VAEEncode, PreviewImage
from nodes import CLIPTextEncode, VAEEncode, SaveImage
from nodes import common_ksampler
@@ -54,9 +56,7 @@ except ImportError as e:
import folder_paths
import torch
import torchvision
-import PIL.Image as Image
import os
-import numpy as np
import hashlib
### info for code completion AI ###
@@ -71,10 +71,22 @@ NODE_CLASS_MAPPINGS = {}
NODE_DISPLAY_NAME_MAPPINGS = {}
+# def ETK_image_base(cls):
+# # cls.FUNCTION = "func"
+# cls.CATEGORY = "ETK/Image"
+# NODE_CLASS_MAPPINGS[cls.__name__] = cls
+# return cls
+
def ETK_image_base(cls):
- # cls.FUNCTION = "func"
cls.CATEGORY = "ETK/Image"
+ # 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
@@ -173,10 +185,10 @@ class TinyTxtToImg:
def __init__(self):
print("ETK> TinyTxtToImg init")
- import random
- self.mdl = TinyTxtToImg.share_mdl
- self.clp = TinyTxtToImg.share_clip
- self.vae = TinyTxtToImg.share_vae
+ self.mdl = None
+ self.clp = None
+ self.vae = None
+
self.vision = None
self.steps = 10
self.cfg = 8
@@ -210,8 +222,11 @@ class TinyTxtToImg:
"default": "tinytxt2img"}
),
"overrides": ("STRING", {"multiline": True,
- "default": ""}),
- "FUNC": ("FUNC",)
+ "default": '{"width":768,"height":768,"steps":20}'}),
+ "latent": ("LATENT",),
+ "pos_cond": ("CONDITIONING",),
+ "neg_cond": ("CONDITIONING",),
+ "image": ("IMAGE",),
}
}
@@ -221,19 +236,13 @@ class TinyTxtToImg:
RETURN_NAMES = ("image", "FUNC", "latent", "positive", "negative",)
FUNCTION = "tinytxt2img"
-
def tinytxt2img(self, **kwargs):
""" use the imports from nodes to generate an image from text """
-
+ import comfy.model_management as mm
import random
import json
import gc
-
- # Create a stream for computation
- # gpu_stream = torch.cuda.Stream(device="cuda")
- # cpu_stream = torch.Stream()
-
- # kwargs = deepcopy(kwargs)
+ torch.cuda.empty_cache()
unload_model = kwargs.get("unload_model", True)
prompt = kwargs.get("prompt", "")
@@ -255,45 +264,37 @@ class TinyTxtToImg:
self.denoise = kwargs.get("denoise", 1.0)
self.steps = kwargs.get("steps", 10)
- if not hasattr(TinyTxtToImg, "share_mdl"):
- TinyTxtToImg.share_mdl = None
- TinyTxtToImg.share_mdl_ckpt_name = None
+ self.mdl, self.clp, self.vae = \
+ CheckpointLoaderSimple.load_checkpoint(None,
+ ckpt_name=ckpt_name
+ )
+ # Check if an input image is provided
+ input_image = kwargs.get("image", None)
+ if input_image is not None:
+ # Encode the input image using the model's VAE
+ with torch.no_grad():
+ self.vae.device = torch.device("cuda:0")
+ self.vae.first_stage_model = self.vae.first_stage_model.to("cuda:0")
+ # self.vae.first_stage_model.decoder = self.vae.first_stage_model.decoder.to("cuda:0")
+ # self.vae.first_stage_model = self.vae.first_stage_model.to("cuda:0")
+ input_image = input_image.to("cuda")
+ q = self.vae.encode(input_image)
+ self.latent_image = {"samples": q}
- if TinyTxtToImg.share_mdl is None:
- self.mdl, self.clp, self.vae = \
- CheckpointLoaderSimple.load_checkpoint(None,
- ckpt_name=ckpt_name
- )
- TinyTxtToImg.share_mdl_ckpt_name = ckpt_name
- CSL = TinyTxtToImg
- CSL.share_mdl = self.mdl
- CSL.share_clip = self.clp
- CSL.share_vae = self.vae
- CSL.share_vision = self.vision
+ self.cfg = kwargs.get("cfg", 8)
- else:
- # check if the checkpoint is the same
- if TinyTxtToImg.share_mdl_ckpt_name != ckpt_name:
- self.mdl, self.clp, self.vae = \
- CheckpointLoaderSimple.load_checkpoint(None,
- ckpt_name=ckpt_name
- )
- TinyTxtToImg.share_mdl_ckpt_name = ckpt_name
- CSL = TinyTxtToImg
- CSL.share_mdl = self.mdl
- CSL.share_clip = self.clp
- CSL.share_vae = self.vae
- CSL.share_vision = self.vision
+ if not "sampler_name" in kwargs:
+ if "lcm" in ckpt_name.lower():
+ self.sampler_name = comfy.samplers.KSampler.SAMPLERS[18]
+ self.scheduler = comfy.samplers.KSampler.SCHEDULERS[0]
else:
+ self.sampler_name = comfy.samplers.KSampler.SAMPLERS[0]
+ self.scheduler = comfy.samplers.KSampler.SCHEDULERS[0]
+ self.scheduler = comfy.samplers.KSampler.SCHEDULERS[0]
+ else:
+ self.sampler_name = kwargs.get("sampler_name", comfy.samplers.KSampler.SAMPLERS[0])
+ self.scheduler = kwargs.get("scheduler", comfy.samplers.KSampler.SCHEDULERS[0])
- self.mdl, self.clp, self.vae = (TinyTxtToImg.share_mdl,
- TinyTxtToImg.share_clip,
- TinyTxtToImg.share_vae
- )
-
- self.cfg = 8
- self.sampler_name = comfy.samplers.KSampler.SAMPLERS[0]
- self.scheduler = comfy.samplers.KSampler.SCHEDULERS[0]
self.positive = prompt
self.negative = neg_prompt
self.width = 512
@@ -338,15 +339,14 @@ class TinyTxtToImg:
print(e)
raise ValueError("advanced clip encoder failed")
self.clp.cond_stage_model = self.clp.cond_stage_model.to("cpu")
- import comfy.model_management as mm
torch.cuda.empty_cache()
gc.collect()
if self.latent_image == None:
- self.latent_image = EmptyLatentImage.generate(None, self.width, self.height, self.batch_size)[0]
+ self.latent_image = EmptyLatentImage().generate(self.width, self.height, self.batch_size)[0]
if "latent" in render_what or "all" in render_what or "image" in render_what:
- self.KSampler = KSampler()
+ self.KSampler = ETKKSampler()
# self.mdl.offload_device = torch.device("cuda:0")
# self.mdl.load_device = torch.device("cuda:0")
self.samples = self.KSampler.sample(
@@ -403,7 +403,18 @@ class TinyTxtToImg:
self.pos_cond,
self.neg_cond
,)
- # ret = deepcopy(ret)
+ if unload_model:
+ del self.mdl
+ del self.clp
+ del self.vae
+ if "vision" in self.__dict__:
+ del self.vision
+ del self.samples
+ del self.FUNC
+ del self.ARGS
+ mm.free_memory(mm.get_total_memory(mm.get_torch_device()), mm.get_torch_device())
+ torch.cuda.empty_cache()
+ gc.collect()
return ret
@@ -1334,7 +1345,6 @@ class StackImages:
def stackme(self, image1, image2=None, number_of_images=2, mode="all", number_of_repeats=1):
# import modules to help clear memory
import gc
- import torch
import torch.cuda
"""
@@ -1502,7 +1512,7 @@ class KSamplerCreateDataset:
@classmethod
def INPUT_TYPES(cls):
- inps = KSampler.INPUT_TYPES()
+ inps = ETKKSampler.INPUT_TYPES()
# make sure to set a default
inps["required"]["dataset_name"] = ("STRING", {"default": "dataset"})
return inps
@@ -1532,7 +1542,7 @@ class KSamplerCreateDataset:
:return:
"""
# create the sampler
- sampler = KSampler()
+ sampler = ETKKSampler()
x = {"model": model, "seed": seed, "steps": steps, "cfg": cfg, "sampler_name": sampler_name,
"scheduler": scheduler, "positive": positive, "negative": negative, "latent_image": latent_image,
@@ -1563,22 +1573,13 @@ class KSamplerCreateDataset:
@ETK_image_base
-class KSampler:
+class ETKKSampler:
@classmethod
def INPUT_TYPES(s):
- return {"required":
- {"model": ("MODEL",),
- "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
- "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
- "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}),
- "sampler_name": (comfy.samplers.KSampler.SAMPLERS,),
- "scheduler": (comfy.samplers.KSampler.SCHEDULERS,),
- "positive": ("CONDITIONING",),
- "negative": ("CONDITIONING",),
- "latent_image": ("LATENT",),
- "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
- "one_seed_per_batch": ([True, False], {"default": False}),
- }}
+ from nodes import KSampler
+ ret = KSampler.INPUT_TYPES()
+ ret["required"]["one_seed_per_batch"] = ([True, False], {"default": False})
+ return ret
RETURN_TYPES = ("LATENT", "FUNC",)
FUNCTION = "sample"
@@ -1587,7 +1588,7 @@ class KSampler:
def sample(self, **kwargs):
- kwargs = deepcopy(kwargs)
+ # kwargs = deepcopy(kwargs)
model = kwargs.get("model", None)
seed = kwargs.get("seed", None)
@@ -1601,8 +1602,8 @@ class KSampler:
denoise = kwargs.get("denoise", None)
one_seed_per_batch = kwargs.get("one_seed_per_batch", None)
- self.FUNC = lambda x: KSampler.sample(self, **x)
- self.ARGS = kwargs
+ # self.FUNC = lambda x: ETKKSampler.sample(self, **x)
+ # self.ARGS = kwargs
try:
ret = common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
@@ -1610,11 +1611,319 @@ class KSampler:
except TypeError as e:
ret = common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
denoise=denoise)
+ except RuntimeError as e:
+ raise e
ret = (ret[0], self,)
return ret
-from PIL import Image, ImageDraw, ImageFont
+@ETK_image_base
+class ListKSampler:
+ """
+ A KSampler that takes a list for every input, if the list has one element it will be repeated for every sample
+ """
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ s1 = ETKKSampler.INPUT_TYPES()["required"]["sampler_name"]
+ s2 = ETKKSampler.INPUT_TYPES()["required"]["scheduler"]
+ return {
+ "required": {
+ "model": ("MODEL",),
+ "seed": ("LIST",),
+ "steps": ("LIST",),
+ "cfg": ("LIST",),
+ "sampler_name": s1,
+ "scheduler": s2,
+ "positive": ("LIST",),
+ "negative": ("LIST",),
+ "latent_image": ("LIST",),
+ "denoise": ("LIST",),
+ "one_seed_per_batch": ("LIST",),
+ }
+ }
+
+ RETURN_TYPES = ("LIST",)
+ FUNCTION = "sample"
+
+ CATEGORY = "sampling"
+
+ def sample(self, **kwargs):
+ import torch
+ # figure out how many samples we need to generate
+ num_samples = 1
+ for key in kwargs:
+ if key == "model":
+ continue
+ elif key == "sampler_name" or key == "scheduler":
+ continue
+ if kwargs[key] is not None:
+ num_samples = max(num_samples, len(kwargs[key]))
+
+ # create a list of dictionaries
+ samples = []
+ for i in range(num_samples):
+ sample = {}
+ for key in kwargs:
+ if not isinstance(kwargs[key], list) and key != "model":
+ kwargs[key] = [kwargs[key]]
+ if key == "model":
+ sample[key] = kwargs[key]
+ elif kwargs[key] is None:
+ sample[key] = None
+ elif len(kwargs[key]) == 1:
+ sample[key] = kwargs[key][0]
+ else:
+ sample[key] = kwargs[key][i]
+ samples.append(sample)
+
+ # create a list of results
+ results = []
+ local_ksampler = ETKKSampler()
+ for sample in samples:
+ results.append(local_ksampler.sample(**sample)[0])
+ torch.cuda.empty_cache()
+
+ return (results,)
+
+
+@ETK_image_base
+class ListCLIPTextEncode:
+ """
+ A CLIPTextEncode that takes a list for every input, if the list has one element it will be repeated for every sample
+ """
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "text": ("LIST",),
+ "clip": ("LIST",),
+
+ }
+ }
+
+ RETURN_TYPES = ("LIST",)
+ FUNCTION = "encode"
+
+ CATEGORY = "CLIP"
+
+ def encode(self, **kwargs):
+ # figure out how many samples we need to generate
+ num_samples = 1
+ for key in kwargs:
+ num_samples = max(num_samples, len(kwargs[key]))
+
+ # create a list of dictionaries
+ samples = []
+ for i in range(num_samples):
+ sample = {}
+ for key in kwargs:
+ if len(kwargs[key]) == 1:
+ sample[key] = kwargs[key][0]
+ else:
+ sample[key] = kwargs[key][i]
+ samples.append(sample)
+
+ # create a list of results
+ results = []
+ for sample in samples:
+ results.append(CLIPTextEncode().encode(**sample)[0])
+
+ return (results,)
+
+
+@ETK_image_base
+class ListVAEDecode:
+ """
+ A VAEDecode that takes a list for every input, if the list has one element it will be repeated for every sample
+ """
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "latent": ("LIST",),
+
+ },
+ "optional": {
+ "vae_list": ("LIST",),
+ "vae": ("VAE",),
+
+ }
+ }
+
+ RETURN_TYPES = ("LIST",)
+ FUNCTION = "decode"
+
+ CATEGORY = "VAE"
+
+ def decode(self, **kwargs):
+ # figure out if we are using the list of vae or just one
+ if kwargs.get("vae_list") is None:
+ kwargs["vae"] = [kwargs.get("vae")]
+ else:
+ kwargs["vae"] = kwargs.pop("vae_list")
+
+ if "vae_list" in kwargs:
+ kwargs.pop("vae_list")
+
+ # figure out how many samples we need to generate
+ num_samples = 1
+ for key in kwargs:
+ num_samples = max(num_samples, len(kwargs[key]))
+
+ # create a list of dictionaries
+ samples = []
+ for i in range(num_samples):
+ sample = {}
+ for key in kwargs:
+ if len(kwargs[key]) == 1:
+ sample[key] = kwargs[key][0]
+ else:
+ sample[key] = kwargs[key][i]
+ samples.append(sample)
+
+ # create a list of results
+ results = []
+ for sample in samples:
+ sample["samples"] = sample.pop("latent")
+ results.append(VAEDecode().decode(**sample)[0])
+
+ return (results,)
+
+
+@ETK_image_base
+class ListVAEEncode:
+ """
+ A VAEEncode that takes a list for every input, if the list has one element it will be repeated for every sample
+ """
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "image": ("LIST",),
+
+ },
+ "optional": {
+ "vae_list": ("LIST",),
+ "vae": ("VAE",),
+
+ }
+ }
+
+ RETURN_TYPES = ("LIST",)
+ FUNCTION = "encode"
+
+ CATEGORY = "VAE"
+
+ def encode(self, **kwargs):
+ if kwargs.get("vae_list", None) is None:
+ kwargs["vae"] = [kwargs.get("vae")]
+ else:
+ kwargs["vae"] = kwargs.pop("vae_list")
+
+ if "vae_list" in kwargs:
+ kwargs.pop("vae_list")
+
+ # figure out how many samples we need to generate
+ num_samples = 1
+ for key in kwargs:
+ num_samples = max(num_samples, len(kwargs[key]))
+
+ # create a list of dictionaries
+ samples = []
+ for i in range(num_samples):
+ sample = {}
+ for key in kwargs:
+ if len(kwargs[key]) == 1:
+ sample[key] = kwargs[key][0]
+ else:
+ sample[key] = kwargs[key][i]
+ samples.append(sample)
+
+ # create a list of results
+ results = []
+ for sample in samples:
+ results.append(VAEEncode().encode(**sample)[0])
+
+ return (results,)
+
+
+@ETK_image_base
+class ImageFromListIdx:
+ """
+ A node that takes a list of images and an index and returns the image at that index
+ """
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "images": ("LIST", {"default": []}),
+ "index": ("INT", {"default": 0, "min": 0, "max": 10e10}),
+ }
+ }
+
+ RETURN_TYPES = ("IMAGE",)
+ FUNCTION = "get_image"
+
+ CATEGORY = "ETK/image"
+
+ def get_image(self, images, index):
+ return (images[index],)
+
+
+@ETK_image_base
+class LatentFromListIdx:
+ """
+ A node that takes a list of latents and an index and returns the latent at that index
+ """
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "latents": ("LIST", {"default": []}),
+ "index": ("INT", {"default": 0, "min": -10e10, "max": 10e10}),
+ }
+ }
+
+ RETURN_TYPES = ("LATENT",)
+ FUNCTION = "get_latent"
+
+ CATEGORY = "ETK/latent"
+
+ def get_latent(self, latents, index):
+ return (latents[index],)
+
+
+@ETK_image_base
+class ListifyAnything:
+ """
+ A node that takes anything and returns it in a list
+ """
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "anything": ("*",),
+ "repeat": ("INT", {"default": 1, "min": 1, "max": 100}),
+ }
+ }
+
+ RETURN_TYPES = ("LIST",)
+ FUNCTION = "listify"
+
+ CATEGORY = "ETK/other"
+
+ def listify(self, anything, repeat):
+ return ([anything] * repeat,)
+
+
+from PIL import Image
import torch
import numpy as np
@@ -1670,13 +1979,16 @@ class TextRender:
"func_only": ([True, False], {"default": False}),
"stroke fill": ("STRING", {"default": "#000000"}),
"stroke width": ("INT", {"default": 1, "min": 1, "max": 100, "step": 1}),
+ },
+ "optional": {
+ "font_override": ("STRING", {"default": None}),
}
}
CATEGORY = "ETK/text"
- RETURN_TYPES = ("IMAGE", "IMAGE", "FUNC",)
- RETURN_NAMES = ("image", "image rgba", "FUNC(**kwargs)",)
+ RETURN_TYPES = ("IMAGE", "IMAGE", "STRING",)
+ RETURN_NAMES = ("image", "image rgba", "possible_fonts",)
FUNCTION = "render_text"
@@ -1687,11 +1999,14 @@ class TextRender:
The text is rendered in the provided font, size, and color.
"""
from PIL import ImageFont, ImageDraw
- import gc
import psutil
process = psutil.Process();
memory_info = process.memory_info();
start_memory = memory_info.rss / (1024 * 1024)
+ font_override = kwargs.get("font_override", None)
+
+ if font_override:
+ font = font_override
sw = kwargs.get("stroke width", None)
sf = kwargs.get("stroke fill", None)
@@ -1758,7 +2073,9 @@ class TextRender:
size,
wrapped_line, x, y_offset + line_spacing,
allow_wrap=False,
- allow_shrink=False)
+ allow_shrink=False,
+ sf=sf,
+ sw=sw)
continue # Skip to the next line
if allow_wrap:
@@ -1771,7 +2088,9 @@ class TextRender:
x,
y_offset + line_spacing,
allow_wrap=False,
- allow_shrink=False)
+ allow_shrink=False,
+ sf=sf,
+ sw=sw)
continue # Skip to the next line
# Draw the line
@@ -1823,10 +2142,10 @@ class TextRender:
image_tensor = torch.from_numpy(image_array).unsqueeze(0).float() / 255.0
# set the alpha to the first channel to make a mask
- image_tensor[..., 3] = image_tensor[..., 0]
+ image_tensor[..., 3] = image_tensor[..., 0].to("cpu")
# now set the text to the correct color
- color_arr = torch.zeros(1, 1, 1, 3)
+ color_arr = torch.zeros(1, 1, 1, 3).to("cpu")
color_arr[0, 0, 0, 0] = int(color[1:3], 16) / 255.0
color_arr[0, 0, 0, 1] = int(color[3:5], 16) / 255.0
color_arr[0, 0, 0, 2] = int(color[5:7], 16) / 255.0
@@ -1847,7 +2166,8 @@ class TextRender:
image_rgb = image_rgb.to("cpu")
image_tensor = image_tensor.to("cpu")
- return (image_rgb, image_tensor, None,)
+ list_of_possible_fonts_str = "\n".join(get_fonts())
+ return (image_rgb, image_tensor, list_of_possible_fonts_str,)
@ETK_image_base
@@ -1904,7 +2224,7 @@ class FuncImageStackToImageStack:
my_locals = locals()
image_x = kwargs.get("image", None)
- #gs = image_x.mean(dim=3).repeat(1, 1, 1, 3)
+ # gs = image_x.mean(dim=3).repeat(1, 1, 1, 3)
my_locals["x"] = func
my_locals["x_image"] = kwargs.get("image", None)
@@ -2048,7 +2368,6 @@ class ImageStackAndMatch:
MAX_RESOLUTION = 4096
-# Define the ComfyUI node class for text detection using CRNN model from Torch Hub
@ETK_image_base
class TextInImage:
@@ -2058,38 +2377,51 @@ class TextInImage:
"required": {
"image": ("IMAGE",),
"scale_factor": ("FLOAT", {"default": 1.0, "min": 0.5, "max": 2.0, "step": 0.1}),
- }
+ },
}
RETURN_TYPES = ("STRING",)
- FUNCTION = "detect_text"
+ FUNCTION = "detect_text_batch"
CATEGORY = "OCR"
+ def detect_text_batch(self, image, scale_factor):
+ # calls the detect_text function for each image in the batch
+ results = []
+ for i in range(image.shape[0]):
+ results.append(self.detect_text(image[i].unsqueeze(0), scale_factor)[0])
+
+ results = '\n\n'.join(results)
+ return (results,)
+
def detect_text(self, image, scale_factor):
- import torch
- from torchvision import transforms
- # the incoming tensor in image is in the format (batch, height, width, channels)
- # we need to convert it to (batch, channels, height, width)
+ from .config import config_settings
+ tesse = config_settings.get("tesseract_location", None)
+ if tesse is None:
+ raise Exception("Tesseract location not set in config")
+ import pytesseract
+ import cv2
+ import numpy as np
+ pytesseract.pytesseract.tesseract_cmd = tesse
+
+ # incomming tensor is (B, H, W, C) so convert it
image = image.permute(0, 3, 1, 2)
- new_size = list(image.shape)
- new_size[2] = int(new_size[2] * scale_factor)
- new_size[3] = int(new_size[3] * scale_factor)
- new_size = tuple(new_size[2:])
- rescaled_image = torch.nn.functional.interpolate(image, size=new_size, mode='bilinear')
- # convert to PIL image
+ # Convert the image to a numpy array
+ image = image[0].cpu().numpy()
+ image = np.transpose(image, (1, 2, 0))
+ image = (image * 255).astype(np.uint8)
- pil_image = transforms.ToPILImage()(rescaled_image.squeeze(0))
- # Load CRNN model from Torch Hub
- # Use a pipeline as a high-level helper
+ # Resize the image
+ height, width, _ = image.shape
+ image = cv2.resize(image, (int(width * scale_factor), int(height * scale_factor)))
- from transformers import pipeline
+ # Convert the image to grayscale
+ gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
- pipe = pipeline("image-to-text", model="microsoft/trocr-large-printed")
-
- detected_text = pipe.predict(pil_image)[0]["generated_text"]
+ # Perform OCR on the image
+ detected_text = pytesseract.image_to_string(gray)
return (detected_text,)
@@ -2166,7 +2498,6 @@ class TileImage:
def tile_image(self, image):
import torch
- import torch.nn.functional as F
# image is already B,H,W,C
@@ -2186,8 +2517,6 @@ class TileImage:
return (image,)
-
-
@ETK_image_base
class FloodFillNode:
node_id = "comfyui.flood_fill"
@@ -2208,7 +2537,6 @@ class FloodFillNode:
RETURN_TYPES = ("IMAGE",)
FUNCTION = "execute"
-
def __init__(self):
self.num_workers = None # Can be set to the number of processes to create
@@ -2264,6 +2592,7 @@ class FloodFillNode:
return images
+
@ETK_image_base
class DecorateHeightMap:
"""
@@ -2284,7 +2613,6 @@ class DecorateHeightMap:
RETURN_TYPES = ("IMAGE",)
FUNCTION = "execute"
-
def execute(self, height_map, images, heights, feather=0.05):
"""
Args:
@@ -2311,10 +2639,9 @@ class DecorateHeightMap:
# if there are two heights then it should be length 2
# if there are three heights then it should be length 3
# etc
- if len(feather) != len(heights)-2:
+ if len(feather) != len(heights) - 2:
raise ValueError("feather must be a single float or a list of floats of the same length as heights")
-
# Check for dimension match between height_map and images
if height_map.shape[1:3] != images.shape[1:3]:
raise ValueError("height_map and images must have the same height and width")
@@ -2338,6 +2665,7 @@ class DecorateHeightMap:
output = torch.clamp(output, 0, 1).unsqueeze(0)
return (output,)
+
@ETK_image_base
class RandomImageFromFolder:
"""
@@ -2348,14 +2676,13 @@ class RandomImageFromFolder:
def INPUT_TYPES(cls):
return {
"required": {
- "folder": ("STRING",{"default": "./output"}),
+ "folder": ("STRING", {"default": "./output"}),
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "execute"
-
def execute(self, folder):
"""
Args:
@@ -2382,8 +2709,347 @@ class RandomImageFromFolder:
return (image,)
-import torch
+@ETK_image_base
+class ZoomInImage:
+ """
+ Zoom in on an image by a specified factor and center point
+ """
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "image": ("IMAGE",),
+ "zoom_factor": ("FLOAT", {"default": 2.0, "min": .01, "max": 1000.0, "step": 0.1}),
+ "center_x": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
+ "center_y": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
+ "interpolation": (["nearest", "bilinear", "bicubic", "area"], {"default": "bilinear"}),
+ "keep_size": (["True", "False"], {"default": "True"}),
+ "animate": ("BOOLEAN", {"default": False}),
+ "exponential_zoom": ("BOOLEAN", {"default": False}),
+ }
+ }
+
+ RETURN_TYPES = ("IMAGE",)
+ FUNCTION = "zoom_in"
+
+ CATEGORY = "ETK/image"
+
+ def zoom_in(self, image, zoom_factor, center_x, center_y, interpolation, keep_size, animate, exponential_zoom):
+ batch_size, height, width, channels = image.shape
+ frames = []
+
+ for frame in range(batch_size):
+ percent_done = frame / batch_size
+ if exponential_zoom:
+ percent_done = percent_done ** 2
+
+ # Calculate the zoom factor for this image
+ frame_zoom_factor = 1 + (zoom_factor * percent_done)
+
+ # Calculate the new dimensions after zooming
+ new_height = int(height / frame_zoom_factor)
+ new_width = int(width / frame_zoom_factor)
+
+ # Calculate the coordinates of the top-left corner of the zoomed region
+ top = int((height - new_height) * center_y)
+ left = int((width - new_width) * center_x)
+
+ # Extract the zoomed region from the image
+ zoomed_image = image[frame, top:top + new_height, left:left + new_width, :].unsqueeze(0)
+
+ if keep_size == "True":
+ # change order of dims for interpolation
+ zoomed_image = zoomed_image.permute(0, 3, 1, 2)
+ # Resize the zoomed region back to the original dimensions
+ zoomed_image = torch.nn.functional.interpolate(zoomed_image, size=(height, width), mode=interpolation,
+ align_corners=False)
+ # change order of dims back
+ zoomed_image = zoomed_image.permute(0, 2, 3, 1)
+
+ frames.append(zoomed_image)
+
+ # Stack all frames into a 4D tensor
+ frames = torch.cat(frames, dim=0)
+
+ return (frames,)
+
+
+@ETK_image_base
+class ListUpscaleLatentBy:
+ """uses from nodes import LatentUpscaleBy"""
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ from nodes import LatentUpscaleBy
+ tmp = LatentUpscaleBy.INPUT_TYPES()["required"]["upscale_method"]
+ return {
+ "required": {
+ "latents": ("LIST", {"default": None}),
+ "upscale_method": tmp,
+
+ },
+ "optional": {
+ "factors": ("LIST", {"default": None}),
+ "factor": ("FLOAT", {"default": 1.5, "min": 0.0, }),
+ }
+ }
+
+ RETURN_TYPES = ("LIST",)
+ FUNCTION = "upscale"
+
+ CATEGORY = "latent"
+
+ def upscale(self, latents, upscale_method, factors=None, factor=1.5):
+ from nodes import LatentUpscaleBy
+
+ if not factors:
+ factors = [factor] * len(latents)
+
+ # Check if latents and factors have the same length
+ if len(latents) != len(factors):
+ raise ValueError("Latents and factors must have the same length")
+
+ # Apply LatentUpscaleBy to each pair of latent and factor
+ upscaled_latents = [LatentUpscaleBy().upscale(latent, upscale_method, factor)[0] for latent, factor in
+ zip(latents, factors)]
+
+ return (upscaled_latents,)
+
+
+@ETK_image_base
+class SBSImage:
+ """
+ Combine two images side by side
+ """
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "image1": ("IMAGE",),
+ "image2": ("IMAGE",),
+ }
+ }
+
+ RETURN_TYPES = ("IMAGE",)
+ FUNCTION = "execute"
+
+ CATEGORY = "ETK/image"
+
+ def execute(self, image1, image2):
+ """
+ Args:
+ image1: torch image b,h,w,3
+ image2: torch image b,h,w,3
+ Returns:
+ torch image 1,h,w,3
+ """
+
+ # Check if the images have the same height
+ if image1.shape[1] != image2.shape[1]:
+ raise ValueError("Images must have the same height")
+
+ # Check if the images have the same number of channels
+ if image1.shape[3] != image2.shape[3]:
+ raise ValueError("Images must have the same number of channels")
+
+ # Stack the images horizontally
+ output = torch.cat([image1, image2], dim=2)
+
+ return (output,)
+
+
+@ETK_image_base
+class TTBImage:
+ """
+ Combine two images top to bottom
+ """
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "image1": ("IMAGE",),
+ "image2": ("IMAGE",),
+ }
+ }
+
+ RETURN_TYPES = ("IMAGE",)
+ FUNCTION = "execute"
+
+ CATEGORY = "ETK/image"
+
+ def execute(self, image1, image2):
+ """
+ Args:
+ image1: torch image b,h,w,3
+ image2: torch image b,h,w,3
+ Returns:
+ torch image 1,h,w,3
+ """
+
+ # Check if the images have the same width
+ if image1.shape[2] != image2.shape[2]:
+ raise ValueError("Images must have the same width")
+
+ # Check if the images have the same number of channels
+ if image1.shape[3] != image2.shape[3]:
+ raise ValueError("Images must have the same number of channels")
+
+ # Stack the images vertically
+ output = torch.cat([image1, image2], dim=1)
+
+ return (output,)
+
+
+@ETK_image_base
+class ImageGrid:
+ """
+ Combine multiple images into a grid
+ """
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "image": ("IMAGE",),
+ "rows": ("INT", {"default": 2, "min": 1}),
+ "cols": ("INT", {"default": 2, "min": 1}),
+ }
+ }
+
+ RETURN_TYPES = ("IMAGE",)
+ FUNCTION = "execute"
+
+ CATEGORY = "ETK/image"
+
+ def get_img_idx(self, row, col, rows, cols):
+ return row * cols + col
+
+ def get_xy_for_idx(self, idx, rows, cols):
+ row = idx % cols
+ col = idx // cols
+ return row, col
+
+ def execute(self, image, rows, cols):
+ # TODO: update this to handle a time series.
+ """
+ Args:
+ image: image (r*c),h,w,3
+ rows: int
+ cols: int
+ Returns:
+ torch image 1,h*r,w*c,3
+
+ """
+
+ # Check if the number of images is equal to the number of rows times the number of columns
+ if len(image) != rows * cols:
+ raise ValueError("Number of images must be equal to rows * cols")
+
+ # check if each row has the same hiegth for each image
+ for row in range(rows):
+ row_idxs = [self.get_img_idx(row, col, rows, cols) for col in range(cols)]
+ heights = [image[idx].shape[1] for idx in row_idxs]
+ if not all(height == heights[0] for height in heights):
+ raise ValueError("Images in the same row must have the same height")
+
+ # check if all columns of images have the same width
+ for col in range(cols):
+ col_idxs = [self.get_img_idx(row, col, rows, cols) for row in range(rows)]
+ widths = [image[idx].shape[2] for idx in col_idxs]
+ if not all(width == widths[0] for width in widths):
+ raise ValueError("Images in the same column must have the same width")
+
+ col_widths = []
+ for column in range(cols):
+ idx = self.get_img_idx(0, column, rows, cols)
+ col_width = image[idx].shape[1]
+ col_widths.append(col_width)
+
+ col_heights = []
+ for row in range(rows):
+ idx = self.get_img_idx(row, 0, rows, cols)
+ col_height = image[idx].shape[0]
+ col_heights.append(col_height)
+
+ col_y_pix = [int(sum(col_heights[:i])) for i in range(len(col_heights))]
+ col_x_pix = [int(sum(col_widths[:i])) for i in range(len(col_widths))]
+
+ total_height = 0
+ for row in range(rows):
+ total_height += image[self.get_img_idx(row, 0, rows, cols)].shape[0]
+
+ total_width = 0
+ for col in range(cols):
+ total_width += image[self.get_img_idx(0, col, rows, cols)].shape[1]
+
+ out_tensor = torch.zeros((1, total_height, total_width, image.shape[3]), dtype=image.dtype,
+ device=image.device)
+
+ for idx in range(len(image)):
+ x, y = self.get_xy_for_idx(idx, rows, cols)
+ x_pix = col_x_pix[x]
+ y_pix = col_y_pix[y]
+ h, w, _ = image[idx].shape
+ out_tensor[0, y_pix:y_pix + h, x_pix:x_pix + w, :] = image[idx]
+
+ image.to("cpu")
+ del image
+ torch.cuda.empty_cache()
+
+ return (out_tensor,)
+
+
+@ETK_image_base
+class QueryImage:
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "prompt": ("STRING", {"multiline": True}),
+ "image": ("IMAGE",)
+ }
+ }
+
+ FUNCTION = "execute"
+
+ CATEGORY = "image_to_text"
+ RETURN_TYPES = ("STRING",)
+
+ def execute(self, prompt, image):
+ from .modules.agent_commands import query_image
+ result = query_image.execute({"prompt": prompt, "image": image})
+
+ return (result["result"],)
+
+
+from typing import Tuple
+
+
+@ETK_image_base
+class PreviewImagePassThrough(PreviewImage):
+ @classmethod
+ def INPUT_TYPES(s):
+ # get the input types from the parent class
+ input_types = super().INPUT_TYPES()
+
+ return input_types
+
+ RETURN_TYPES = ("IMAGE",)
+ FUNCTION = "pass_through"
+
+ def pass_through(self, **kwargs):
+ images = kwargs.pop("images", None)
+ # first call the parent class's save_image method
+ result = super().save_images(images, **kwargs)
+ result["result"] = (images,)
+ return result
+
+
+import torch
if __name__ == "__main__":
# test TextRender
diff --git a/modules/image_utils.py b/modules/image_utils.py
new file mode 100644
index 0000000..a9ef174
--- /dev/null
+++ b/modules/image_utils.py
@@ -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()
diff --git a/modules/musicgencli.py b/modules/musicgencli.py
new file mode 100644
index 0000000..20193f1
--- /dev/null
+++ b/modules/musicgencli.py
@@ -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()
diff --git a/modules/noise.py b/modules/noise.py
index afe885d..04b33cb 100644
--- a/modules/noise.py
+++ b/modules/noise.py
@@ -1,3 +1,4 @@
+# TODO: replace pyfastnoisesimd with pyfastnoiselite (works with rpi)
import pyfastnoisesimd as fns
diff --git a/modules/xttscli.py b/modules/xttscli.py
new file mode 100644
index 0000000..062a974
--- /dev/null
+++ b/modules/xttscli.py
@@ -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()
\ No newline at end of file
diff --git a/pytorch_nodes.py b/pytorch_nodes.py
index 8da397f..c8d17f8 100644
--- a/pytorch_nodes.py
+++ b/pytorch_nodes.py
@@ -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
diff --git a/requirements.txt b/requirements.txt
index 3ab8378..2f8e601 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -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
\ No newline at end of file
+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
\ No newline at end of file
diff --git a/server_endpoints.py b/server_endpoints.py
index ec80fcc..e966941 100644
--- a/server_endpoints.py
+++ b/server_endpoints.py
@@ -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"]
diff --git a/tests/test_video.py b/tests/test_video.py
index 65f2f9f..eae27db 100644
--- a/tests/test_video.py
+++ b/tests/test_video.py
@@ -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
diff --git a/video.py b/video.py
index db51dab..f78471f 100644
--- a/video.py
+++ b/video.py
@@ -6,6 +6,9 @@ all image inputs are (B,W,H,C)
avoid numpy and PIL as much as possible
"""
+import shutil
+
+import fractions
import os
from abc import ABC, ABCMeta, abstractmethod
from typing import List, Tuple, Union
@@ -15,9 +18,9 @@ import torchaudio
# from cv2 import VideoWriter, VideoWriter_fourcc, VideoCapture
# from line_profiler import profile
from torch import dtype
-from custom_nodes.EternalKernelLiteGraphNodes.image import torch_image_show
+#from custom_nodes.EternalKernelLiteGraphNodes.image import torch_image_show
from torchvision.models.optical_flow import raft_large, Raft_Large_Weights, raft_small, Raft_Small_Weights
-import imageio_ffmpeg as ffmpeg
+from custom_nodes.EternalKernelLiteGraphNodes.modules.image_utils import lanczos_resize
import os
import torch
from torchvision.utils import draw_keypoints
@@ -27,6 +30,7 @@ import numpy as np
import cv2
import subprocess
import imageio_ffmpeg as ffmpeg
+from comfy_extras.nodes_video_model import SVD_img2vid_Conditioning
NODE_CLASS_MAPPINGS = {} # this is the dictionary that will be used to register the nodes
@@ -64,6 +68,8 @@ def run_ffmpeg_command(args):
Parameters:
args (list): List of command-line arguments for ffmpeg.
"""
+ if isinstance(args, str):
+ args = args.split()
cmd = [ffmpeg.get_ffmpeg_exe()] + args
subprocess.run(cmd)
@@ -551,6 +557,8 @@ class GetImageStackStatistics(ABCABCVideoFolderToImage, metaclass=ABCWidgetMetac
if method == "MSE":
stack_mse_B = torch.mean(image_stack_BCWH_DIFF ** 2, dim=(1, 2, 3))
+ # find the max possible value of the mean squared error
+
else:
# not implimented
raise NotImplementedError
@@ -561,13 +569,13 @@ class GetImageStackStatistics(ABCABCVideoFolderToImage, metaclass=ABCWidgetMetac
mean = torch.mean(stack_mse_B)
std = torch.std(stack_mse_B)
median = torch.median(stack_mse_B)
- max = torch.max(stack_mse_B)
- min = torch.min(stack_mse_B)
+ max_mse = torch.max(stack_mse_B)
+ min_mse = torch.min(stack_mse_B)
- out_str = f"mean: {mean}\n std: {std}\n median: {median}\n max: {max}\n min: {min}\n len: {image_stack_diff_len}"
+ out_str = f"mean: {mean}\n std: {std}\n median: {median}\n max: {max_mse}\n min: {min_mse}\n len: {image_stack_diff_len}"
# create the histogram
- hist_values, _ = np.histogram(stack_mse_B.numpy(), bins=100)
+ hist_values, _ = np.histogram(stack_mse_B.numpy(), bins=min(100, image_stack_diff_len))
hist_values = hist_values / hist_values.max() # Normalize to [0, 1]
# Create an image tensor with shape (1, 512, required_width, 3) filled with zeros
@@ -579,18 +587,21 @@ class GetImageStackStatistics(ABCABCVideoFolderToImage, metaclass=ABCWidgetMetac
# Fill the image tensor with the histogram
for i, value in enumerate(hist_values):
- height = int(value * scale_height)
+ height = 512 - int(value * scale_height)
start_x = int(i * scale_width)
end_x = int((i + 1) * scale_width)
out_hist_image[0, height:, start_x:end_x - 1, :] = 1 # Color as white
+ # repeat each element of the stack_mse_B tensor by scale_width
+ stack_mse_B = stack_mse_B.repeat_interleave(scale_width)
+
# Create a plot for the difference for each frame pair
plot_offset = 512
for i, value in enumerate(stack_mse_B):
- height = int(value * scale_height / max) # Normalize by the max value
+ height = 512 - int(value * (scale_height / .5)) # Normalize by the max value
start_x = plot_offset + i
end_x = plot_offset + i + 1
- out_hist_image[0, height:, start_x:end_x, :] = 1 # Color as white
+ out_hist_image[0, height:, start_x:end_x, :] = (1 * height) / 512 # Color as white
return (out_str, out_hist_image,)
@@ -674,7 +685,7 @@ class ImageStackToVideoFramesFolder(metaclass=ABCWidgetMetaclass):
},
}
- RETURN_TYPES = ()
+ RETURN_TYPES = ("STRING",)
CATEGORY = "video"
OUTPUT_NODE = True
FUNCTION = "handler"
@@ -717,7 +728,7 @@ class ImageStackToVideoFramesFolder(metaclass=ABCWidgetMetaclass):
image_stack = None
del image_stack
- return ()
+ return (folder_out,)
class TemporalSpatialSmoothing(metaclass=ABCWidgetMetaclass):
@@ -1126,6 +1137,17 @@ class CombineAudioAndVideoFiles(metaclass=ABCWidgetMetaclass):
OUTPUT_NODE = True
FUNCTION = "handler"
+ @classmethod
+ def has_audio_stream(cls, video_file):
+ """Check if the video file has an audio stream."""
+ command = ['ffprobe', '-v', 'error', '-select_streams', 'a', '-show_entries', 'stream=codec_name', '-of',
+ 'default=nw=1:nk=1', video_file]
+ try:
+ output = subprocess.check_output(command, stderr=subprocess.STDOUT)
+ return bool(output.strip())
+ except subprocess.CalledProcessError:
+ return False
+
def handler(self, **kwargs):
from os.path import splitext, join
import shutil
@@ -1141,15 +1163,33 @@ class CombineAudioAndVideoFiles(metaclass=ABCWidgetMetaclass):
raise RuntimeError("FFmpeg could not be found.")
for video_file, audio_file in video_audio_pairs:
+
video_out_file_fp = splitext(video_file)[0] + "_v_and_a.mp4"
input_args = ['-i', video_file, '-i', audio_file]
- output_args = ['-c:v', 'copy', '-c:a', 'aac', '-strict', 'experimental', '-y', video_out_file_fp]
- command = [ffmpeg_exe] + input_args + output_args
- os.system(' '.join(command))
+
+ if self.has_audio_stream(video_file):
+ # The video has an audio stream, proceed with mixing
+ filter_complex_arg = ['-filter_complex', '[0:a][1:a]amix=inputs=2:duration=longest[a]']
+ output_args = ['-c:v', 'copy', '-map', '0:v:0', '-map', '[a]', '-c:a', 'aac', '-strict', 'experimental',
+ '-y', video_out_file_fp]
+
+
+ else:
+ filter_complex_arg = []
+ output_args = ['-c:v', 'copy', '-map', '0:v:0', '-map', '1:a:0', '-c:a', 'aac', '-strict',
+ 'experimental', '-y', video_out_file_fp]
+
+ command = [ffmpeg_exe] + input_args + filter_complex_arg + output_args
+ cmd = ' '.join(command)
+ print(cmd)
+ os.system(cmd)
if move_to_folder:
destination_path = join(move_to_folder, os.path.basename(video_out_file_fp))
- shutil.move(video_out_file_fp, destination_path)
+
+ if video_out_file_fp != destination_path:
+ shutil.move(video_out_file_fp, destination_path)
+
combined_video_paths.append(destination_path)
else:
combined_video_paths.append(video_out_file_fp)
@@ -1168,7 +1208,7 @@ class CombineFoldersWithFramesToVideoFiles(metaclass=ABCWidgetMetaclass):
"fps/length": ("FLOAT", {"default": 30, "min": 1, "max": 10000000000}),
"bitrate": ("STRING", {"default": "20000k"}),
"codec": (
- ["libx264", "libx265", "mpeg4", "vp9", "huffyuv", "flv1"],
+ ["libx264", "libx265", "mpeg4", "vp9", "huffyuv", "flv1", "h264_nvenc", "hevc_nvenc"],
{"default": "libx264"}
),
"file_extension": (
@@ -1176,7 +1216,10 @@ class CombineFoldersWithFramesToVideoFiles(metaclass=ABCWidgetMetaclass):
{"default": ".mp4"}
),
"fps or length": (["fps", "length"], {"default": "fps"}),
- }
+ },
+ "optional": {
+ "[duration or fps]": ("LIST", {"default": []}),
+ },
}
RETURN_TYPES = ("LIST", "STRING",)
@@ -1195,14 +1238,15 @@ class CombineFoldersWithFramesToVideoFiles(metaclass=ABCWidgetMetaclass):
codec = kwargs["codec"]
file_extension = kwargs["file_extension"]
fps_or_length = kwargs["fps or length"]
+ duration_or_fps_list = kwargs.get("[duration or fps]", [])
combined_video_paths = []
messages = []
- for folder_path in folders_with_frames:
+ for i, folder_path in enumerate(folders_with_frames):
try:
out_file = os.path.join(os.path.dirname(folder_path),
- f"{os.path.basename(folder_path)}{file_extension}")
+ f"{os.path.basename(folder_path)}tmp_{i}{file_extension}")
frame_files = sorted(
[f for f in os.listdir(folder_path) if f.lower().endswith(('.jpg', '.jpeg', '.png'))])
@@ -1212,9 +1256,15 @@ class CombineFoldersWithFramesToVideoFiles(metaclass=ABCWidgetMetaclass):
if fps_or_length == "fps":
fps = fps_length_value
+ if duration_or_fps_list:
+ fps = duration_or_fps_list[i]
+
else:
total_frames = len(frame_files)
fps = total_frames / fps_length_value if fps_length_value != 0 else 30
+ if duration_or_fps_list:
+ # calculate the fps based on the duration
+ fps = total_frames / duration_or_fps_list[i]
writer = imageio.get_writer(out_file, fps=fps, codec=codec, bitrate=bitrate)
@@ -1270,7 +1320,7 @@ class ImageStackToVideoFile(metaclass=ABCWidgetMetaclass):
import cv2
import torchvision
- use_torchvision = True
+ use_torchvision = False
# video_dir = kwargs["video_dir"]
@@ -1319,6 +1369,9 @@ class ImageStackToVideoFile(metaclass=ABCWidgetMetaclass):
if use_torchvision == False:
writer = imageio.get_writer(video_out_path, fps=float(video_fps))
+ # convert the image stack to uint8
+ image_stack = (image_stack * 255).to(torch.uint8)
+
for frame_data in image_stack:
frame = np.array(frame_data) # Convert torch tensor to numpy array if needed
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
@@ -1352,12 +1405,15 @@ class ImageStackToVideoFile(metaclass=ABCWidgetMetaclass):
if use_torchvision:
audio_tensor = None
if audio_file is not None:
- aud_sr = torchaudio.io.StreamReader(audio_file)
-
+ audio_tensor, aud_sr = torchaudio.load(audio_file)
+ fps_frac = fractions.Fraction(float(video_fps))
torchvision.io.write_video(filename=video_out_path,
- video_array=image_stack,
- fps=video_fps,
+ video_array=image_stack * 255,
+ fps=fps_frac,
audio_array=audio_tensor,
+ video_codec=codec,
+ audio_codec="aac",
+ audio_fps=aud_sr
)
return (self, video_out_path,)
@@ -1919,6 +1975,7 @@ class PanByStep(metaclass=ABCWidgetMetaclass):
ret = {"required": {"ImageStack": ("IMAGE", {"default": None}),
"pan x": ("INT", {"default": 0, "min": -2048}),
"pan y": ("INT", {"default": 0, "min": -2048}),
+ "wrap": ("BOOLEAN", {"default": False}),
},
}
return ret
@@ -1938,12 +1995,14 @@ class PanByStep(metaclass=ABCWidgetMetaclass):
x = kwargs["pan x"]
y = kwargs["pan y"]
+ wrap = kwargs["wrap"]
+
time_index, H, W, C = x_image.shape # Assuming x_image has shape (time_index, H, W, C)
out_list = []
stop_x = x
stop_y = y
- dx_step = stop_x / time_index - 1
+ dx_step = 0 if x == 0 else stop_x / time_index - 1
dy_step = stop_y / time_index - 1
for t in range(time_index):
@@ -1954,8 +2013,17 @@ class PanByStep(metaclass=ABCWidgetMetaclass):
# Change from (H, W, C) to (C, H, W)
current_frame = current_frame.permute(2, 0, 1)
- transformed_frame = TF.affine(current_frame, angle=0, translate=(tx, ty), scale=1,
- shear=[0, 0]).detach().clone()
+ # transformed_frame = TF.affine(current_frame, angle=0, translate=(tx, ty), scale=1,
+ # shear=[0, 0]).detach().clone()
+
+ # If wrap is True, apply the wrapping operation
+ if wrap:
+ transformed_frame = current_frame.roll(
+ (int(tx) % W, int(ty) % H), dims=(2, 1)
+ )
+ else:
+ transformed_frame = TF.affine(current_frame, angle=0, translate=(tx, ty), scale=1,
+ shear=[0, 0]).detach().clone()
# Convert back to (H, W, C)
transformed_frame = transformed_frame.permute(1, 2, 0)
@@ -1964,3 +2032,1097 @@ class PanByStep(metaclass=ABCWidgetMetaclass):
x_image[t] = transformed_frame
return (x_image.detach().clone(),)
+
+
+class VideoFolderProvider(metaclass=ABCWidgetMetaclass):
+ """
+ given a folder path, creates the folder or empties it entirely returns the folder path as
+ a string but as the node type "VIDEOFOLDER"
+ """
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ ret = {
+ "required": {
+ "folder_path": ("STRING", {"default": None}),
+ },
+ "optional": {
+ "activator:": ("*", {"default": None})
+ }
+ }
+
+ return ret
+
+ RETURN_TYPES = ("VIDEOFOLDER",)
+ FUNCTION = "video_folder_provider"
+
+ def video_folder_provider(self, **kwargs):
+ import shutil
+ import os
+ # check if the folder_path is none
+ if kwargs["folder_path"] is None:
+ raise ValueError("folder_path must be specified")
+
+ # set folder exists flag by using the os.path.exists function
+ folder_exists = os.path.exists(kwargs["folder_path"])
+
+ if not folder_exists:
+ os.makedirs(kwargs["folder_path"])
+ folder_exists = True
+
+ # set files in folder flag by using the os.listdir function
+ files_in_folder = os.listdir(kwargs["folder_path"])
+
+ # set accessible flag by using the os.access function
+ accessible = os.access(kwargs["folder_path"], os.W_OK)
+ if not accessible:
+ raise ValueError("Folder for video files is not accessible")
+
+ # if the folder exists, delete all files in the folder
+ if folder_exists:
+ for file in files_in_folder:
+ # use shutil to remove the file
+ fl = os.path.join(kwargs["folder_path"], file)
+ if os.path.isfile(fl):
+ os.remove(fl)
+ # use shutil to remove the entire path
+ check_this_path = file
+ if len(check_this_path) > 3:
+ shutil.rmtree(check_this_path, ignore_errors=True)
+ else:
+ raise ValueError("Folder path is too short")
+
+ # if the folder does not exist, create the folder
+
+ return (kwargs["folder_path"],)
+
+
+class AudioFolderProvider(metaclass=ABCWidgetMetaclass):
+ """
+ given a folder path, creates the folder or empties it entirely returns the folder path as
+ a string but as the node type "AUDIOFOLDER"
+ """
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ ret = {"required": {"folder_path": ("STRING", {"default": None}),
+ },
+ }
+ return ret
+
+ RETURN_TYPES = ("AUDIOFOLDER",)
+ FUNCTION = "audio_folder_provider"
+
+ def audio_folder_provider(self, **kwargs):
+ import shutil
+ import os
+ # check if the folder_path is none
+ if kwargs["folder_path"] is None:
+ raise ValueError("folder_path must be specified")
+
+ # set folder exists flag by using the os.path.exists function
+ folder_exists = os.path.exists(kwargs["folder_path"])
+
+ if not folder_exists:
+ os.makedirs(kwargs["folder_path"])
+ folder_exists = True
+
+ # set files in folder flag by using the os.listdir function
+ files_in_folder = os.listdir(kwargs["folder_path"])
+
+ # set accessible flag by using the os.access function
+ accessible = os.access(kwargs["folder_path"], os.W_OK)
+ if not accessible:
+ raise ValueError("Folder for audio files is not accessible")
+
+ # if the folder exists, delete all files in the folder
+ if folder_exists:
+ for file in files_in_folder:
+ # use shutil to remove the file
+ fl = os.path.join(kwargs["folder_path"], file)
+ if os.path.isfile(fl):
+ os.remove(fl)
+
+ # if the folder does not exist, create the folder
+
+ return (kwargs["folder_path"],)
+
+
+class VideoDefinitionProvider(metaclass=ABCWidgetMetaclass):
+ """
+ given a VIDEOFOLDER node, and inputs for height, width, and fps, creates a VIDEOFOLDERDEF variable
+ """
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ ret = {"required": {"video_folder": ("VIDEOFOLDER", {"default": None}),
+ "height": ("INT", {"default": 1080}),
+ "width": ("INT", {"default": 1920}),
+ "fps": ("FLOAT", {"default": 30}),
+ "frame format": (["JPG", "PNG"], {"default": "JPG"}),
+ },
+ }
+ return ret
+
+ RETURN_TYPES = ("VIDEO_FOLDER_DEF",)
+ FUNCTION = "video_definition_provider"
+
+ def video_definition_provider(self, **kwargs):
+ from addict import Dict as adict
+ ret = adict()
+ ret["video_folder"] = kwargs["video_folder"]
+ ret["height"] = kwargs["height"]
+ ret["width"] = kwargs["width"]
+ ret["fps"] = kwargs["fps"]
+ ret["frame format"] = kwargs["frame format"]
+ # a function that gets the frame count from the disk
+ ret.get_frame_count = lambda: len([f for f in os.listdir(ret["video_folder"]) if "frame_" in f])
+ # a function that gets the next file name and full path dont forget to add 10 zeros to the frame count
+ ret.get_next_file_name = lambda x=0: os.path.join(ret['video_folder'],
+ f"frame_{str(ret.get_frame_count()+x).zfill(10)}.{ret['frame format'].lower()}")
+ ret.get_last_file_name = lambda offset=0: os.path.join(ret['video_folder'],
+ f"frame_{str(max(ret.get_frame_count() - offset - 1, 0)).zfill(10)}.{ret['frame format'].lower()}")
+ ret.get_current_file_name = lambda: os.path.join(ret['video_folder'], f"frame_{str(ret.get_frame_count() - 1).zfill(10)}.{ret['frame format'].lower()}")
+
+ def reserve_next_frame():
+ fn = ret.get_next_file_name()
+ open(fn, "w").close()
+ return fn
+
+ ret.reserve_next_frame = reserve_next_frame
+
+ return (ret,)
+
+
+class AudioDefinitionProvider(metaclass=ABCWidgetMetaclass):
+ """
+ given a AUDIOFOLDER node, and inputs for sample rate, and channels, creates a AUDIOFOLDERDEF variable
+ """
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ ret = {"required": {"audio_folder": ("AUDIOFOLDER", {"default": None}),
+ "sample rate": ("INT", {"default": 44100}),
+ "channels": ("INT", {"default": 2}),
+ "audio format": (["WAV", "MP3"], {"default": "WAV"}),
+ },
+ }
+ return ret
+
+ RETURN_TYPES = ("AUDIO_FOLDER_DEF",)
+ FUNCTION = "audio_definition_provider"
+
+ def get_num_audio_files_that_match(self, audio_folder):
+ audio_files = self.get_all_audio_files(audio_folder)
+ return len(audio_files)
+
+ def get_all_audio_files(self, audio_folder):
+ import os
+ audio_files = os.listdir(audio_folder)
+ audio_files = [os.path.join(audio_folder, x) for x in audio_files if "audio_" in x]
+ return audio_files
+
+ def audio_definition_provider(self, **kwargs):
+ from addict import Dict as adict
+ ret = adict()
+ ret["audio_folder"] = kwargs["audio_folder"]
+ ret["sample rate"] = kwargs["sample rate"]
+ ret["channels"] = kwargs["channels"]
+ ret["audio format"] = kwargs["audio format"]
+ # a function that gets the frame count from the disk
+ ret.get_file_count = lambda: self.get_num_audio_files_that_match(ret["audio_folder"])
+ # a function that gets the next file name and full path dont forget to zfill
+ ret.get_next_file_name = lambda: f"{ret['audio_folder']}{os.sep}audio_{str(ret.get_file_count()).zfill(10)}.{ret['audio format'].lower()}"
+ ret.get_last_file_name = lambda: f"{ret['audio_folder']}{os.sep}audio_{str(ret.get_file_count() - 1).zfill(10)}.{ret['audio format'].lower()}"
+ ret.get_all_audio_files = lambda: self.get_all_audio_files(ret["audio_folder"])
+
+ return (ret,)
+
+
+class AddAudioToAudioFolder(metaclass=ABCWidgetMetaclass):
+ """
+ given a AUDIO_FOLDER_DEF node, and a list of audio files, adds the audio files to the audio folder
+ """
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ ret = {"required": {"audio_folder_def": ("AUDIO_FOLDER_DEF", {"default": None}),
+ "AUDIO": ("LIST", {"default": None}),
+ "delete after": ("BOOLEAN", {"default": False}),
+ },
+ }
+ return ret
+
+ RETURN_TYPES = ("AUDIO_FOLDER_DEF",)
+ FUNCTION = "add_audio_to_audio_folder"
+
+ def add_audio_to_audio_folder(self, **kwargs):
+ # check each file in the list to make sure we can convert it to the correct format
+ import shutil
+ import os
+
+ ffmpeg_exe = ffmpeg.get_ffmpeg_exe()
+
+ audio_folder_def = kwargs["audio_folder_def"]
+ audio_files = kwargs["AUDIO"]
+
+ delete = kwargs["delete after"]
+
+ for audio_file in audio_files:
+ # use ffmpeg_exe to convert the audio file to the correct format which is audio_folder_def["audio format"]
+ # use the audio_folder_def["audio format"] to determine the output file name
+ # use the audio_folder_def["audio folder"] to determine the output folder
+ # use the audio_folder_def["sample rate"] to determine the sample rate
+ # use the audio_folder_def["channels"] to determine the number of channels
+
+ # get the output file name
+ output_file_name = audio_folder_def.get_next_file_name()
+ input_file_name = audio_file
+
+ CMD = f"{ffmpeg_exe} -y -i {input_file_name} -ar {audio_folder_def['sample rate']} -ac {audio_folder_def['channels']} {output_file_name}"
+
+ os.system(CMD)
+ if delete:
+ os.remove(input_file_name)
+
+ return (audio_folder_def,)
+
+
+class AddSilenceToAudioFolder(metaclass=ABCWidgetMetaclass):
+ """create and add a new audio file that is silence for the specified duration"""
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ ret = {"required": {"audio_folder_def": ("AUDIO_FOLDER_DEF", {"default": None}),
+ "duration": ("FLOAT", {"default": 1.0}),
+ },
+ }
+ return ret
+
+ RETURN_TYPES = ("AUDIO_FOLDER_DEF",)
+ FUNCTION = "add_silence_to_audio_folder"
+
+ def add_silence_to_audio_folder(self, **kwargs):
+ # this is to be multi-platform
+ import shutil
+ import os
+
+ audio_folder_def = kwargs["audio_folder_def"]
+ duration = kwargs["duration"]
+
+ # get the output file name
+ output_file_name = audio_folder_def.get_next_file_name()
+ sample_rate = audio_folder_def["sample rate"]
+ channels = audio_folder_def["channels"]
+ fmt = audio_folder_def["audio format"]
+
+ # create the silence audio file
+
+ if fmt == "WAV":
+ fmt_part = "pcm_s16le"
+ elif fmt == "MP3":
+ fmt_part = "mp3"
+ else:
+ fmt_part = "pcm_s16le"
+
+ ffmpeg_exe = ffmpeg.get_ffmpeg_exe()
+ # "-i anullsrc" is the source of the silence this will work on all platforms
+ # "-c:a" is the codec for the audio file
+ # "-f" is the format of the audio file, "lavfi" is equivalent to "ffmpeg filter"
+ # "-i anullsrc=r=44100:cl=2:d=1" is the source of the silence
+
+ # CMD = f"{ffmpeg_exe} -f lavfi -i aevalsrc=0:d={duration}:cl={channels}:s={sample_rate} -c:a {fmt_part} {output_file_name}"
+ # CMD = f"{ffmpeg_exe} -f lavfi -i anullsrc=channel_layout={channels}:sample_rate={sample_rate}:duration={duration} -c:a {fmt_part} {output_file_name}"
+ CMD = f"{ffmpeg_exe} -f lavfi -i aevalsrc=0:sample_rate={sample_rate}:channel_layout={channels},atrim=0:{duration} -c:a {fmt_part} {output_file_name}"
+ # now we call it in a way to that we will know if there is an error
+
+ import subprocess
+ # the plan is to call the command and wait for it to finish
+ process = subprocess.Popen(CMD, shell=True, stderr=subprocess.PIPE)
+ process.wait()
+ # check if the process was successful
+ if process.returncode != 0:
+ # if the process was not successful raise an exception
+ raise Exception(f"ffmpeg failed with error code {process.returncode}")
+
+ return (audio_folder_def,)
+
+
+class GetAudioFilesLengthsFromAudioFolder(metaclass=ABCWidgetMetaclass):
+ """Get the lengths of all audio files in a given audio folder"""
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {"required": {"audio_folder_def": ("AUDIO_FOLDER_DEF", {"default": None})}}
+
+ RETURN_TYPES = ("AUDIO_FOLDER_DEF", "LIST",)
+ RETURN_NAMES = ("audio_folder_def", "[lens(seconds)]",)
+ FUNCTION = "get_audio_files_lengths"
+
+ def get_audio_files_lengths(self, audio_folder_def):
+ import os
+ import torchaudio
+
+ audio_folder = audio_folder_def["audio_folder"]
+ audio_files = os.listdir(audio_folder)
+ audio_files_lengths = []
+
+ for audio_file in audio_files:
+ audio_path = os.path.join(audio_folder, audio_file)
+ waveform, sample_rate = torchaudio.load(audio_path)
+ duration = waveform.shape[1] / sample_rate
+ audio_files_lengths.append(duration)
+
+ return (audio_folder_def, audio_files_lengths,)
+
+
+class AudioFolderRenderAudio(metaclass=ABCWidgetMetaclass):
+ """ Joins all the audio files """
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "audio_folder_def": ("AUDIO_FOLDER_DEF", {"default": None}),
+ "output_path": ("STRING", {"multiline": False, "default": ""}),
+ }
+ }
+
+ RETURN_TYPES = ("STRING",)
+ RETURN_NAMES = ("output_path",)
+ FUNCTION = 'join_audio_files'
+ CATEGORY = 'ETK/audio'
+
+ def join_audio_files(self, **kwargs):
+ import librosa
+ import copy
+ import soundfile as sf
+
+ audio_folder_def = kwargs['audio_folder_def']
+ kwargs = copy.deepcopy(kwargs)
+ paths = audio_folder_def.get_all_audio_files()
+ output_path = kwargs['output_path']
+
+ # Read the first file to get metadata
+ audio_data_list = []
+ for path in paths:
+ audio_data, sr = librosa.load(path, sr=None)
+ audio_data_list.append(audio_data)
+
+ # Concatenate all audio data
+ joined_audio_data = np.concatenate(audio_data_list)
+
+ # TODO: need to implement writing in the correct format
+ # Save the concatenated audio
+ sf.write(output_path, joined_audio_data, sr)
+
+ return (output_path,)
+
+
+import subprocess
+import platform
+
+import subprocess
+import platform
+import os
+
+
+def duplicate_file(source_file, target_files):
+ """
+ Duplicates a given source file into multiple target files using a single CLI command
+ in a non-blocking way. The base directory is automatically determined from the
+ source file's location.
+
+ Args:
+ - source_file (str): Path to the source file to be duplicated.
+ - target_files (list of str): A list of target file paths where the source file should be duplicated.
+
+ Returns:
+ - None
+ """
+ # Determine the base directory from the source file's directory
+ base_dir = os.path.dirname(os.path.abspath(source_file))
+
+ # Change the current working directory to the base directory
+ os.chdir(base_dir)
+
+ # Get the source file name (relative to the base directory)
+ source_file_name = os.path.basename(source_file)
+
+ # get the relative paths of the target files
+ target_files = [os.path.relpath(target_file, base_dir) for target_file in target_files]
+
+ # Determine the shell command based on the operating system
+ copy_command = 'copy' if os.name == 'nt' else 'cp'
+
+ # Construct the command string
+ command_parts = [f'{copy_command} "{source_file_name}" "{target}"' for target in target_files]
+ command = ' && '.join(command_parts)
+
+ command = command.replace("\"", "")
+ command = command.replace("\'", "")
+ # print(len(command))
+
+ try:
+ # Execute the command in a non-blocking way
+ process = subprocess.Popen(command, shell=True, stderr=subprocess.DEVNULL)
+ process.wait()
+
+ except subprocess.CalledProcessError as e:
+ pass
+ except ValueError as e:
+ pass
+ except Exception as e:
+ pass
+
+ if "process" not in locals():
+ pass
+ if "process" not in locals() or process.returncode != 0:
+ # If the command is too long, split it into multiple commands
+ # split the command into commands that do 64 at a time
+ command_parts = [command_parts[i:i + 64] for i in range(0, len(command_parts), 64)]
+ processes = []
+ for command_part in command_parts:
+ command = ' && '.join(command_part)
+ process = subprocess.Popen(command, shell=True, stderr=subprocess.DEVNULL)
+ processes.append(process)
+
+ for process in processes:
+ process.wait()
+
+ return process
+
+
+class AddFramesToVideoFolder(metaclass=ABCWidgetMetaclass):
+ """
+ given a VIDEO_FOLDER_DEF node, and a list of frames, adds the frames to the video folder
+ """
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ ret = {"required": {"video_folder_def": ("VIDEO_FOLDER_DEF", {"default": None}),
+ "IMAGE": ("IMAGE", {"default": None}),
+ "Rescale method": (
+ ["NEAREST", "BOX", "BILINEAR", "HAMMING", "BICUBIC", "LANCZOS"],
+ {"default": "LANCZOS"}),
+ "Release vram?": ("BOOLEAN", {"default": True}),
+ },
+ }
+ return ret
+
+ RETURN_TYPES = ("VIDEO_FOLDER_DEF",)
+ FUNCTION = "add_frames_to_video_folder"
+
+ @staticmethod
+ def save_or_copy_frame(current_frame, previous_frame, frame_fp, previous_frame_fp):
+ from torchvision.utils import save_image
+ import shutil
+ from PIL import Image
+ if previous_frame is None or not torch.equal(current_frame, previous_frame):
+ # If the current frame is different from the previous frame, save it to disk
+ current_frame = current_frame
+
+ if current_frame.dtype != torch.uint8: # handle float type frames
+ ndarr = current_frame.permute(2, 0, 1).mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu",
+ torch.uint8).numpy()
+ else: # handle uint8 frames
+ ndarr = current_frame.permute(2, 0, 1).permute(1, 2, 0).numpy()
+ im = Image.fromarray(ndarr)
+ im.save(frame_fp)
+
+ previous_frame = current_frame
+ else:
+ # If the current frame is the same as the previous frame, copy the previously saved frame file
+ shutil.copy2(previous_frame_fp, frame_fp)
+ return previous_frame
+
+ @staticmethod
+ def process_frames(frames, rescale_method, video_folder_def):
+ import gc
+ num_frames = frames.shape[0]
+ previous_frame = None
+ frame_fp = None
+ frames_area_all_the_same = torch.all(frames[0] == frames)
+ if frames_area_all_the_same and frames.shape[0] > 1:
+ # get a list of file names using reservenexframe
+ frame_f_names = [video_folder_def.reserve_next_frame() for _ in range(num_frames)]
+ # save the first frame
+ AddFramesToVideoFolder.save_or_copy_frame(frames[0], previous_frame, frame_f_names[0], None)
+ # copy the first frame to the rest of the frames using os by calling the cli
+ duplicate_file(frame_f_names[0], frame_f_names[1:])
+
+ else:
+ for i in range(num_frames):
+ next_fn = video_folder_def.get_next_file_name()
+
+ previous_frame_fp = frame_fp
+ frame_fp = next_fn
+
+ current_frame = frames[i]
+
+ previous_frame = AddFramesToVideoFolder.save_or_copy_frame(current_frame, previous_frame, frame_fp,
+ previous_frame_fp)
+
+ del frames
+ gc.collect()
+ return None
+
+ def add_frames_to_video_folder(self, **kwargs):
+ from glob import glob
+ from PIL import Image
+ import os
+ import torch
+ import gc
+ from torchvision.transforms.functional import resize
+ import psutil
+
+ initial_memory = psutil.virtual_memory().available
+ print(f"Initial free memory: {initial_memory / (1024 * 1024):.2f} MB")
+
+ import shutil
+
+ # get the video folder path
+ video_folder = kwargs["video_folder_def"]["video_folder"]
+
+ # get the frame format
+ frame_format = kwargs["video_folder_def"]["frame format"]
+
+ # get the height
+ height = kwargs["video_folder_def"]["height"]
+
+ # get the width
+ width = kwargs["video_folder_def"]["width"]
+
+ # get the fps
+ fps = kwargs["video_folder_def"]["fps"]
+
+ # get the frames
+ frames = kwargs["IMAGE"]
+
+ # get the rescale method
+ rescale_method = kwargs.get("Rescale method", None)
+ if not rescale_method:
+ rescale_method = kwargs.get("Rescale_method", None)
+
+ free_vram = kwargs.get("Release vram?", None)
+ if not free_vram:
+ free_vram = kwargs.get("Release_vram_", True)
+
+ orig_frames = frames
+ # num_files = len(glob(os.path.join(video_folder, f"frame_*.{frame_format.lower()}")))
+
+ # rescale if needed to the correct size
+ if height != frames.shape[1] or width != frames.shape[2]:
+ # first since transforms.Resize requires a tensor of shape B,C, H, W
+ # we need to permute the tensor to B, C, H, W
+ frames = frames.permute(0, 3, 1, 2)
+ frames = frames.to("cuda")
+
+ # get the right value for interpolation
+ interpolation_methods = {
+ "NEAREST": Image.Resampling.NEAREST,
+ "BOX": Image.Resampling.BOX,
+ "BILINEAR": Image.Resampling.BILINEAR,
+ "HAMMING": Image.Resampling.HAMMING,
+ "BICUBIC": Image.Resampling.BICUBIC,
+ "LANCZOS": Image.Resampling.LANCZOS
+ }
+
+ rescale_type = interpolation_methods[rescale_method]
+ # rescale_method = transforms.Resize((height, width), interpolation=rescale_type) # this doesnt work so well lets use a different function
+
+ if rescale_type != Image.Resampling.LANCZOS:
+ rescale_method = lambda x: resize(x, (height, width), interpolation=rescale_type)
+ frames = rescale_method(frames)
+ else:
+ rescale_method = lambda x: lanczos_resize(x, (width, height))
+ frames = rescale_method(frames)
+
+ # now we need to permute back to B, H, W, C
+ frames = frames.permute(0, 2, 3, 1)
+ self.process_frames(frames, rescale_method, kwargs["video_folder_def"])
+ else:
+ self.process_frames(frames, rescale_method, kwargs["video_folder_def"])
+
+ gc.collect()
+
+ if free_vram:
+ import gc
+ if isinstance(frames, torch.Tensor):
+ frames.to("cpu").detach().numpy()
+ del frames
+
+ if isinstance(orig_frames, torch.Tensor):
+ orig_frames = orig_frames.to("cpu").detach().numpy()
+ del orig_frames
+
+ torch.cuda.empty_cache()
+ gc.collect()
+ torch.cuda.empty_cache()
+ gc.collect()
+
+ initial_memory = psutil.virtual_memory().available
+ print(f"after adding frames free memory: {initial_memory / (1024 * 1024):.2f} MB")
+ # return the VIDEO_FOLDER_DEF
+ return (kwargs["video_folder_def"],)
+
+
+class RepeatFrames(metaclass=ABCWidgetMetaclass):
+ """
+ does nothing more than repeat the frames a number of times
+ """
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ ret = {"required": {"IMAGE": ("IMAGE", {"default": None}),
+ "repeat": ("INT", {"default": 1}),
+ },
+ }
+ return ret
+
+ RETURN_TYPES = ("IMAGE",)
+ RETURN_NAMES = ("repeated images",)
+ FUNCTION = "repeat_frames"
+
+ def repeat_frames(self, **kwargs):
+ import torch
+ frames = kwargs["IMAGE"]
+ repeat = kwargs["repeat"]
+ return (frames.repeat(repeat, 1, 1, 1),)
+
+
+class HoldFrameForSecsInVideoFolder(metaclass=ABCWidgetMetaclass):
+ """
+ given a VIDEO_FOLDER_DEF node, and a list of frames, first uses RepeatFrames to repeat the frames
+ for the correct amount and then adds the frames to the video folder using AddFramesToVideoFolder
+ """
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ ret = {"required": {"video_folder_def": ("VIDEO_FOLDER_DEF", {"default": None}),
+ "IMAGE": ("IMAGE", {"default": None}),
+ "seconds": ("FLOAT", {"default": 1.0}),
+ "Rescale method": (
+ ["NEAREST", "BOX", "BILINEAR", "HAMMING", "BICUBIC", "LANCZOS"],
+ {"default": "LANCZOS"}),
+ "Release vram?": ("BOOLEAN", {"default": True}),
+ },
+ }
+ return ret
+
+ RETURN_TYPES = ("VIDEO_FOLDER_DEF",)
+ FUNCTION = "hold_frame_for_secs_in_video_folder"
+
+ def saver(self, x, kwargs):
+ # this just saves in an async way
+ rescale_method = kwargs.get("Rescale method", None)
+ if rescale_method is None:
+ rescale_method = kwargs.get("Rescale_method", None)
+
+ rel_vram = kwargs.get("Release vram?", None)
+ if rel_vram is None:
+ rel_vram = kwargs.get("Release_vram_", True)
+
+ AddFramesToVideoFolder().add_frames_to_video_folder(
+ video_folder_def=kwargs["video_folder_def"],
+ IMAGE=x,
+ Rescale_method=rescale_method,
+ Release_vram_=rel_vram
+ )
+
+ def hold_frame_for_secs_in_video_folder(self, **kwargs):
+ import torch
+ import asyncio
+
+ vide_folder_ref = kwargs["video_folder_def"]
+
+ video_folder = kwargs["video_folder_def"]["video_folder"]
+
+ # get the frame format
+ frame_format = kwargs["video_folder_def"]["frame format"]
+
+ # get the height
+ height = kwargs["video_folder_def"]["height"]
+
+ # get the width
+ width = kwargs["video_folder_def"]["width"]
+
+ # get the fps
+ fps = kwargs["video_folder_def"]["fps"]
+
+ # get the frames
+ frames = kwargs["IMAGE"]
+
+ # get the number of frames
+ num_frames = frames.shape[0]
+
+ # get the rescale method
+ rescale_method = kwargs.get("Rescale method", None)
+ if rescale_method is None:
+ rescale_method = kwargs.get("Rescale_method", None)
+
+ # get the number of seconds
+ seconds = kwargs["seconds"]
+
+ # get the number of frames to hold
+ num_frames_to_hold = int(fps * seconds)
+
+ num_frames_in_this_stack = frames.shape[0]
+
+ # get the number of repeats needed
+ num_repeats_needed = num_frames_to_hold // num_frames_in_this_stack
+
+ # get the number of frames left over
+ num_frames_left_over = num_frames_to_hold % num_frames_in_this_stack
+
+ # write the first set of frames to disk
+ self.saver(frames, kwargs)
+
+ # get the name of the last frame
+ source_last_frame_name = vide_folder_ref.get_last_file_name()
+ # from the last frame find the rest of the frames names/paths
+ source_last_frame_num = int(source_last_frame_name.split("_")[-1].split(".")[0])
+ source_first_frame_num = source_last_frame_num - num_frames_in_this_stack + 1
+
+ template = source_last_frame_name.replace(str(source_last_frame_num).zfill(10), "{}")
+
+ range_of_source_frame_numbers = list(range(source_first_frame_num, source_last_frame_num + 1))
+
+ source_frame_names = [template.format(str(x).zfill(10)) for x in range_of_source_frame_numbers]
+
+ target_done_frame_num = source_last_frame_num + num_frames_to_hold - num_frames_in_this_stack
+ num_frames_remaining = num_frames_to_hold - num_frames_in_this_stack
+ target_start_frame_num = source_last_frame_num + 1
+
+ target_frame_range = list(range(target_start_frame_num, target_done_frame_num + 1))
+
+ target_frame_names = [template.format(str(x).zfill(10)) for x in target_frame_range]
+
+ set_todo = {fn: [] for fn in source_frame_names}
+ # we can now repeat the frames
+ for i, target_frame_name in enumerate(target_frame_names):
+ # keep track of what stack frame we are on
+ stack_frame_idx = i % num_frames_in_this_stack
+ stack_frame_name_to_use = source_frame_names[stack_frame_idx]
+ set_todo[stack_frame_name_to_use].append(target_frame_name)
+
+ # now we can repeat the frames
+ for frame_name in set_todo:
+ src_frame = frame_name
+ target_frames = set_todo[src_frame]
+ duplicate_file(src_frame, target_frames)
+
+ return (kwargs["video_folder_def"],)
+
+
+class Listed_HoldFramesForSecsInVideoFolder(metaclass=ABCWidgetMetaclass):
+ """This just calls HoldFrameForSecsInVideoFolder multiple times for a list of frames and a list of durations"""
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ ret = {"required": {"video_folder_def": ("VIDEO_FOLDER_DEF", {"default": None}),
+ "IMAGE": ("LIST", {"default": None}),
+ "seconds": ("LIST", {"default": None}),
+ "Rescale method": (
+ ["NEAREST", "BOX", "BILINEAR", "HAMMING", "BICUBIC", "LANCZOS"],
+ {"default": "LANCZOS"}),
+ "Release vram?": ("BOOLEAN", {"default": True}),
+ },
+ }
+ return ret
+
+ RETURN_TYPES = ("VIDEO_FOLDER_DEF",)
+ FUNCTION = "listed_hold_frames_for_secs_in_video_folder"
+
+ def listed_hold_frames_for_secs_in_video_folder(self, **kwargs):
+ # the plan is to call HoldFrameForSecsInVideoFolder for each frame and duration
+ # and then return the video_folder_def
+
+ for frame, duration in zip(kwargs["IMAGE"], kwargs["seconds"]):
+ HoldFrameForSecsInVideoFolder().hold_frame_for_secs_in_video_folder(
+ video_folder_def=kwargs["video_folder_def"],
+ IMAGE=frame,
+ seconds=duration,
+ Rescale_method=kwargs["Rescale method"],
+ Release_vram_=kwargs["Release vram?"])
+
+ return (kwargs["video_folder_def"],)
+
+
+class VideoFolderDefInfo(metaclass=ABCWidgetMetaclass):
+ """
+ Returns everything in the VIDEO_FOLDER_DEF, and also returns information about the status of the folder
+ """
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ ret = {"required": {"video_folder_def": ("VIDEO_FOLDER_DEF", {"default": None}),
+ },
+ }
+ return ret
+
+ RETURN_TYPES = ("VIDEO_FOLDER_DEF", "STRING", "INT", "INT", "FLOAT", "STRING", "INT", "LIST",)
+ RETURN_NAMES = ("VFD", "video_folder", "height", "width", "fps", "frame format", "num_files", "list of files")
+
+ OUTPUT_NODE = True
+ FUNCTION = "video_folder_def_info"
+
+ def video_folder_def_info(self, **kwargs):
+ import os
+
+ video_folder = kwargs["video_folder_def"]["video_folder"]
+ height = kwargs["video_folder_def"]["height"]
+ width = kwargs["video_folder_def"]["width"]
+ fps = kwargs["video_folder_def"]["fps"]
+ frame_format = kwargs["video_folder_def"]["frame format"]
+ files = os.listdir(video_folder)
+ num_files = len(files)
+
+ return (kwargs["video_folder_def"], video_folder, height, width, fps, frame_format, num_files, files,)
+
+
+class FadeToInVideoFolder(metaclass=ABCWidgetMetaclass):
+ """
+ given a VIDEO_FOLDER_DEF node, and a stack of frames, fades the frames in and adds them to the video folder
+ """
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ ret = {"required": {"video_folder_def": ("VIDEO_FOLDER_DEF", {"default": None}),
+ "IMAGE": ("IMAGE", {"default": None}),
+ "seconds": ("FLOAT", {"default": 1.0}),
+ "transition ms": ("FLOAT", {"default": 0.1}),
+ "Rescale method": (
+ ["NEAREST", "BOX", "BILINEAR", "HAMMING", "BICUBIC", "LANCZOS"],
+ {"default": "LANCZOS"}),
+ "Release vram?": ("BOOLEAN", {"default": True}),
+ },
+ }
+ return ret
+
+ RETURN_TYPES = ("VIDEO_FOLDER_DEF",)
+ FUNCTION = "fade_to_in_video_folder"
+
+ def fade_to_in_video_folder(self, **kwargs):
+ from PIL import Image
+ import torch
+
+ # usee HoldFrameForSecsInVideoFolder to hold the frame after creating and AddFramesToVideoFolder the transition
+
+ transition_length_secs = int(kwargs["transition ms"] * kwargs["video_folder_def"]["fps"]) / 1000
+ hold_frame_length_secs = int(kwargs["seconds"] * kwargs["video_folder_def"]["fps"]) - transition_length_secs
+
+ transition_frame_count = int(kwargs["video_folder_def"]["fps"] * transition_length_secs)
+
+ base_frame_fn = kwargs["video_folder_def"].get_last_file_name()
+ # load it into a pytorch tensor (B, H, W, C)
+ base_frame = Image.open(base_frame_fn)
+ base_frame = torch.tensor(base_frame).permute(2, 0, 1).unsqueeze(0).to("cuda")
+
+ transition_stack = torch.zeros((transition_frame_count,
+ base_frame.shape[1],
+ base_frame.shape[2],
+ base_frame.shape[3]),
+ dtype=torch.float16, device="cuda")
+
+ # fade from base_frame to the first frame in the stack over the transition length
+ new_frame = kwargs["IMAGE"][0].to("cuda")
+ for i in range(transition_frame_count):
+ alpha = i / transition_frame_count
+ new_frame = (1 - alpha) * base_frame + alpha * new_frame
+ new_frame = new_frame.to("cpu")
+ transition_stack[i] = new_frame
+
+ # add to the video folder
+ AddFramesToVideoFolder().add_frames_to_video_folder(
+ video_folder_def=kwargs["video_folder_def"],
+ IMAGE=transition_stack,
+ Rescale_method=kwargs["Rescale method"],
+ Release_vram_=kwargs["Release vram?"])
+
+ # hold the frame
+ HoldFrameForSecsInVideoFolder().hold_frame_for_secs_in_video_folder(
+ video_folder_def=kwargs["video_folder_def"],
+ IMAGE=kwargs["IMAGE"],
+ seconds=hold_frame_length_secs,
+ Rescale_method=kwargs["Rescale method"],
+ Release_vram_=kwargs["Release vram?"])
+
+ return (kwargs["video_folder_def"],)
+
+
+class ListedSVDConditioning(metaclass=ABCWidgetMetaclass):
+ """
+ repeats SVD_img2vid_Conditioning for each set of params
+ """
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ base = SVD_img2vid_Conditioning.INPUT_TYPES()
+ base["required"]["init_image"] = ("LIST", {"default": None})
+ # base["required"]["clip_vision"] = ("LIST", {"default": None})
+ base["required"]["video_frames"] = ("LIST", {"default": None})
+
+ return base
+
+ RETURN_TYPES = ("LIST", "LIST", "LIST")
+ RETURN_NAMES = ("positive", "negative", "latent")
+
+ FUNCTION = "listed_SVD_conditioning"
+
+ def listed_SVD_conditioning(self, **kwargs):
+ # just like Listed_HoldFramesForSecsInVideoFolder, we will call SVD_img2vid_Conditioning for each set of params
+
+ positives = []
+ negatives = []
+ latents = []
+
+ init_images = kwargs["init_image"]
+ clip_vision = kwargs["clip_vision"]
+ # vae
+ vae = kwargs["vae"]
+ # width and height
+ width = kwargs["width"]
+ height = kwargs["height"]
+ video_frames = kwargs["video_frames"]
+ motion_bucket_id = kwargs["motion_bucket_id"]
+ fps = kwargs["fps"]
+ augmentation_level = kwargs["augmentation_level"]
+
+ lens = [len(init_images), len(video_frames)]
+
+ # check if one length is divisible by the other
+ if lens[0] % lens[1] != 0 and lens[1] % lens[0] != 0:
+ raise ValueError("The length of init_image and video_frames must be divisible by each other")
+
+ # repeat the shorter
+ if lens[0] < lens[1]:
+ init_images = init_images * (lens[1] // lens[0])
+ else:
+ video_frames = video_frames * (lens[0] // lens[1])
+
+ # at this point the lengths should be the same
+
+ for init_image, video_frame in zip(init_images, video_frames):
+ pos, neg, latent = SVD_img2vid_Conditioning().encode(
+ clip_vision=clip_vision,
+ init_image=init_image,
+ video_frames=video_frame,
+ vae=vae,
+ width=width,
+ height=height,
+ motion_bucket_id=motion_bucket_id,
+ fps=fps,
+ augmentation_level=augmentation_level
+ )
+
+ positives.append(pos)
+ negatives.append(neg)
+ latents.append(latent)
+
+ return (positives, negatives, latents,)
+
+
+class VideoFolderFFMPEGVideoInterpolate(metaclass=ABCWidgetMetaclass):
+ """
+ use run_ffmpeg to interpolate the frames in the video folder
+ """
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ ret = {"required": {"video_folder_def": ("VIDEO_FOLDER_DEF", {"default": None}),
+ "last_n_frames": ("INT", {"default": 2}),
+ "target_n_frames": ("INT", {"default": 3}),
+ },
+ }
+ return ret
+
+ RETURN_TYPES = ("VIDEO_FOLDER_DEF",)
+ RETURN_NAMES = ("video_folder_def",)
+ FUNCTION = "video_folder_ffmpeg_video_interpolate"
+
+ def video_folder_ffmpeg_video_interpolate(self, **kwargs):
+ import os
+ import shutil
+
+ video_folder_def = kwargs["video_folder_def"]
+ last_n_frames = kwargs["last_n_frames"]
+ target_n_frames = kwargs["target_n_frames"]
+
+ # Get the current frame count
+ current_frame_count = video_folder_def.get_frame_count()
+
+ # Calculate the target frame rate for interpolation
+ target_fps = video_folder_def["fps"] * (target_n_frames / last_n_frames)
+
+ # Get the file names for the last n frames
+ last_frame_names = [video_folder_def.get_last_file_name(offset=i) for i in range(last_n_frames)]
+ last_frame_names = list(set(last_frame_names))
+ last_frame_names.sort()
+
+ # Create a temporary folder off of the video frame folder's base path
+ temp_folder = os.path.join(video_folder_def["video_folder"], "temp")
+ os.makedirs(temp_folder, exist_ok=True)
+
+ # put the input frames into the new folder
+ for frame_name in last_frame_names:
+ # do this in a way that overwrites the files and MOVES them
+ shutil.move(frame_name, os.path.join(temp_folder, os.path.basename(frame_name)))
+
+ # new input frame names list
+ last_frame_names = [os.path.join(temp_folder, os.path.basename(frame_name)) for frame_name in last_frame_names]
+ last_frame_names.sort()
+
+ # Create a file with the list of input frame paths in the temporary folder
+ input_frame_list_file = os.path.join(temp_folder, "input_frames.txt")
+ input_frame_list_file = os.path.abspath(input_frame_list_file)
+ with open(input_frame_list_file, 'w') as file:
+ # need to do this in a manner that puts the first frame at the top of the file.
+ file.write('\n'.join([f"file '{frame_name}'" for frame_name in last_frame_names]))
+
+ # Get the output frame names for the interpolated frames
+ # output_frame_names = [video_folder_def.reserve_next_frame() for _ in range(target_n_frames)]
+ output_frame_names = [video_folder_def.get_next_file_name(i) for i in range(target_n_frames)]
+
+ # delete the frames on the list
+ #for frame_name in output_frame_names:
+ # os.remove(frame_name)
+
+ temp_output_frame_pattern = os.path.join(video_folder_def["video_folder"], "temp",
+ "tmp_frame_%010d." + video_folder_def["frame format"].lower())
+ temp_output_frame_pattern=os.path.abspath(temp_output_frame_pattern)
+
+ # Create the FFmpeg command for motion interpolation using the frame list file
+ ffmpeg_interpolate_cmd = [
+ "-f", "concat", "-safe", "0",
+ "-i", input_frame_list_file,
+ "-vsync", "0", "-filter_complex",
+ # f"[0:v]minterpolate=fps={target_fps}:mi_mode=mci:mc_mode=aobmc:me_mode=bidir:me=ds:mb_size=16:search_param=8:vsbmc=1[v]",
+ # f"[0:v]framerate=fps={target_fps}:interp_start=0:interp_end=255:scene=0[v]",
+ # f"[0:v]minterpolate=fps={target_fps}:vsbmc=1[v]",
+ f"[0:v]minterpolate=fps={target_fps}:mi_mode=blend[v]",
+ "-map", "[v]", "-q:v", "2", "-r", str(target_fps), "-frames", str(target_n_frames),
+ "-threads", "0",
+ "-y", temp_output_frame_pattern
+ ]
+
+ run_ffmpeg_command(ffmpeg_interpolate_cmd)
+
+ # Move the interpolated frames to the video folder with the correct names
+ for i, frame_name in enumerate(output_frame_names):
+ temp_frame_name = os.path.join(temp_folder,
+ f"tmp_frame_{i + 1:010d}.{video_folder_def['frame format'].lower()}")
+ try:
+ shutil.move(temp_frame_name, frame_name)
+ except FileNotFoundError:
+ pass
+
+ # Remove the temporary folder and frame list file
+ os.remove(input_frame_list_file)
+ # remove this without generating an error even if it is not empty
+ try:
+ os.rmdir(temp_folder)
+ except OSError:
+ shutil.rmtree(temp_folder)
+
+ return (video_folder_def,)
diff --git a/video_v2.py b/video_v2.py
new file mode 100644
index 0000000..e69de29
diff --git a/web/marked.js b/web/marked.js
new file mode 100644
index 0000000..3128e26
--- /dev/null
+++ b/web/marked.js
@@ -0,0 +1,1191 @@
+/**
+ * marked v12.0.1 - a markdown parser
+ * Copyright (c) 2011-2024, Christopher Jeffrey. (MIT Licensed)
+ * https://github.com/markedjs/marked
+ */
+!function (e, t) {
+ "object" == typeof exports && "undefined" != typeof module ? t(exports) : "function" == typeof define && define.amd ? define(["exports"], t) : t((e = "undefined" != typeof globalThis ? globalThis : e || self).marked = {})
+}(this, (function (e) {
+ "use strict";
+
+ function t() {
+ return {
+ async: !1,
+ breaks: !1,
+ extensions: null,
+ gfm: !0,
+ hooks: null,
+ pedantic: !1,
+ renderer: null,
+ silent: !1,
+ tokenizer: null,
+ walkTokens: null
+ }
+ }
+
+ function n(t) {
+ e.defaults = t
+ }
+
+ e.defaults = {
+ async: !1,
+ breaks: !1,
+ extensions: null,
+ gfm: !0,
+ hooks: null,
+ pedantic: !1,
+ renderer: null,
+ silent: !1,
+ tokenizer: null,
+ walkTokens: null
+ };
+ const s = /[&<>"']/, r = new RegExp(s.source, "g"), i = /[<>"']|&(?!(#\d{1,7}|#[Xx][a-fA-F0-9]{1,6}|\w+);)/,
+ l = new RegExp(i.source, "g"), o = {"&": "&", "<": "<", ">": ">", '"': """, "'": "'"},
+ a = e => o[e];
+
+ function c(e, t) {
+ if (t) {
+ if (s.test(e)) return e.replace(r, a)
+ } else if (i.test(e)) return e.replace(l, a);
+ return e
+ }
+
+ const h = /&(#(?:\d+)|(?:#x[0-9A-Fa-f]+)|(?:\w+));?/gi;
+
+ function p(e) {
+ return e.replace(h, ((e, t) => "colon" === (t = t.toLowerCase()) ? ":" : "#" === t.charAt(0) ? "x" === t.charAt(1) ? String.fromCharCode(parseInt(t.substring(2), 16)) : String.fromCharCode(+t.substring(1)) : ""))
+ }
+
+ const u = /(^|[^\[])\^/g;
+
+ function k(e, t) {
+ let n = "string" == typeof e ? e : e.source;
+ t = t || "";
+ const s = {
+ replace: (e, t) => {
+ let r = "string" == typeof t ? t : t.source;
+ return r = r.replace(u, "$1"), n = n.replace(e, r), s
+ }, getRegex: () => new RegExp(n, t)
+ };
+ return s
+ }
+
+ function g(e) {
+ try {
+ e = encodeURI(e).replace(/%25/g, "%")
+ } catch (e) {
+ return null
+ }
+ return e
+ }
+
+ const f = {exec: () => null};
+
+ function d(e, t) {
+ const n = e.replace(/\|/g, ((e, t, n) => {
+ let s = !1, r = t;
+ for (; --r >= 0 && "\\" === n[r];) s = !s;
+ return s ? "|" : " |"
+ })).split(/ \|/);
+ let s = 0;
+ if (n[0].trim() || n.shift(), n.length > 0 && !n[n.length - 1].trim() && n.pop(), t) if (n.length > t) n.splice(t); else for (; n.length < t;) n.push("");
+ for (; s < n.length; s++) n[s] = n[s].trim().replace(/\\\|/g, "|");
+ return n
+ }
+
+ function x(e, t, n) {
+ const s = e.length;
+ if (0 === s) return "";
+ let r = 0;
+ for (; r < s;) {
+ const i = e.charAt(s - r - 1);
+ if (i !== t || n) {
+ if (i === t || !n) break;
+ r++
+ } else r++
+ }
+ return e.slice(0, s - r)
+ }
+
+ function b(e, t, n, s) {
+ const r = t.href, i = t.title ? c(t.title) : null, l = e[1].replace(/\\([\[\]])/g, "$1");
+ if ("!" !== e[0].charAt(0)) {
+ s.state.inLink = !0;
+ const e = {type: "link", raw: n, href: r, title: i, text: l, tokens: s.inlineTokens(l)};
+ return s.state.inLink = !1, e
+ }
+ return {type: "image", raw: n, href: r, title: i, text: c(l)}
+ }
+
+ class w {
+ options;
+ rules;
+ lexer;
+
+ constructor(t) {
+ this.options = t || e.defaults
+ }
+
+ space(e) {
+ const t = this.rules.block.newline.exec(e);
+ if (t && t[0].length > 0) return {type: "space", raw: t[0]}
+ }
+
+ code(e) {
+ const t = this.rules.block.code.exec(e);
+ if (t) {
+ const e = t[0].replace(/^ {1,4}/gm, "");
+ return {
+ type: "code",
+ raw: t[0],
+ codeBlockStyle: "indented",
+ text: this.options.pedantic ? e : x(e, "\n")
+ }
+ }
+ }
+
+ fences(e) {
+ const t = this.rules.block.fences.exec(e);
+ if (t) {
+ const e = t[0], n = function (e, t) {
+ const n = e.match(/^(\s+)(?:```)/);
+ if (null === n) return t;
+ const s = n[1];
+ return t.split("\n").map((e => {
+ const t = e.match(/^\s+/);
+ if (null === t) return e;
+ const [n] = t;
+ return n.length >= s.length ? e.slice(s.length) : e
+ })).join("\n")
+ }(e, t[3] || "");
+ return {
+ type: "code",
+ raw: e,
+ lang: t[2] ? t[2].trim().replace(this.rules.inline.anyPunctuation, "$1") : t[2],
+ text: n
+ }
+ }
+ }
+
+ heading(e) {
+ const t = this.rules.block.heading.exec(e);
+ if (t) {
+ let e = t[2].trim();
+ if (/#$/.test(e)) {
+ const t = x(e, "#");
+ this.options.pedantic ? e = t.trim() : t && !/ $/.test(t) || (e = t.trim())
+ }
+ return {type: "heading", raw: t[0], depth: t[1].length, text: e, tokens: this.lexer.inline(e)}
+ }
+ }
+
+ hr(e) {
+ const t = this.rules.block.hr.exec(e);
+ if (t) return {type: "hr", raw: t[0]}
+ }
+
+ blockquote(e) {
+ const t = this.rules.block.blockquote.exec(e);
+ if (t) {
+ const e = x(t[0].replace(/^ *>[ \t]?/gm, ""), "\n"), n = this.lexer.state.top;
+ this.lexer.state.top = !0;
+ const s = this.lexer.blockTokens(e);
+ return this.lexer.state.top = n, {type: "blockquote", raw: t[0], tokens: s, text: e}
+ }
+ }
+
+ list(e) {
+ let t = this.rules.block.list.exec(e);
+ if (t) {
+ let n = t[1].trim();
+ const s = n.length > 1,
+ r = {type: "list", raw: "", ordered: s, start: s ? +n.slice(0, -1) : "", loose: !1, items: []};
+ n = s ? `\\d{1,9}\\${n.slice(-1)}` : `\\${n}`, this.options.pedantic && (n = s ? n : "[*+-]");
+ const i = new RegExp(`^( {0,3}${n})((?:[\t ][^\\n]*)?(?:\\n|$))`);
+ let l = "", o = "", a = !1;
+ for (; e;) {
+ let n = !1;
+ if (!(t = i.exec(e))) break;
+ if (this.rules.block.hr.test(e)) break;
+ l = t[0], e = e.substring(l.length);
+ let s = t[2].split("\n", 1)[0].replace(/^\t+/, (e => " ".repeat(3 * e.length))),
+ c = e.split("\n", 1)[0], h = 0;
+ this.options.pedantic ? (h = 2, o = s.trimStart()) : (h = t[2].search(/[^ ]/), h = h > 4 ? 1 : h, o = s.slice(h), h += t[1].length);
+ let p = !1;
+ if (!s && /^ *$/.test(c) && (l += c + "\n", e = e.substring(c.length + 1), n = !0), !n) {
+ const t = new RegExp(`^ {0,${Math.min(3,h-1)}}(?:[*+-]|\\d{1,9}[.)])((?:[ \t][^\\n]*)?(?:\\n|$))`),
+ n = new RegExp(`^ {0,${Math.min(3,h-1)}}((?:- *){3,}|(?:_ *){3,}|(?:\\* *){3,})(?:\\n+|$)`),
+ r = new RegExp(`^ {0,${Math.min(3,h-1)}}(?:\`\`\`|~~~)`),
+ i = new RegExp(`^ {0,${Math.min(3,h-1)}}#`);
+ for (; e;) {
+ const a = e.split("\n", 1)[0];
+ if (c = a, this.options.pedantic && (c = c.replace(/^ {1,4}(?=( {4})*[^ ])/g, " ")), r.test(c)) break;
+ if (i.test(c)) break;
+ if (t.test(c)) break;
+ if (n.test(e)) break;
+ if (c.search(/[^ ]/) >= h || !c.trim()) o += "\n" + c.slice(h); else {
+ if (p) break;
+ if (s.search(/[^ ]/) >= 4) break;
+ if (r.test(s)) break;
+ if (i.test(s)) break;
+ if (n.test(s)) break;
+ o += "\n" + c
+ }
+ p || c.trim() || (p = !0), l += a + "\n", e = e.substring(a.length + 1), s = c.slice(h)
+ }
+ }
+ r.loose || (a ? r.loose = !0 : /\n *\n *$/.test(l) && (a = !0));
+ let u, k = null;
+ this.options.gfm && (k = /^\[[ xX]\] /.exec(o), k && (u = "[ ] " !== k[0], o = o.replace(/^\[[ xX]\] +/, ""))), r.items.push({
+ type: "list_item",
+ raw: l,
+ task: !!k,
+ checked: u,
+ loose: !1,
+ text: o,
+ tokens: []
+ }), r.raw += l
+ }
+ r.items[r.items.length - 1].raw = l.trimEnd(), r.items[r.items.length - 1].text = o.trimEnd(), r.raw = r.raw.trimEnd();
+ for (let e = 0; e < r.items.length; e++) if (this.lexer.state.top = !1, r.items[e].tokens = this.lexer.blockTokens(r.items[e].text, []), !r.loose) {
+ const t = r.items[e].tokens.filter((e => "space" === e.type)),
+ n = t.length > 0 && t.some((e => /\n.*\n/.test(e.raw)));
+ r.loose = n
+ }
+ if (r.loose) for (let e = 0; e < r.items.length; e++) r.items[e].loose = !0;
+ return r
+ }
+ }
+
+ html(e) {
+ const t = this.rules.block.html.exec(e);
+ if (t) {
+ return {
+ type: "html",
+ block: !0,
+ raw: t[0],
+ pre: "pre" === t[1] || "script" === t[1] || "style" === t[1],
+ text: t[0]
+ }
+ }
+ }
+
+ def(e) {
+ const t = this.rules.block.def.exec(e);
+ if (t) {
+ const e = t[1].toLowerCase().replace(/\s+/g, " "),
+ n = t[2] ? t[2].replace(/^<(.*)>$/, "$1").replace(this.rules.inline.anyPunctuation, "$1") : "",
+ s = t[3] ? t[3].substring(1, t[3].length - 1).replace(this.rules.inline.anyPunctuation, "$1") : t[3];
+ return {type: "def", tag: e, raw: t[0], href: n, title: s}
+ }
+ }
+
+ table(e) {
+ const t = this.rules.block.table.exec(e);
+ if (!t) return;
+ if (!/[:|]/.test(t[2])) return;
+ const n = d(t[1]), s = t[2].replace(/^\||\| *$/g, "").split("|"),
+ r = t[3] && t[3].trim() ? t[3].replace(/\n[ \t]*$/, "").split("\n") : [],
+ i = {type: "table", raw: t[0], header: [], align: [], rows: []};
+ if (n.length === s.length) {
+ for (const e of s) /^ *-+: *$/.test(e) ? i.align.push("right") : /^ *:-+: *$/.test(e) ? i.align.push("center") : /^ *:-+ *$/.test(e) ? i.align.push("left") : i.align.push(null);
+ for (const e of n) i.header.push({text: e, tokens: this.lexer.inline(e)});
+ for (const e of r) i.rows.push(d(e, i.header.length).map((e => ({
+ text: e,
+ tokens: this.lexer.inline(e)
+ }))));
+ return i
+ }
+ }
+
+ lheading(e) {
+ const t = this.rules.block.lheading.exec(e);
+ if (t) return {
+ type: "heading",
+ raw: t[0],
+ depth: "=" === t[2].charAt(0) ? 1 : 2,
+ text: t[1],
+ tokens: this.lexer.inline(t[1])
+ }
+ }
+
+ paragraph(e) {
+ const t = this.rules.block.paragraph.exec(e);
+ if (t) {
+ const e = "\n" === t[1].charAt(t[1].length - 1) ? t[1].slice(0, -1) : t[1];
+ return {type: "paragraph", raw: t[0], text: e, tokens: this.lexer.inline(e)}
+ }
+ }
+
+ text(e) {
+ const t = this.rules.block.text.exec(e);
+ if (t) return {type: "text", raw: t[0], text: t[0], tokens: this.lexer.inline(t[0])}
+ }
+
+ escape(e) {
+ const t = this.rules.inline.escape.exec(e);
+ if (t) return {type: "escape", raw: t[0], text: c(t[1])}
+ }
+
+ tag(e) {
+ const t = this.rules.inline.tag.exec(e);
+ if (t) return !this.lexer.state.inLink && /^/i.test(t[0]) && (this.lexer.state.inLink = !1), !this.lexer.state.inRawBlock && /^<(pre|code|kbd|script)(\s|>)/i.test(t[0]) ? this.lexer.state.inRawBlock = !0 : this.lexer.state.inRawBlock && /^<\/(pre|code|kbd|script)(\s|>)/i.test(t[0]) && (this.lexer.state.inRawBlock = !1), {
+ type: "html",
+ raw: t[0],
+ inLink: this.lexer.state.inLink,
+ inRawBlock: this.lexer.state.inRawBlock,
+ block: !1,
+ text: t[0]
+ }
+ }
+
+ link(e) {
+ const t = this.rules.inline.link.exec(e);
+ if (t) {
+ const e = t[2].trim();
+ if (!this.options.pedantic && /^$/.test(e)) return;
+ const t = x(e.slice(0, -1), "\\");
+ if ((e.length - t.length) % 2 == 0) return
+ } else {
+ const e = function (e, t) {
+ if (-1 === e.indexOf(t[1])) return -1;
+ let n = 0;
+ for (let s = 0; s < e.length; s++) if ("\\" === e[s]) s++; else if (e[s] === t[0]) n++; else if (e[s] === t[1] && (n--, n < 0)) return s;
+ return -1
+ }(t[2], "()");
+ if (e > -1) {
+ const n = (0 === t[0].indexOf("!") ? 5 : 4) + t[1].length + e;
+ t[2] = t[2].substring(0, e), t[0] = t[0].substring(0, n).trim(), t[3] = ""
+ }
+ }
+ let n = t[2], s = "";
+ if (this.options.pedantic) {
+ const e = /^([^'"]*[^\s])\s+(['"])(.*)\2/.exec(n);
+ e && (n = e[1], s = e[3])
+ } else s = t[3] ? t[3].slice(1, -1) : "";
+ return n = n.trim(), /^$/.test(e) ? n.slice(1) : n.slice(1, -1)), b(t, {
+ href: n ? n.replace(this.rules.inline.anyPunctuation, "$1") : n,
+ title: s ? s.replace(this.rules.inline.anyPunctuation, "$1") : s
+ }, t[0], this.lexer)
+ }
+ }
+
+ reflink(e, t) {
+ let n;
+ if ((n = this.rules.inline.reflink.exec(e)) || (n = this.rules.inline.nolink.exec(e))) {
+ const e = t[(n[2] || n[1]).replace(/\s+/g, " ").toLowerCase()];
+ if (!e) {
+ const e = n[0].charAt(0);
+ return {type: "text", raw: e, text: e}
+ }
+ return b(n, e, n[0], this.lexer)
+ }
+ }
+
+ emStrong(e, t, n = "") {
+ let s = this.rules.inline.emStrongLDelim.exec(e);
+ if (!s) return;
+ if (s[3] && n.match(/[\p{L}\p{N}]/u)) return;
+ if (!(s[1] || s[2] || "") || !n || this.rules.inline.punctuation.exec(n)) {
+ const n = [...s[0]].length - 1;
+ let r, i, l = n, o = 0;
+ const a = "*" === s[0][0] ? this.rules.inline.emStrongRDelimAst : this.rules.inline.emStrongRDelimUnd;
+ for (a.lastIndex = 0, t = t.slice(-1 * e.length + n); null != (s = a.exec(t));) {
+ if (r = s[1] || s[2] || s[3] || s[4] || s[5] || s[6], !r) continue;
+ if (i = [...r].length, s[3] || s[4]) {
+ l += i;
+ continue
+ }
+ if ((s[5] || s[6]) && n % 3 && !((n + i) % 3)) {
+ o += i;
+ continue
+ }
+ if (l -= i, l > 0) continue;
+ i = Math.min(i, i + l + o);
+ const t = [...s[0]][0].length, a = e.slice(0, n + s.index + t + i);
+ if (Math.min(n, i) % 2) {
+ const e = a.slice(1, -1);
+ return {type: "em", raw: a, text: e, tokens: this.lexer.inlineTokens(e)}
+ }
+ const c = a.slice(2, -2);
+ return {type: "strong", raw: a, text: c, tokens: this.lexer.inlineTokens(c)}
+ }
+ }
+ }
+
+ codespan(e) {
+ const t = this.rules.inline.code.exec(e);
+ if (t) {
+ let e = t[2].replace(/\n/g, " ");
+ const n = /[^ ]/.test(e), s = /^ /.test(e) && / $/.test(e);
+ return n && s && (e = e.substring(1, e.length - 1)), e = c(e, !0), {
+ type: "codespan",
+ raw: t[0],
+ text: e
+ }
+ }
+ }
+
+ br(e) {
+ const t = this.rules.inline.br.exec(e);
+ if (t) return {type: "br", raw: t[0]}
+ }
+
+ del(e) {
+ const t = this.rules.inline.del.exec(e);
+ if (t) return {type: "del", raw: t[0], text: t[2], tokens: this.lexer.inlineTokens(t[2])}
+ }
+
+ autolink(e) {
+ const t = this.rules.inline.autolink.exec(e);
+ if (t) {
+ let e, n;
+ return "@" === t[2] ? (e = c(t[1]), n = "mailto:" + e) : (e = c(t[1]), n = e), {
+ type: "link",
+ raw: t[0],
+ text: e,
+ href: n,
+ tokens: [{type: "text", raw: e, text: e}]
+ }
+ }
+ }
+
+ url(e) {
+ let t;
+ if (t = this.rules.inline.url.exec(e)) {
+ let e, n;
+ if ("@" === t[2]) e = c(t[0]), n = "mailto:" + e; else {
+ let s;
+ do {
+ s = t[0], t[0] = this.rules.inline._backpedal.exec(t[0])?.[0] ?? ""
+ } while (s !== t[0]);
+ e = c(t[0]), n = "www." === t[1] ? "http://" + t[0] : t[0]
+ }
+ return {type: "link", raw: t[0], text: e, href: n, tokens: [{type: "text", raw: e, text: e}]}
+ }
+ }
+
+ inlineText(e) {
+ const t = this.rules.inline.text.exec(e);
+ if (t) {
+ let e;
+ return e = this.lexer.state.inRawBlock ? t[0] : c(t[0]), {type: "text", raw: t[0], text: e}
+ }
+ }
+ }
+
+ const m = /^ {0,3}((?:-[\t ]*){3,}|(?:_[ \t]*){3,}|(?:\*[ \t]*){3,})(?:\n+|$)/, y = /(?:[*+-]|\d{1,9}[.)])/,
+ $ = k(/^(?!bull |blockCode|fences|blockquote|heading|html)((?:.|\n(?!\s*?\n|bull |blockCode|fences|blockquote|heading|html))+?)\n {0,3}(=+|-+) *(?:\n+|$)/).replace(/bull/g, y).replace(/blockCode/g, / {4}/).replace(/fences/g, / {0,3}(?:`{3,}|~{3,})/).replace(/blockquote/g, / {0,3}>/).replace(/heading/g, / {0,3}#{1,6}/).replace(/html/g, / {0,3}<[^\n>]+>\n/).getRegex(),
+ z = /^([^\n]+(?:\n(?!hr|heading|lheading|blockquote|fences|list|html|table| +\n)[^\n]+)*)/,
+ T = /(?!\s*\])(?:\\.|[^\[\]\\])+/,
+ R = k(/^ {0,3}\[(label)\]: *(?:\n *)?([^<\s][^\s]*|<.*?>)(?:(?: +(?:\n *)?| *\n *)(title))? *(?:\n+|$)/).replace("label", T).replace("title", /(?:"(?:\\"?|[^"\\])*"|'[^'\n]*(?:\n[^'\n]+)*\n?'|\([^()]*\))/).getRegex(),
+ _ = k(/^( {0,3}bull)([ \t][^\n]+?)?(?:\n|$)/).replace(/bull/g, y).getRegex(),
+ A = "address|article|aside|base|basefont|blockquote|body|caption|center|col|colgroup|dd|details|dialog|dir|div|dl|dt|fieldset|figcaption|figure|footer|form|frame|frameset|h[1-6]|head|header|hr|html|iframe|legend|li|link|main|menu|menuitem|meta|nav|noframes|ol|optgroup|option|p|param|search|section|summary|table|tbody|td|tfoot|th|thead|title|tr|track|ul",
+ S = /|$))/,
+ I = k("^ {0,3}(?:<(script|pre|style|textarea)[\\s>][\\s\\S]*?(?:[^\\n]*\\n+|$)|comment[^\\n]*(\\n+|$)|<\\?[\\s\\S]*?(?:\\?>\\n*|$)|\\n*|$)|\\n*|$)|)[\\s\\S]*?(?:(?:\\n *)+\\n|$)|<(?!script|pre|style|textarea)([a-z][\\w-]*)(?:attribute)*? */?>(?=[ \\t]*(?:\\n|$))[\\s\\S]*?(?:(?:\\n *)+\\n|$)|(?=[ \\t]*(?:\\n|$))[\\s\\S]*?(?:(?:\\n *)+\\n|$))", "i").replace("comment", S).replace("tag", A).replace("attribute", / +[a-zA-Z:_][\w.:-]*(?: *= *"[^"\n]*"| *= *'[^'\n]*'| *= *[^\s"'=<>`]+)?/).getRegex(),
+ E = k(z).replace("hr", m).replace("heading", " {0,3}#{1,6}(?:\\s|$)").replace("|lheading", "").replace("|table", "").replace("blockquote", " {0,3}>").replace("fences", " {0,3}(?:`{3,}(?=[^`\\n]*\\n)|~{3,})[^\\n]*\\n").replace("list", " {0,3}(?:[*+-]|1[.)]) ").replace("html", ")|<(?:script|pre|style|textarea|!--)").replace("tag", A).getRegex(),
+ q = {
+ blockquote: k(/^( {0,3}> ?(paragraph|[^\n]*)(?:\n|$))+/).replace("paragraph", E).getRegex(),
+ code: /^( {4}[^\n]+(?:\n(?: *(?:\n|$))*)?)+/,
+ def: R,
+ fences: /^ {0,3}(`{3,}(?=[^`\n]*(?:\n|$))|~{3,})([^\n]*)(?:\n|$)(?:|([\s\S]*?)(?:\n|$))(?: {0,3}\1[~`]* *(?=\n|$)|$)/,
+ heading: /^ {0,3}(#{1,6})(?=\s|$)(.*)(?:\n+|$)/,
+ hr: m,
+ html: I,
+ lheading: $,
+ list: _,
+ newline: /^(?: *(?:\n|$))+/,
+ paragraph: E,
+ table: f,
+ text: /^[^\n]+/
+ },
+ Z = k("^ *([^\\n ].*)\\n {0,3}((?:\\| *)?:?-+:? *(?:\\| *:?-+:? *)*(?:\\| *)?)(?:\\n((?:(?! *\\n|hr|heading|blockquote|code|fences|list|html).*(?:\\n|$))*)\\n*|$)").replace("hr", m).replace("heading", " {0,3}#{1,6}(?:\\s|$)").replace("blockquote", " {0,3}>").replace("code", " {4}[^\\n]").replace("fences", " {0,3}(?:`{3,}(?=[^`\\n]*\\n)|~{3,})[^\\n]*\\n").replace("list", " {0,3}(?:[*+-]|1[.)]) ").replace("html", ")|<(?:script|pre|style|textarea|!--)").replace("tag", A).getRegex(),
+ L = {
+ ...q,
+ table: Z,
+ paragraph: k(z).replace("hr", m).replace("heading", " {0,3}#{1,6}(?:\\s|$)").replace("|lheading", "").replace("table", Z).replace("blockquote", " {0,3}>").replace("fences", " {0,3}(?:`{3,}(?=[^`\\n]*\\n)|~{3,})[^\\n]*\\n").replace("list", " {0,3}(?:[*+-]|1[.)]) ").replace("html", ")|<(?:script|pre|style|textarea|!--)").replace("tag", A).getRegex()
+ }, P = {
+ ...q,
+ html: k("^ *(?:comment *(?:\\n|\\s*$)|<(tag)[\\s\\S]+? *(?:\\n{2,}|\\s*$)|\\s]*)*?/?> *(?:\\n{2,}|\\s*$))").replace("comment", S).replace(/tag/g, "(?!(?:a|em|strong|small|s|cite|q|dfn|abbr|data|time|code|var|samp|kbd|sub|sup|i|b|u|mark|ruby|rt|rp|bdi|bdo|span|br|wbr|ins|del|img)\\b)\\w+(?!:|[^\\w\\s@]*@)\\b").getRegex(),
+ def: /^ *\[([^\]]+)\]: *]+)>?(?: +(["(][^\n]+[")]))? *(?:\n+|$)/,
+ heading: /^(#{1,6})(.*)(?:\n+|$)/,
+ fences: f,
+ lheading: /^(.+?)\n {0,3}(=+|-+) *(?:\n+|$)/,
+ paragraph: k(z).replace("hr", m).replace("heading", " *#{1,6} *[^\n]").replace("lheading", $).replace("|table", "").replace("blockquote", " {0,3}>").replace("|fences", "").replace("|list", "").replace("|html", "").replace("|tag", "").getRegex()
+ }, Q = /^\\([!"#$%&'()*+,\-./:;<=>?@\[\]\\^_`{|}~])/, v = /^( {2,}|\\)\n(?!\s*$)/, B = "\\p{P}\\p{S}",
+ C = k(/^((?![*_])[\spunctuation])/, "u").replace(/punctuation/g, B).getRegex(),
+ M = k(/^(?:\*+(?:((?!\*)[punct])|[^\s*]))|^_+(?:((?!_)[punct])|([^\s_]))/, "u").replace(/punct/g, B).getRegex(),
+ O = k("^[^_*]*?__[^_*]*?\\*[^_*]*?(?=__)|[^*]+(?=[^*])|(?!\\*)[punct](\\*+)(?=[\\s]|$)|[^punct\\s](\\*+)(?!\\*)(?=[punct\\s]|$)|(?!\\*)[punct\\s](\\*+)(?=[^punct\\s])|[\\s](\\*+)(?!\\*)(?=[punct])|(?!\\*)[punct](\\*+)(?!\\*)(?=[punct])|[^punct\\s](\\*+)(?=[^punct\\s])", "gu").replace(/punct/g, B).getRegex(),
+ D = k("^[^_*]*?\\*\\*[^_*]*?_[^_*]*?(?=\\*\\*)|[^_]+(?=[^_])|(?!_)[punct](_+)(?=[\\s]|$)|[^punct\\s](_+)(?!_)(?=[punct\\s]|$)|(?!_)[punct\\s](_+)(?=[^punct\\s])|[\\s](_+)(?!_)(?=[punct])|(?!_)[punct](_+)(?!_)(?=[punct])", "gu").replace(/punct/g, B).getRegex(),
+ j = k(/\\([punct])/, "gu").replace(/punct/g, B).getRegex(),
+ H = k(/^<(scheme:[^\s\x00-\x1f<>]*|email)>/).replace("scheme", /[a-zA-Z][a-zA-Z0-9+.-]{1,31}/).replace("email", /[a-zA-Z0-9.!#$%&'*+/=?^_`{|}~-]+(@)[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?(?:\.[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?)+(?![-_])/).getRegex(),
+ U = k(S).replace("(?:--\x3e|$)", "--\x3e").getRegex(),
+ X = k("^comment|^|^<[a-zA-Z][\\w-]*(?:attribute)*?\\s*/?>|^<\\?[\\s\\S]*?\\?>|^|^").replace("comment", U).replace("attribute", /\s+[a-zA-Z:_][\w.:-]*(?:\s*=\s*"[^"]*"|\s*=\s*'[^']*'|\s*=\s*[^\s"'=<>`]+)?/).getRegex(),
+ F = /(?:\[(?:\\.|[^\[\]\\])*\]|\\.|`[^`]*`|[^\[\]\\`])*?/,
+ N = k(/^!?\[(label)\]\(\s*(href)(?:\s+(title))?\s*\)/).replace("label", F).replace("href", /<(?:\\.|[^\n<>\\])+>|[^\s\x00-\x1f]*/).replace("title", /"(?:\\"?|[^"\\])*"|'(?:\\'?|[^'\\])*'|\((?:\\\)?|[^)\\])*\)/).getRegex(),
+ G = k(/^!?\[(label)\]\[(ref)\]/).replace("label", F).replace("ref", T).getRegex(),
+ J = k(/^!?\[(ref)\](?:\[\])?/).replace("ref", T).getRegex(), K = {
+ _backpedal: f,
+ anyPunctuation: j,
+ autolink: H,
+ blockSkip: /\[[^[\]]*?\]\([^\(\)]*?\)|`[^`]*?`|<[^<>]*?>/g,
+ br: v,
+ code: /^(`+)([^`]|[^`][\s\S]*?[^`])\1(?!`)/,
+ del: f,
+ emStrongLDelim: M,
+ emStrongRDelimAst: O,
+ emStrongRDelimUnd: D,
+ escape: Q,
+ link: N,
+ nolink: J,
+ punctuation: C,
+ reflink: G,
+ reflinkSearch: k("reflink|nolink(?!\\()", "g").replace("reflink", G).replace("nolink", J).getRegex(),
+ tag: X,
+ text: /^(`+|[^`])(?:(?= {2,}\n)|[\s\S]*?(?:(?=[\\ t + " ".repeat(n.length))); e;) if (!(this.options.extensions && this.options.extensions.block && this.options.extensions.block.some((s => !!(n = s.call({lexer: this}, e, t)) && (e = e.substring(n.raw.length), t.push(n), !0))))) if (n = this.tokenizer.space(e)) e = e.substring(n.raw.length), 1 === n.raw.length && t.length > 0 ? t[t.length - 1].raw += "\n" : t.push(n); else if (n = this.tokenizer.code(e)) e = e.substring(n.raw.length), s = t[t.length - 1], !s || "paragraph" !== s.type && "text" !== s.type ? t.push(n) : (s.raw += "\n" + n.raw, s.text += "\n" + n.text, this.inlineQueue[this.inlineQueue.length - 1].src = s.text); else if (n = this.tokenizer.fences(e)) e = e.substring(n.raw.length), t.push(n); else if (n = this.tokenizer.heading(e)) e = e.substring(n.raw.length), t.push(n); else if (n = this.tokenizer.hr(e)) e = e.substring(n.raw.length), t.push(n); else if (n = this.tokenizer.blockquote(e)) e = e.substring(n.raw.length), t.push(n); else if (n = this.tokenizer.list(e)) e = e.substring(n.raw.length), t.push(n); else if (n = this.tokenizer.html(e)) e = e.substring(n.raw.length), t.push(n); else if (n = this.tokenizer.def(e)) e = e.substring(n.raw.length), s = t[t.length - 1], !s || "paragraph" !== s.type && "text" !== s.type ? this.tokens.links[n.tag] || (this.tokens.links[n.tag] = {
+ href: n.href,
+ title: n.title
+ }) : (s.raw += "\n" + n.raw, s.text += "\n" + n.raw, this.inlineQueue[this.inlineQueue.length - 1].src = s.text); else if (n = this.tokenizer.table(e)) e = e.substring(n.raw.length), t.push(n); else if (n = this.tokenizer.lheading(e)) e = e.substring(n.raw.length), t.push(n); else {
+ if (r = e, this.options.extensions && this.options.extensions.startBlock) {
+ let t = 1 / 0;
+ const n = e.slice(1);
+ let s;
+ this.options.extensions.startBlock.forEach((e => {
+ s = e.call({lexer: this}, n), "number" == typeof s && s >= 0 && (t = Math.min(t, s))
+ })), t < 1 / 0 && t >= 0 && (r = e.substring(0, t + 1))
+ }
+ if (this.state.top && (n = this.tokenizer.paragraph(r))) s = t[t.length - 1], i && "paragraph" === s.type ? (s.raw += "\n" + n.raw, s.text += "\n" + n.text, this.inlineQueue.pop(), this.inlineQueue[this.inlineQueue.length - 1].src = s.text) : t.push(n), i = r.length !== e.length, e = e.substring(n.raw.length); else if (n = this.tokenizer.text(e)) e = e.substring(n.raw.length), s = t[t.length - 1], s && "text" === s.type ? (s.raw += "\n" + n.raw, s.text += "\n" + n.text, this.inlineQueue.pop(), this.inlineQueue[this.inlineQueue.length - 1].src = s.text) : t.push(n); else if (e) {
+ const t = "Infinite loop on byte: " + e.charCodeAt(0);
+ if (this.options.silent) {
+ console.error(t);
+ break
+ }
+ throw new Error(t)
+ }
+ }
+ return this.state.top = !0, t
+ }
+
+ inline(e, t = []) {
+ return this.inlineQueue.push({src: e, tokens: t}), t
+ }
+
+ inlineTokens(e, t = []) {
+ let n, s, r, i, l, o, a = e;
+ if (this.tokens.links) {
+ const e = Object.keys(this.tokens.links);
+ if (e.length > 0) for (; null != (i = this.tokenizer.rules.inline.reflinkSearch.exec(a));) e.includes(i[0].slice(i[0].lastIndexOf("[") + 1, -1)) && (a = a.slice(0, i.index) + "[" + "a".repeat(i[0].length - 2) + "]" + a.slice(this.tokenizer.rules.inline.reflinkSearch.lastIndex))
+ }
+ for (; null != (i = this.tokenizer.rules.inline.blockSkip.exec(a));) a = a.slice(0, i.index) + "[" + "a".repeat(i[0].length - 2) + "]" + a.slice(this.tokenizer.rules.inline.blockSkip.lastIndex);
+ for (; null != (i = this.tokenizer.rules.inline.anyPunctuation.exec(a));) a = a.slice(0, i.index) + "++" + a.slice(this.tokenizer.rules.inline.anyPunctuation.lastIndex);
+ for (; e;) if (l || (o = ""), l = !1, !(this.options.extensions && this.options.extensions.inline && this.options.extensions.inline.some((s => !!(n = s.call({lexer: this}, e, t)) && (e = e.substring(n.raw.length), t.push(n), !0))))) if (n = this.tokenizer.escape(e)) e = e.substring(n.raw.length), t.push(n); else if (n = this.tokenizer.tag(e)) e = e.substring(n.raw.length), s = t[t.length - 1], s && "text" === n.type && "text" === s.type ? (s.raw += n.raw, s.text += n.text) : t.push(n); else if (n = this.tokenizer.link(e)) e = e.substring(n.raw.length), t.push(n); else if (n = this.tokenizer.reflink(e, this.tokens.links)) e = e.substring(n.raw.length), s = t[t.length - 1], s && "text" === n.type && "text" === s.type ? (s.raw += n.raw, s.text += n.text) : t.push(n); else if (n = this.tokenizer.emStrong(e, a, o)) e = e.substring(n.raw.length), t.push(n); else if (n = this.tokenizer.codespan(e)) e = e.substring(n.raw.length), t.push(n); else if (n = this.tokenizer.br(e)) e = e.substring(n.raw.length), t.push(n); else if (n = this.tokenizer.del(e)) e = e.substring(n.raw.length), t.push(n); else if (n = this.tokenizer.autolink(e)) e = e.substring(n.raw.length), t.push(n); else if (this.state.inLink || !(n = this.tokenizer.url(e))) {
+ if (r = e, this.options.extensions && this.options.extensions.startInline) {
+ let t = 1 / 0;
+ const n = e.slice(1);
+ let s;
+ this.options.extensions.startInline.forEach((e => {
+ s = e.call({lexer: this}, n), "number" == typeof s && s >= 0 && (t = Math.min(t, s))
+ })), t < 1 / 0 && t >= 0 && (r = e.substring(0, t + 1))
+ }
+ if (n = this.tokenizer.inlineText(r)) e = e.substring(n.raw.length), "_" !== n.raw.slice(-1) && (o = n.raw.slice(-1)), l = !0, s = t[t.length - 1], s && "text" === s.type ? (s.raw += n.raw, s.text += n.text) : t.push(n); else if (e) {
+ const t = "Infinite loop on byte: " + e.charCodeAt(0);
+ if (this.options.silent) {
+ console.error(t);
+ break
+ }
+ throw new Error(t)
+ }
+ } else e = e.substring(n.raw.length), t.push(n);
+ return t
+ }
+ }
+
+ class se {
+ options;
+
+ constructor(t) {
+ this.options = t || e.defaults
+ }
+
+ code(e, t, n) {
+ const s = (t || "").match(/^\S*/)?.[0];
+ return e = e.replace(/\n$/, "") + "\n", s ? '' + (n ? e : c(e, !0)) + "
\n" : "" + (n ? e : c(e, !0)) + "
\n"
+ }
+
+ blockquote(e) {
+ return `\n${e}
\n`
+ }
+
+ html(e, t) {
+ return e
+ }
+
+ heading(e, t, n) {
+ return `${e}\n`
+ }
+
+ hr() {
+ return "
\n"
+ }
+
+ list(e, t, n) {
+ const s = t ? "ol" : "ul";
+ return "<" + s + (t && 1 !== n ? ' start="' + n + '"' : "") + ">\n" + e + "\n"
+ }
+
+ listitem(e, t, n) {
+ return `${e}\n`
+ }
+
+ checkbox(e) {
+ return "'
+ }
+
+ paragraph(e) {
+ return `${e}
\n`
+ }
+
+ table(e, t) {
+ return t && (t = `${t}`), "\n"
+ }
+
+ tablerow(e) {
+ return `\n${e}
\n`
+ }
+
+ tablecell(e, t) {
+ const n = t.header ? "th" : "td";
+ return (t.align ? `<${n} align="${t.align}">` : `<${n}>`) + e + `\n`
+ }
+
+ strong(e) {
+ return `${e}`
+ }
+
+ em(e) {
+ return `${e}`
+ }
+
+ codespan(e) {
+ return `${e}`
+ }
+
+ br() {
+ return "
"
+ }
+
+ del(e) {
+ return `${e}`
+ }
+
+ link(e, t, n) {
+ const s = g(e);
+ if (null === s) return n;
+ let r = '", r
+ }
+
+ image(e, t, n) {
+ const s = g(e);
+ if (null === s) return n;
+ let r = `
0 && "paragraph" === n.tokens[0].type ? (n.tokens[0].text = e + " " + n.tokens[0].text, n.tokens[0].tokens && n.tokens[0].tokens.length > 0 && "text" === n.tokens[0].tokens[0].type && (n.tokens[0].tokens[0].text = e + " " + n.tokens[0].tokens[0].text)) : n.tokens.unshift({
+ type: "text",
+ text: e + " "
+ }) : o += e + " "
+ }
+ o += this.parse(n.tokens, i), l += this.renderer.listitem(o, r, !!s)
+ }
+ n += this.renderer.list(l, t, s);
+ continue
+ }
+ case"html": {
+ const e = r;
+ n += this.renderer.html(e.text, e.block);
+ continue
+ }
+ case"paragraph": {
+ const e = r;
+ n += this.renderer.paragraph(this.parseInline(e.tokens));
+ continue
+ }
+ case"text": {
+ let i = r, l = i.tokens ? this.parseInline(i.tokens) : i.text;
+ for (; s + 1 < e.length && "text" === e[s + 1].type;) i = e[++s], l += "\n" + (i.tokens ? this.parseInline(i.tokens) : i.text);
+ n += t ? this.renderer.paragraph(l) : l;
+ continue
+ }
+ default: {
+ const e = 'Token with "' + r.type + '" type was not found.';
+ if (this.options.silent) return console.error(e), "";
+ throw new Error(e)
+ }
+ }
+ }
+ return n
+ }
+
+ parseInline(e, t) {
+ t = t || this.renderer;
+ let n = "";
+ for (let s = 0; s < e.length; s++) {
+ const r = e[s];
+ if (this.options.extensions && this.options.extensions.renderers && this.options.extensions.renderers[r.type]) {
+ const e = this.options.extensions.renderers[r.type].call({parser: this}, r);
+ if (!1 !== e || !["escape", "html", "link", "image", "strong", "em", "codespan", "br", "del", "text"].includes(r.type)) {
+ n += e || "";
+ continue
+ }
+ }
+ switch (r.type) {
+ case"escape": {
+ const e = r;
+ n += t.text(e.text);
+ break
+ }
+ case"html": {
+ const e = r;
+ n += t.html(e.text);
+ break
+ }
+ case"link": {
+ const e = r;
+ n += t.link(e.href, e.title, this.parseInline(e.tokens, t));
+ break
+ }
+ case"image": {
+ const e = r;
+ n += t.image(e.href, e.title, e.text);
+ break
+ }
+ case"strong": {
+ const e = r;
+ n += t.strong(this.parseInline(e.tokens, t));
+ break
+ }
+ case"em": {
+ const e = r;
+ n += t.em(this.parseInline(e.tokens, t));
+ break
+ }
+ case"codespan": {
+ const e = r;
+ n += t.codespan(e.text);
+ break
+ }
+ case"br":
+ n += t.br();
+ break;
+ case"del": {
+ const e = r;
+ n += t.del(this.parseInline(e.tokens, t));
+ break
+ }
+ case"text": {
+ const e = r;
+ n += t.text(e.text);
+ break
+ }
+ default: {
+ const e = 'Token with "' + r.type + '" type was not found.';
+ if (this.options.silent) return console.error(e), "";
+ throw new Error(e)
+ }
+ }
+ }
+ return n
+ }
+ }
+
+ class le {
+ options;
+
+ constructor(t) {
+ this.options = t || e.defaults
+ }
+
+ static passThroughHooks = new Set(["preprocess", "postprocess", "processAllTokens"]);
+
+ preprocess(e) {
+ return e
+ }
+
+ postprocess(e) {
+ return e
+ }
+
+ processAllTokens(e) {
+ return e
+ }
+ }
+
+ class oe {
+ defaults = {
+ async: !1,
+ breaks: !1,
+ extensions: null,
+ gfm: !0,
+ hooks: null,
+ pedantic: !1,
+ renderer: null,
+ silent: !1,
+ tokenizer: null,
+ walkTokens: null
+ };
+ options = this.setOptions;
+ parse = this.#e(ne.lex, ie.parse);
+ parseInline = this.#e(ne.lexInline, ie.parseInline);
+ Parser = ie;
+ Renderer = se;
+ TextRenderer = re;
+ Lexer = ne;
+ Tokenizer = w;
+ Hooks = le;
+
+ constructor(...e) {
+ this.use(...e)
+ }
+
+ walkTokens(e, t) {
+ let n = [];
+ for (const s of e) switch (n = n.concat(t.call(this, s)), s.type) {
+ case"table": {
+ const e = s;
+ for (const s of e.header) n = n.concat(this.walkTokens(s.tokens, t));
+ for (const s of e.rows) for (const e of s) n = n.concat(this.walkTokens(e.tokens, t));
+ break
+ }
+ case"list": {
+ const e = s;
+ n = n.concat(this.walkTokens(e.items, t));
+ break
+ }
+ default: {
+ const e = s;
+ this.defaults.extensions?.childTokens?.[e.type] ? this.defaults.extensions.childTokens[e.type].forEach((s => {
+ const r = e[s].flat(1 / 0);
+ n = n.concat(this.walkTokens(r, t))
+ })) : e.tokens && (n = n.concat(this.walkTokens(e.tokens, t)))
+ }
+ }
+ return n
+ }
+
+ use(...e) {
+ const t = this.defaults.extensions || {renderers: {}, childTokens: {}};
+ return e.forEach((e => {
+ const n = {...e};
+ if (n.async = this.defaults.async || n.async || !1, e.extensions && (e.extensions.forEach((e => {
+ if (!e.name) throw new Error("extension name required");
+ if ("renderer" in e) {
+ const n = t.renderers[e.name];
+ t.renderers[e.name] = n ? function (...t) {
+ let s = e.renderer.apply(this, t);
+ return !1 === s && (s = n.apply(this, t)), s
+ } : e.renderer
+ }
+ if ("tokenizer" in e) {
+ if (!e.level || "block" !== e.level && "inline" !== e.level) throw new Error("extension level must be 'block' or 'inline'");
+ const n = t[e.level];
+ n ? n.unshift(e.tokenizer) : t[e.level] = [e.tokenizer], e.start && ("block" === e.level ? t.startBlock ? t.startBlock.push(e.start) : t.startBlock = [e.start] : "inline" === e.level && (t.startInline ? t.startInline.push(e.start) : t.startInline = [e.start]))
+ }
+ "childTokens" in e && e.childTokens && (t.childTokens[e.name] = e.childTokens)
+ })), n.extensions = t), e.renderer) {
+ const t = this.defaults.renderer || new se(this.defaults);
+ for (const n in e.renderer) {
+ if (!(n in t)) throw new Error(`renderer '${n}' does not exist`);
+ if ("options" === n) continue;
+ const s = n, r = e.renderer[s], i = t[s];
+ t[s] = (...e) => {
+ let n = r.apply(t, e);
+ return !1 === n && (n = i.apply(t, e)), n || ""
+ }
+ }
+ n.renderer = t
+ }
+ if (e.tokenizer) {
+ const t = this.defaults.tokenizer || new w(this.defaults);
+ for (const n in e.tokenizer) {
+ if (!(n in t)) throw new Error(`tokenizer '${n}' does not exist`);
+ if (["options", "rules", "lexer"].includes(n)) continue;
+ const s = n, r = e.tokenizer[s], i = t[s];
+ t[s] = (...e) => {
+ let n = r.apply(t, e);
+ return !1 === n && (n = i.apply(t, e)), n
+ }
+ }
+ n.tokenizer = t
+ }
+ if (e.hooks) {
+ const t = this.defaults.hooks || new le;
+ for (const n in e.hooks) {
+ if (!(n in t)) throw new Error(`hook '${n}' does not exist`);
+ if ("options" === n) continue;
+ const s = n, r = e.hooks[s], i = t[s];
+ le.passThroughHooks.has(n) ? t[s] = e => {
+ if (this.defaults.async) return Promise.resolve(r.call(t, e)).then((e => i.call(t, e)));
+ const n = r.call(t, e);
+ return i.call(t, n)
+ } : t[s] = (...e) => {
+ let n = r.apply(t, e);
+ return !1 === n && (n = i.apply(t, e)), n
+ }
+ }
+ n.hooks = t
+ }
+ if (e.walkTokens) {
+ const t = this.defaults.walkTokens, s = e.walkTokens;
+ n.walkTokens = function (e) {
+ let n = [];
+ return n.push(s.call(this, e)), t && (n = n.concat(t.call(this, e))), n
+ }
+ }
+ this.defaults = {...this.defaults, ...n}
+ })), this
+ }
+
+ setOptions(e) {
+ return this.defaults = {...this.defaults, ...e}, this
+ }
+
+ lexer(e, t) {
+ return ne.lex(e, t ?? this.defaults)
+ }
+
+ parser(e, t) {
+ return ie.parse(e, t ?? this.defaults)
+ }
+
+ #e(e, t) {
+ return (n, s) => {
+ const r = {...s}, i = {...this.defaults, ...r};
+ !0 === this.defaults.async && !1 === r.async && (i.silent || console.warn("marked(): The async option was set to true by an extension. The async: false option sent to parse will be ignored."), i.async = !0);
+ const l = this.#t(!!i.silent, !!i.async);
+ if (null == n) return l(new Error("marked(): input parameter is undefined or null"));
+ if ("string" != typeof n) return l(new Error("marked(): input parameter is of type " + Object.prototype.toString.call(n) + ", string expected"));
+ if (i.hooks && (i.hooks.options = i), i.async) return Promise.resolve(i.hooks ? i.hooks.preprocess(n) : n).then((t => e(t, i))).then((e => i.hooks ? i.hooks.processAllTokens(e) : e)).then((e => i.walkTokens ? Promise.all(this.walkTokens(e, i.walkTokens)).then((() => e)) : e)).then((e => t(e, i))).then((e => i.hooks ? i.hooks.postprocess(e) : e)).catch(l);
+ try {
+ i.hooks && (n = i.hooks.preprocess(n));
+ let s = e(n, i);
+ i.hooks && (s = i.hooks.processAllTokens(s)), i.walkTokens && this.walkTokens(s, i.walkTokens);
+ let r = t(s, i);
+ return i.hooks && (r = i.hooks.postprocess(r)), r
+ } catch (e) {
+ return l(e)
+ }
+ }
+ }
+
+ #t(e, t) {
+ return n => {
+ if (n.message += "\nPlease report this to https://github.com/markedjs/marked.", e) {
+ const e = "An error occurred:
" + c(n.message + "", !0) + "
";
+ return t ? Promise.resolve(e) : e
+ }
+ if (t) return Promise.reject(n);
+ throw n
+ }
+ }
+ }
+
+ const ae = new oe;
+
+ function ce(e, t) {
+ return ae.parse(e, t)
+ }
+
+ ce.options = ce.setOptions = function (e) {
+ return ae.setOptions(e), ce.defaults = ae.defaults, n(ce.defaults), ce
+ }, ce.getDefaults = t, ce.defaults = e.defaults, ce.use = function (...e) {
+ return ae.use(...e), ce.defaults = ae.defaults, n(ce.defaults), ce
+ }, ce.walkTokens = function (e, t) {
+ return ae.walkTokens(e, t)
+ }, ce.parseInline = ae.parseInline, ce.Parser = ie, ce.parser = ie.parse, ce.Renderer = se, ce.TextRenderer = re, ce.Lexer = ne, ce.lexer = ne.lex, ce.Tokenizer = w, ce.Hooks = le, ce.parse = ce;
+ const he = ce.options, pe = ce.setOptions, ue = ce.use, ke = ce.walkTokens, ge = ce.parseInline, fe = ce,
+ de = ie.parse, xe = ne.lex;
+ e.Hooks = le, e.Lexer = ne, e.Marked = oe, e.Parser = ie, e.Renderer = se, e.TextRenderer = re, e.Tokenizer = w, e.getDefaults = t, e.lexer = xe, e.marked = ce, e.options = he, e.parse = fe, e.parseInline = ge, e.parser = de, e.setOptions = pe, e.use = ue, e.walkTokens = ke
+}));
diff --git a/web/showMarkdown.js b/web/showMarkdown.js
new file mode 100644
index 0000000..5b32173
--- /dev/null
+++ b/web/showMarkdown.js
@@ -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);
+ };
+ }
+ },
+});
\ No newline at end of file
diff --git a/youtube_nodes.py b/youtube_nodes.py
index e6f1d9a..7414488 100644
--- a/youtube_nodes.py
+++ b/youtube_nodes.py
@@ -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]