diff --git a/__init__.py b/__init__.py index 47c8a36..c2f2b3c 100644 --- a/__init__.py +++ b/__init__.py @@ -29,7 +29,7 @@ OLD_NODE_CLASS_MAPPINGS = { OLD_NODE_CLASS_MAPPINGS = {v: k for k, v in OLD_NODE_CLASS_MAPPINGS.items()} NODE_CLASS_MAPPINGS = {} - +NODE_DISPLAY_NAME_MAPPINGS = {} # Try importing optional dependencies @@ -44,35 +44,33 @@ else: # if not a unit test print("ETK> Failed to import comfy.samplers, assuming no comfyui skipping SD related nodes") # Load classes from modules in the custom_nodes.EternalKernelLiteGraphNodes package - package_dir = os.path.dirname(os.path.abspath(__file__)) - package_name = "custom_nodes.EternalKernelLiteGraphNodes" - for module_file in os.listdir(package_dir): - if module_file.endswith('.py') and not module_file.startswith('__'): - module_name = package_name + '.' + module_file[:-3] # Removing .py extension - try: - # import directly from the file - full_path = os.path.join(package_dir, module_file) - with open(full_path, 'r') as f: - code = f.read() - exec(code, globals(), locals()) - for class_name in dir(): - class_obj = locals()[class_name] - if isinstance(class_obj, type) and hasattr(class_obj, 'INPUT_TYPES'): - NODE_CLASS_MAPPINGS[class_name] = class_obj + folder = os.path.dirname(os.path.abspath(__file__)) - # module = importlib.import_module(module_name) - # for class_name in dir(module): - # class_obj = getattr(module, class_name) - # if isinstance(class_obj, type) and hasattr(class_obj, 'INPUT_TYPES'): # Checking if it's a class and has INPUT_TYPES attribute - # NODE_CLASS_MAPPINGS[class_name] = class_obj + for file in os.listdir(folder): + if file == "nodes.py": + raise ValueError("nodes.py is not an allowed name in this folder") + if file.endswith(".py") and file[0] != "_": + # import the NODE_CLASS_MAPPINGS from the file + try: + module = importlib.import_module(f".{file[:-3]}", __name__) except ImportError: - if not use_generic_nodes: - print(f"Failed to import {module_name}, it may require comfy.samplers. Skipping...") + try: + module = importlib.import_module(f"{file[:-3]}") # needed for testing? + except ModuleNotFoundError: + module = importlib.import_module(f".{file[:-3]}", __name__) + + if hasattr(module, "NODE_CLASS_MAPPINGS"): + mappings = module.NODE_CLASS_MAPPINGS + NODE_CLASS_MAPPINGS.update(mappings) + print(f"ETK: updated with {mappings}") + + if hasattr(module, "NODE_DISPLAY_NAME_MAPPINGS"): + NODE_DISPLAY_NAME_MAPPINGS.update(module.NODE_DISPLAY_NAME_MAPPINGS) # use OLD_NODE_CLASS_MAPPINGS to add the old names to the new nodes new_node_class_mappings = {} for k, v in OLD_NODE_CLASS_MAPPINGS.items(): - if NODE_CLASS_MAPPINGS[k]: + if k in NODE_CLASS_MAPPINGS: new_node_class_mappings[v] = NODE_CLASS_MAPPINGS[k] # mark it for removeal NODE_CLASS_MAPPINGS[k] = None diff --git a/audio.py b/audio.py new file mode 100644 index 0000000..c28ee2d --- /dev/null +++ b/audio.py @@ -0,0 +1,73 @@ +import copy + +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} + + +def ETK_audio_base(cls): + cls.CATEGORY = "ETK/audio" + NODE_CLASS_MAPPINGS[cls.__name__] = cls + return cls + + +@ETK_audio_base +class GetAudioMetaData: + """ Returns a dictionary of audio metadata """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "path": ("STRING", {"multiline": False, "default": ""}), + } + + } + + RETURN_TYPES = ("DICT",) + FUNCTION = 'get_audio_metadata' + CATEGORY = 'ETK/audio' + + def get_audio_metadata(self, **kwargs): + import audioread + kwargs = copy.deepcopy(kwargs) + path = kwargs.get('path') + metadata = {} + with audioread.audio_open(path) as f: + metadata['duration'] = f.duration + metadata['samplerate'] = f.samplerate + metadata['channels'] = f.channels + + return (metadata,) + + +@ETK_audio_base +class AudioFileListFromPathPattern: + """ Returns a list of audio files from a path pattern """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "path": ("STRING", {"multiline": False, "default": ""}), + } + + } + + RETURN_TYPES = ("LIST", "LIST", "LIST",) + RETURN_NAMES = ("list_of_fpaths", "list_of_bytes", "list_of_metadata") + FUNCTION = 'get_audio_file_list_from_path_pattern' + CATEGORY = 'ETK/audio' + + def get_audio_file_list_from_path_pattern(self, **kwargs): + import glob + kwargs = copy.deepcopy(kwargs) + path = kwargs.get('path') + l_fpaths = glob.glob(path) + l_bytes = [] + l_metadata = [] + for fpath in l_fpaths: + with open(fpath, "rb") as f: + l_bytes.append(f.read()) + l_metadata.append(GetAudioMetaData().get_audio_metadata(path=fpath)[0]) + + return (l_fpaths, l_bytes, l_metadata,) diff --git a/descript.py b/descript.py index d40560c..baefb53 100644 --- a/descript.py +++ b/descript.py @@ -1,6 +1,7 @@ import copy NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} def ETK_descript_base(cls): @@ -9,6 +10,7 @@ def ETK_descript_base(cls): NODE_CLASS_MAPPINGS[cls.__name__] = cls return cls + @ETK_descript_base class DecodeTranscriptNode: """ @@ -22,13 +24,13 @@ class DecodeTranscriptNode: return { "required": { "string": ("STRING", {"multiline": True, - "default": "" - }), + "default": "" + }), } } - RETURN_TYPES = ("LIST","LIST","LIST") - RETURN_NAMES = ("speakers","durations","paragraphs") + RETURN_TYPES = ("LIST", "LIST", "LIST") + RETURN_NAMES = ("speakers", "durations", "paragraphs") OUTPUT_NODE = True def func(self, string): @@ -51,23 +53,23 @@ class DecodeTranscriptNode: continue # first find the first ] and split on that part1 = para.find("]") - time_part = para[:part1+1] - remainder = para[part1+1:] + time_part = para[:part1 + 1] + remainder = para[part1 + 1:] # now find first : in the remainder part2 = remainder.find(":") - speaker_part = remainder[:part2+1].replace(":","").strip() - paragraph_part = remainder[part2+1:] + speaker_part = remainder[:part2 + 1].replace(":", "").strip() + paragraph_part = remainder[part2 + 1:] speakers.append(speaker_part.strip()) # now just remove [ and ] from time_part - time_part = time_part.replace("[","").replace("]","").strip() + time_part = time_part.replace("[", "").replace("]", "").strip() - #split it on : + # split it on : time_part = time_part.split(":") # now convert to seconds - time_part = int(time_part[0])*60*60 + int(time_part[1])*60 + int(time_part[2]) + time_part = int(time_part[0]) * 60 * 60 + int(time_part[1]) * 60 + int(time_part[2]) # save here for post-processing later durations.append(time_part) @@ -76,13 +78,10 @@ class DecodeTranscriptNode: import copy start_seconds = copy.copy(durations) - durations=[] - for i in range(len(start_seconds)-1): - durations.append(start_seconds[i+1] - start_seconds[i]) + durations = [] + for i in range(len(start_seconds) - 1): + durations.append(start_seconds[i + 1] - start_seconds[i]) # now append the last value - durations.append(int(sum(durations)/len(durations))) - - return (speakers,durations,paragraphs,) - - + durations.append(int(sum(durations) / len(durations))) + return (speakers, durations, paragraphs,) diff --git a/functional.py b/functional.py index 5d951df..a4c3e2f 100644 --- a/functional.py +++ b/functional.py @@ -1,9 +1,12 @@ import hashlib +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} _list = {"list": ("LIST",)} _func = {"func": ("FUNC", {"default": None})} _func_extra = {"func_extra": ("FUNC",)} _text = {"code": ("STRING", {"multiline": True, "default": "y=x()"})} +_image = {"image": ("IMAGE",)} _string = {"string": ("STRING", {"default": None})} optional = lambda x: {"optional": {k: v for k, v in x.items()}} @@ -12,6 +15,14 @@ both = lambda a, b: {**a, **b} many = lambda *dicts: {k: v for d in dicts for k, v in d.items()} # non tested +def ETK_functional_base(cls): + cls.FUNCTION = "func" + cls.CATEGORY = "ETK/functional" + NODE_CLASS_MAPPINGS[cls.__name__] = cls + return cls + + +@ETK_functional_base class FuncBase: """ Base class for func nodes @@ -68,6 +79,7 @@ class FuncBase: return (self.func,) +@ETK_functional_base class FuncRender: """ as input takes a function and a string of code @@ -122,6 +134,7 @@ class FuncRender: return (locals["y_float"], locals["y_string"], locals["y_int"], locals["y_list"],) +@ETK_functional_base class FuncRenderImage: def __init__(self): from custom_nodes.EternalKernelLiteGraphNodes.image import common_ksampler, TinyTxtToImg @@ -157,6 +170,7 @@ class FuncRenderImage: return (locals["y"],) +@ETK_functional_base class ExecWidget: import torch @@ -271,6 +285,7 @@ class ExecWidget: return m.digest().hex() +@ETK_functional_base class FuncListToList: """runs eval on the given text, with a list as input and output""" @@ -306,6 +321,7 @@ class FuncListToList: return (my_locals["y_list"],) +@ETK_functional_base class FuncStrToStr: """runs eval on the given text, with a string as input and outputs a string""" @@ -323,6 +339,7 @@ class FuncStrToStr: def func(self, **kwargs): text = kwargs.get("string", None) + x = text code = kwargs.get("code", None) my_globals = globals() my_locals = locals() @@ -332,6 +349,7 @@ class FuncStrToStr: return (my_locals["y"],) +@ETK_functional_base class FuncStrToList: def __init__(self): pass @@ -348,6 +366,7 @@ class FuncStrToList: def func(self, **kwargs): text = kwargs.get("string", None) code = kwargs.get("code", None) + x = text my_globals = globals() my_locals = locals() @@ -356,6 +375,33 @@ class FuncStrToList: return (my_locals["y_list"],) +@ETK_functional_base +class FuncListToStr: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + req = required(both(_list, _text)) + return req + + CATEGORY = "ETK/func" + RETURN_TYPES = ("STRING",) + FUNCTION = "func" + + def func(self, **kwargs): + text = kwargs.get("list", None) + code = kwargs.get("code", None) + x = text + my_globals = globals() + my_locals = locals() + + exec(code, my_globals, my_locals) + + return (my_locals["y"],) + + +@ETK_functional_base class GetFirstCodeBlock(): """ returns the first code block in the text which is surrounded with ``` or ```python or diff --git a/git.py b/git.py index 114447b..78c1f74 100644 --- a/git.py +++ b/git.py @@ -1,10 +1,20 @@ from git import Repo, InvalidGitRepositoryError import os import copy - import re +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} + +def ETK_git_base(cls): + cls.FUNCTION = "func" + cls.CATEGORY = "ETK/git" + NODE_CLASS_MAPPINGS[cls.__name__] = cls + return cls + + +@ETK_git_base class PublicGitRepo: @classmethod def INPUT_TYPES(cls): @@ -33,6 +43,7 @@ def validate_git_repo_url(self, **kwargs): return match is not None and match.group() == repo_url +@ETK_git_base class CloneRepoNode: @classmethod def INPUT_TYPES(cls): @@ -45,7 +56,7 @@ class CloneRepoNode: } RETURN_TYPES = ("STRING", "STRING", "LIST",) - RETURN_NAMES = ("path_passthrough","file_structure", "file_list",) + RETURN_NAMES = ("path_passthrough", "file_structure", "file_list",) FUNCTION = 'clone_repo' CATEGORY = 'ETK/git' @@ -61,6 +72,7 @@ class CloneRepoNode: return ("".join(ld), ld,) +@ETK_git_base class PullRepoNode: @classmethod def INPUT_TYPES(cls): @@ -89,6 +101,7 @@ def pull_repo(self, **kwargs): # origin.pull() +@ETK_git_base class PushRepoNode: @classmethod def INPUT_TYPES(cls): @@ -114,6 +127,7 @@ def push_repo(self, **kwargs): origin.push() +@ETK_git_base class SendCommitNode: @classmethod def INPUT_TYPES(cls): @@ -139,6 +153,7 @@ def send_commit(self, **kwargs): repo.index.commit(commit_message) +@ETK_git_base class CreateRepoNode: @classmethod def INPUT_TYPES(cls): diff --git a/image.py b/image.py index 0d795be..e406e58 100644 --- a/image.py +++ b/image.py @@ -1,5 +1,10 @@ import os +try: + from utils import etk_deep_copy as deepcopy +except ImportError as e: + from custom_nodes.EternalKernelLiteGraphNodes.utils import etk_deep_copy as deepcopy + testing = os.environ.get("ETERNAL_KERNEL_LITEGRAPH_NODES_TEST", None) if testing == "True": testing = True @@ -55,6 +60,15 @@ all image inputs are (B,W,H,C) avoid numpy and PIL as much as possible """ +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 get_fonts(): @@ -137,24 +151,26 @@ def torch_image_show(image): # try to rearrange the channels using np image = np.moveaxis(image, 0, -1) - #image.permute(1, 2, 0) + # image.permute(1, 2, 0) image = Image.fromarray(image) # Create a PIL image image.show() +@ETK_image_base class TinyTxtToImg: """small text to image generator""" share_clip = None share_mdl = None share_vae = None + def __init__(self): + print("ETK> TinyTxtToImg init") import random self.mdl = None self.clp = None self.vae = None self.vision = None - self.seed = random.randint(0, 2 ** 32 - 1) self.steps = 10 self.cfg = 8 self.sampler_name = comfy.samplers.KSampler.SAMPLERS[0] @@ -168,7 +184,6 @@ class TinyTxtToImg: self.latent_image = None self.samples = None - @classmethod def INPUT_TYPES(cls): return {"required": @@ -179,7 +194,7 @@ class TinyTxtToImg: "token_normalization": (["none", "mean", "length", "length+mean"],), "weight_interpretation": (["comfy", "A1111", "compel", "comfy++"],), "ckpt_name": (folder_paths.get_filename_list("checkpoints"),), - + "one_seed": ([True, False], {"default": False}), }, "optional": { @@ -194,29 +209,44 @@ class TinyTxtToImg: CATEGORY = "ETK" - RETURN_TYPES = ("IMAGE", "FUNC",) + RETURN_TYPES = ("IMAGE", "FUNC", "LATENT", "CONDITIONING", "CONDITIONING",) + RETURN_NAMES = ("image", "FUNC", "latent", "positive", "negative",) FUNCTION = "tinytxt2img" - def tinytxt2img(self, prompt, neg_prompt, name="tinytxt2img", overrides: str = "", - clip_encoder="comfy -ignore below", - token_normalization="length+mean", - weight_interpretation="comfy++", - FUNC: callable = None, - ckpt_name="v1-5-pruned-emaonly.safetensors", - ): + def tinytxt2img(self, **kwargs): """ use the imports from nodes to generate an image from text """ import random import json - if not hasattr( TinyTxtToImg, "share_mdl"): + # kwargs = deepcopy(kwargs) + + prompt = kwargs.get("prompt", "") + neg_prompt = kwargs.get("neg_prompt", "") + clip_encoder = kwargs.get("clip_encoder", "comfy -ignore below") + token_normalization = kwargs.get("token_normalization", "length+mean") + weight_interpretation = kwargs.get("weight_interpretation", "comfy++") + ckpt_name = kwargs.get("ckpt_name", "v1-5-pruned-emaonly.safetensors") + overrides = kwargs.get("overrides", "") + name = kwargs.get("name", "tinytxt2img") + FUNC = kwargs.get("FUNC", None) + self.latent_image = kwargs.get("latent", None) + self.pos_cond = kwargs.get("pos_cond", None) + self.neg_cond = kwargs.get("neg_cond", None) + self.seed = kwargs.get("seed", random.randint(0, 2 ** 32 - 1)) + self.one_seed_per_batch = kwargs.get("one_seed", False) + self.denoise = kwargs.get("denoise", 1.0) + + if not hasattr(TinyTxtToImg, "share_mdl"): TinyTxtToImg.share_mdl = None + TinyTxtToImg.share_mdl_ckpt_name = None if TinyTxtToImg.share_mdl is None: self.mdl, self.clp, self.vae, self.vision = \ 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 @@ -224,12 +254,25 @@ class TinyTxtToImg: CSL.share_vision = self.vision else: - self.mdl, self.clp, self.vae = (TinyTxtToImg.share_mdl, - TinyTxtToImg.share_clip, - TinyTxtToImg.share_vae - ) + # check if the checkpoint is the same + if TinyTxtToImg.share_mdl_ckpt_name != ckpt_name: + self.mdl, self.clp, self.vae, self.vision = \ + 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 + else: + + self.mdl, self.clp, self.vae = (TinyTxtToImg.share_mdl, + TinyTxtToImg.share_clip, + TinyTxtToImg.share_vae + ) - self.seed = random.randint(0, 2 ** 32 - 1) self.steps = 10 self.cfg = 8 self.sampler_name = comfy.samplers.KSampler.SAMPLERS[0] @@ -239,8 +282,8 @@ class TinyTxtToImg: self.width = 512 self.height = 512 self.batch_size = 1 - self.denoise = 1.0 - self.one_seed_per_batch = False + if not self.denoise: + self.denoise = 1.0 if overrides: if overrides != "": @@ -256,45 +299,61 @@ class TinyTxtToImg: ## TODO: probably need to prepare the execution better to make sure that applying the incomming func can change everything if FUNC is not None: self.__dict__.update(FUNC(self.__dict__)) - if clip_encoder == "comfy -ignore below": - self.pos_cond = CLIPTextEncode.encode(None, self.clp, self.positive)[0] - self.neg_cond = CLIPTextEncode.encode(None, self.clp, self.negative)[0] - elif clip_encoder == "advanced": - try: - self.pos_cond = CLIPTextEncodeAdvanced.encode(None, - self.clp, - self.positive, - token_normalization, - weight_interpretation)[0] - self.neg_cond = CLIPTextEncodeAdvanced.encode(None, - self.clp, - self.negative, - token_normalization, - weight_interpretation)[0] - except Exception as e: - print(e) - raise ValueError("advanced clip encoder failed") + if not self.pos_cond: + if clip_encoder == "comfy -ignore below": + self.clp_encode = lambda x: CLIPTextEncode.encode(None, self.clp, x)[0] + self.pos_cond = self.clp_encode(self.positive) + self.neg_cond = self.clp_encode(self.negative) - self.latent_image = EmptyLatentImage.generate(None, self.width, self.height, self.batch_size)[0] + elif clip_encoder == "advanced": + self.clp_encode = lambda x: CLIPTextEncodeAdvanced.encode(None, + self.clp, + x, + token_normalization, + weight_interpretation)[0] + try: + self.pos_cond = self.clp_encode(self.positive) + self.neg_cond = self.clp_encode(self.negative) - samples = KSampler.sample(None, - self.mdl, - self.seed, self.steps, self.cfg, self.sampler_name, self.scheduler, - self.pos_cond, self.neg_cond, - self.latent_image, - denoise=self.denoise, - one_seed_per_batch=self.one_seed_per_batch)[0] + except Exception as e: + print(e) + raise ValueError("advanced clip encoder failed") - image = VAEDecode.decode(None, self.vae, samples)[0] + if not self.latent_image: + self.latent_image = EmptyLatentImage.generate(None, self.width, self.height, self.batch_size)[0] + + self.KSampler = KSampler() + self.samples = self.KSampler.sample( + model=self.mdl, + seed=self.seed, + steps=self.steps, + cfg=self.cfg, + sampler_name=self.sampler_name, + scheduler=self.scheduler, + latent_image=self.latent_image, + pos_cond=self.pos_cond, + neg_cond=self.neg_cond, + positive=self.pos_cond, + negative=self.neg_cond, + denoise=self.denoise, + one_seed_per_batch=self.one_seed_per_batch)[0] + # self.vae.offload_device = torch.device("cuda") + image = VAEDecode.decode(None, self.vae, self.samples)[0] image = image.detach().cpu() self.FUNC = self.tinytxt2img - self.ARGS = (prompt, neg_prompt, name, overrides, clip_encoder, token_normalization, weight_interpretation, FUNC) - return (image, - self - ,) + self.ARGS = kwargs + ret = (image, + self, + self.samples, + self.pos_cond, + self.neg_cond + ,) + # ret = deepcopy(ret) + return ret +@ETK_image_base class PreviewImageTest(SaveImage): def __init__(self): self.output_dir = folder_paths.get_temp_directory() @@ -327,8 +386,7 @@ class PreviewImageTest(SaveImage): return my_ret - - +@ETK_image_base class PromptTemplate: """replaces the text in the given text string with a given other text at some key positions""" @@ -354,6 +412,7 @@ class PromptTemplate: return (text.replace(key1, replacement),) +@ETK_image_base class LoadImage: @classmethod def INPUT_TYPES(s): @@ -452,6 +511,7 @@ class LoadImage: return m.digest().hex() +@ETK_image_base class SelectFromRGBSimilarity: """ selects pixel positions based on similarity to a given color @@ -556,6 +616,7 @@ class SelectFromRGBSimilarity: return (out,) +@ETK_image_base class Quantize: def __init__(self): pass @@ -626,6 +687,7 @@ class Quantize: return (result,) +@ETK_image_base class select_from_batch: """selects a single image from a batch""" @@ -686,6 +748,7 @@ class select_from_batch: return (oimg, olat,) +@ETK_image_base class ImageDistanceMask: def __init__(self): pass @@ -773,6 +836,7 @@ class ImageDistanceMask: return (dist,) +@ETK_image_base class rgba_lower_clip: """Clip RGBA values below a threshold""" @@ -844,6 +908,65 @@ class rgba_lower_clip: return (image,) +@ETK_image_base +class ImageStackBlendByMask: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "stack a": ("IMAGE",), + "stack b": ("IMAGE",), + "img mask": ("IMAGE",), + + }, + } + + RETURN_TYPES = ("IMAGE", "IMAGE",) + RETURN_NAMES = ("image", "image rgba") + FUNCTION = "stack_mask" + + CATEGORY = "ETK/Image" + + def stack_mask(self, **kwargs): + import torch + # use torch to blend two images by a mask + # stacks are already (B,W,H,C) + + stack_a = kwargs["stack a"] + stack_b = kwargs["stack b"] + img_mask = kwargs["img mask"] + # check if they all have alpha + if stack_a.shape[3] == 3: + alpha = torch.ones_like(stack_a[:, :, :, 0]) + stack_a = torch.cat((stack_a, alpha.unsqueeze(3)), dim=3) + + if stack_b.shape[3] == 3: + alpha = torch.ones_like(stack_b[:, :, :, 0]) + stack_b = torch.cat((stack_b, alpha.unsqueeze(3)), dim=3) + + # allocate new image + new_image = torch.zeros_like(stack_a) + + # blend rgb channels + + for i in range(3): + new_image[:, :, :, i] = \ + (stack_a[:, :, :, i] * img_mask[:, :, :, 0]) \ + + \ + (stack_b[:, :, :, i] * (1 - img_mask[:, :, :, 0])) + + # alpha should be the mean of the two alphas + new_image[:, :, :, 3] = (stack_a[:, :, :, 3] + stack_b[:, :, :, 3]) / 2 + rgba = new_image + rgb = new_image[:, :, :, 0:3] + + return (rgb, rgba,) + + +@ETK_image_base class rgba_upper_clip: """Clip RGBA values above a threshold""" @@ -915,6 +1038,7 @@ class rgba_upper_clip: return (image,) +@ETK_image_base class RGBA_MOD: """provide sliders to modify RGBA values""" @@ -975,6 +1099,7 @@ class RGBA_MOD: return (image,) +@ETK_image_base class ImageBC: """Normalize an images brightness,droping lower and higher values of n perctile""" @@ -1046,6 +1171,7 @@ class ImageBC: return (image,) +@ETK_image_base class PadToMatch: """ a node that will pad the smaller of two images to match the size of the larger one """ @@ -1124,6 +1250,7 @@ class PadToMatch: return (padded1, padded2) +@ETK_image_base class StackImages: """ a node that will stack two images on top of each other """ @@ -1139,6 +1266,9 @@ class StackImages: "optional": { "image2": ("IMAGE",), "number_of_images": ("INT", {"default": 2, "min": 2, "max": 100}), + # an optionbox with "all" or "upto", model after "one seed per batch" in the sampler + "mode": (["all", "upto"], {"default": "all"}), + }, } RETURN_TYPES = ("IMAGE",) @@ -1146,7 +1276,7 @@ class StackImages: CATEGORY = "ETK/image" - def stackme(self, image1, image2=None, number_of_images=2): + def stackme(self, image1, image2=None, number_of_images=2, mode="all"): """ stacks two images on top of each other images are (B, W, H, C) torch tensors @@ -1159,15 +1289,21 @@ class StackImages: else: stacked = torch.cat((image1, image1), dim=0) - if number_of_images > 2: - while stacked.shape[0] < number_of_images: - stacked = torch.cat((stacked, stacked), dim=0) - print(stacked.shape) - out = stacked[:number_of_images, :, :, :] - print(out.shape) + if mode == "all": + if number_of_images > 2: + while stacked.shape[0] < number_of_images: + stacked = torch.cat((stacked, stacked), dim=0) + print(stacked.shape) + out = stacked[:number_of_images, :, :, :] + else: + # stack them all together regardless of the first dimension + out = torch.cat((image1, image2), dim=0) + print(stacked.shape) + return (out,) +@ETK_image_base class rgba_merge: """merges the channels of an RGBA image""" @@ -1215,6 +1351,7 @@ class rgba_merge: return (torch.cat((r, g, b, a), dim=3),) +@ETK_image_base class rgba_split: """splits the channels of an RGBA image""" @@ -1266,6 +1403,73 @@ def get_image_size(image): return (width * height) +@ETK_image_base +class KSamplerCreateDataset: + """ use KSampler to create a dataset """ + + @classmethod + def INPUT_TYPES(cls): + inps = KSampler.INPUT_TYPES() + # make sure to set a default + inps["required"]["dataset_name"] = ("STRING", {"default": "dataset"}) + return inps + + CATEGORY = "ETK/training" + + RETURN_TYPES = ("LATENT", "STRING",) + RETURN_NAMES = ("LATENT", "dataset_path",) + + FUNCTION = "create_dataset" + + def create_dataset(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, + denoise, one_seed_per_batch, dataset_name): + """ + creates a dataset using KSampler + :param model: + :param seed: + :param steps: + :param cfg: + :param sampler_name: + :param scheduler: + :param positive: + :param negative: + :param latent_image: + :param denoise: + :param dataset_name: + :return: + """ + # create the sampler + sampler = KSampler() + + x = {"model": model, "seed": seed, "steps": steps, "cfg": cfg, "sampler_name": sampler_name, + "scheduler": scheduler, "positive": positive, "negative": negative, "latent_image": latent_image, + "denoise": denoise, "one_seed_per_batch": one_seed_per_batch} + + # now save x + # just use torch .save? + import torch + # create a random name + import hashlib + import time + + # create a random name + name = hashlib.md5(str(time.time()).encode()).hexdigest()[0:8] + name = dataset_name + "_" + name + + o = folder_paths.get_output_directory() + import os + + dataset_path = os.path.join(o, name) + torch.save(x, dataset_path + "_x.pt") + + results = sampler.sample(**x) + # save the results + torch.save(results[0], dataset_path + "_y.pt") + # return the dataset path + return (results[0], dataset_path,) + + +@ETK_image_base class KSampler: @classmethod def INPUT_TYPES(s): @@ -1283,20 +1487,38 @@ class KSampler: "one_seed_per_batch": ([True, False], {"default": False}), }} - RETURN_TYPES = ("LATENT",) + RETURN_TYPES = ("LATENT", "FUNC",) FUNCTION = "sample" CATEGORY = "sampling" - def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0, - one_seed_per_batch=False): - try: - return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, - denoise=denoise,one_seed_per_batch=one_seed_per_batch) - except TypeError as e: - return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, - denoise=denoise) + def sample(self, **kwargs): + kwargs = deepcopy(kwargs) + + model = kwargs.get("model", None) + seed = kwargs.get("seed", None) + steps = kwargs.get("steps", None) + cfg = kwargs.get("cfg", None) + sampler_name = kwargs.get("sampler_name", None) + scheduler = kwargs.get("scheduler", None) + positive = kwargs.get("positive", None) + negative = kwargs.get("negative", None) + latent_image = kwargs.get("latent_image", None) + 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 + + try: + ret = common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, + denoise=denoise, one_seed_per_batch=one_seed_per_batch) + except TypeError as e: + ret = common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, + denoise=denoise) + ret = (ret[0], self,) + return ret from PIL import Image, ImageDraw, ImageFont @@ -1304,6 +1526,7 @@ import torch import numpy as np +@ETK_image_base class TextRender: """ renders text to an image @@ -1357,6 +1580,7 @@ class TextRender: CATEGORY = "ETK/text" RETURN_TYPES = ("IMAGE", "IMAGE",) + RETURN_NAMES = ("image", "image rgba",) FUNCTION = "render_text" @@ -1476,6 +1700,70 @@ class TextRender: return y_offset +@ETK_image_base +class StripAlphaChannel: + """ + This function strips the alpha channel from the image. + """ + + @classmethod + def INPUT_TYPES(cls): + reqd = {"required": {"image": ("IMAGE",)}} + return reqd + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + + FUNCTION = "strip_alpha_channel" + + def strip_alpha_channel(self, image): + """ + This function strips the alpha channel from the image. + + >>> sac = StripAlphaChannel() + >>> result = sac.strip_alpha_channel(image) + >>> torch_image_show(result[0]) + """ + + return (deepcopy(image[..., :3]),) + + +@ETK_image_base +class FuncImageStackToImageStack: + """ runs exec on giving the user the input image stack and the ability to define the output image stack""" + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + req = {"required": { + "image": ("IMAGE",), + "code": ("STRING", {"multiline": True, "default": "y=x()"}) + }} + return req + + CATEGORY = "ETK/func" + RETURN_TYPES = ("IMAGE",) + FUNCTION = "func" + + def func(self, **kwargs): + func = kwargs.get("func", None) + code = kwargs.get("code", None) + my_globals = globals() + my_locals = locals() + + image_x = kwargs.get("image", None) + gs = image_x.mean(dim=3).repeat(1, 1, 1, 3) + + my_locals["x"] = func + my_locals["x_image"] = kwargs.get("image", None) + + exec(code, my_globals, my_locals) + + return (my_locals["y_image"],) + + if __name__ == "__main__": # test TextRender tr = TextRender() @@ -1483,8 +1771,10 @@ if __name__ == "__main__": torch_image_show(result) print(result.shape) + def main(): pass + if __name__ == "__main__": main() diff --git a/modules/youtube.py b/modules/youtube.py new file mode 100644 index 0000000..9762711 --- /dev/null +++ b/modules/youtube.py @@ -0,0 +1,193 @@ +import httplib2 +import os +import random +import sys +import time + +from apiclient.discovery import build +from apiclient.errors import HttpError +from apiclient.http import MediaFileUpload +from oauth2client.client import flow_from_clientsecrets +from oauth2client.file import Storage +from oauth2client.tools import argparser, run_flow + +# Explicitly tell the underlying HTTP transport library not to retry, since +# we are handling retry logic ourselves. +httplib2.RETRIES = 1 + +# Maximum number of times to retry before giving up. +MAX_RETRIES = 10 + +# Always retry when these exceptions are raised. +RETRIABLE_EXCEPTIONS = (httplib2.HttpLib2Error, IOError) + +# Always retry when an apiclient.errors.HttpError with one of these status +# codes is raised. +RETRIABLE_STATUS_CODES = [500, 502, 503, 504] + +# The CLIENT_SECRETS_FILE variable specifies the name of a file that contains +# the OAuth 2.0 information for this application, including its client_id and +# client_secret. You can acquire an OAuth 2.0 client ID and client secret from +# the Google API Console at +# https://console.cloud.google.com/. +# Please ensure that you have enabled the YouTube Data API for your project. +# For more information about using OAuth2 to access the YouTube Data API, see: +# https://developers.google.com/youtube/v3/guides/authentication +# For more information about the client_secrets.json file format, see: +# https://developers.google.com/api-client-library/python/guide/aaa_client_secrets +try: + import comfy +except ModuleNotFoundError: + print("comfy not found, using fallback") + comfy = lambda: None + comfy.__path__ = lambda: None + comfy.__path__._path = ["."] + +comfy_dir = comfy.__path__._path[0] +# the base dir we need is one level down use os.path.dirname to get it +base_dir = os.path.dirname(comfy_dir) +# now we can use os.path.join to get the full path to the file +CLIENT_SECRETS_FILE = os.path.join(base_dir, "google-credentials.json") + +# This OAuth 2.0 access scope allows an application to upload files to the +# authenticated user's YouTube channel, but doesn't allow other types of access. +YOUTUBE_UPLOAD_SCOPE = "https://www.googleapis.com/auth/youtube.upload" +YOUTUBE_API_SERVICE_NAME = "youtube" +YOUTUBE_API_VERSION = "v3" + +# This variable defines a message to display if the CLIENT_SECRETS_FILE is +# missing. +MISSING_CLIENT_SECRETS_MESSAGE = """ +WARNING: Please configure OAuth 2.0 + +To make this sample run you will need to populate the client_secrets.json file +found at: + + %s + +with information from the API Console +https://console.cloud.google.com/ + +For more information about the client_secrets.json file format, please visit: +https://developers.google.com/api-client-library/python/guide/aaa_client_secrets +""" % os.path.abspath(os.path.join(os.path.dirname(__file__), + CLIENT_SECRETS_FILE)) + +VALID_PRIVACY_STATUSES = ("public", "private", "unlisted") + + +def get_authenticated_service(args): + flow = flow_from_clientsecrets(CLIENT_SECRETS_FILE, + scope=YOUTUBE_UPLOAD_SCOPE, + message=MISSING_CLIENT_SECRETS_MESSAGE) + + storage = Storage("%s-oauth2.json" % sys.argv[0]) + credentials = storage.get() + + if credentials is None or credentials.invalid: + credentials = run_flow(flow, storage, args) + + return build(YOUTUBE_API_SERVICE_NAME, YOUTUBE_API_VERSION, + http=credentials.authorize(httplib2.Http())) + + +def initialize_upload(youtube, options): + tags = None + if options.keywords: + tags = options.keywords.split(",") + + body = dict( + snippet=dict( + title=options.title, + description=options.description, + tags=tags, + categoryId=options.category + ), + status=dict( + privacyStatus=options.privacyStatus + ) + ) + + # Call the API's videos.insert method to create and upload the video. + insert_request = youtube.videos().insert( + part=",".join(body.keys()), + body=body, + # The chunksize parameter specifies the size of each chunk of data, in + # bytes, that will be uploaded at a time. Set a higher value for + # reliable connections as fewer chunks lead to faster uploads. Set a lower + # value for better recovery on less reliable connections. + # + # Setting "chunksize" equal to -1 in the code below means that the entire + # file will be uploaded in a single HTTP request. (If the upload fails, + # it will still be retried where it left off.) This is usually a best + # practice, but if you're using Python older than 2.6 or if you're + # running on App Engine, you should set the chunksize to something like + # 1024 * 1024 (1 megabyte). + media_body=MediaFileUpload(options.file, chunksize=-1, resumable=True) + ) + + resumable_upload(insert_request) + + +# This method implements an exponential backoff strategy to resume a +# failed upload. + + +def resumable_upload(insert_request): + response = None + error = None + retry = 0 + while response is None: + try: + print("Uploading file...") + status, response = insert_request.next_chunk() + if response is not None: + if 'id' in response: + print("Video id '%s' was successfully uploaded." % + response['id']) + else: + exit("The upload failed with an unexpected response: %s" % response) + except HttpError as e: + if e.resp.status in RETRIABLE_STATUS_CODES: + error = "A retriable HTTP error %d occurred:\n%s" % (e.resp.status, + e.content) + else: + raise + except RETRIABLE_EXCEPTIONS as e: + error = "A retriable error occurred: %s" % e + + if error is not None: + print(error) + retry += 1 + if retry > MAX_RETRIES: + exit("No longer attempting to retry.") + + max_sleep = 2 ** retry + sleep_seconds = random.random() * max_sleep + print("Sleeping %f seconds and then retrying..." % sleep_seconds) + time.sleep(sleep_seconds) + + +if __name__ == '__main__': + argparser.add_argument("--file", required=True, + help="Video file to upload") + argparser.add_argument("--title", help="Video title", default="Test Title") + argparser.add_argument("--description", help="Video description", + default="Test Description") + argparser.add_argument("--category", default="22", + help="Numeric video category. " + + "See https://developers.google.com/youtube/v3/docs/videoCategories/list") + argparser.add_argument("--keywords", help="Video keywords, comma separated", + default="") + argparser.add_argument("--privacyStatus", choices=VALID_PRIVACY_STATUSES, + default=VALID_PRIVACY_STATUSES[0], help="Video privacy status.") + args = argparser.parse_args() + + if not os.path.exists(args.file): + exit("Please specify a valid file using the --file= parameter.") + + youtube = get_authenticated_service(args) + try: + initialize_upload(youtube, args) + except HttpError as e: + print("An HTTP error %d occurred:\n%s" % (e.resp.status, e.content)) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..55f3ca0 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,5 @@ +google-api-python-client +google-auth-oauthlib +google-auth-httplib2 +oauth2client +httplib2 \ No newline at end of file diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 0000000..703f24b --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,41 @@ +import torch +import numpy as np +from custom_nodes.EternalKernelLiteGraphNodes.utils import etk_deep_copy + + +def test_etk_deep_copy(): + original = { + "list": [1, 2, 3, [4, 5, 6], {"inner_list": [7, 8, 9]}], + "dict": {"inner_dict": {"key": "value"}} + } + torch_test = {"torch_tensor": torch.tensor([1, 2, 3])} + np_test = {"np_array": np.array([1, 2, 3])} + + test(original) + test(torch_test) + test(np_test) + + +def test(original): + # Creating a deep copy + copy = etk_deep_copy(original) + + # Modifying the original structure + if "list" in original: + original["list"][0] = 100 + if isinstance(original["list"][3], list): + original["list"][3][0] = 200 + if "inner_list" in original["list"][4]: + original["list"][4]["inner_list"][0] = 300 + + if "dict" in original and "inner_dict" in original["dict"]: + original["dict"]["inner_dict"]["key"] = "modified" + + if "torch_tensor" in original: + original["torch_tensor"][0] = 100 + + if "np_array" in original: + original["np_array"][0] = 200 + + # Returning both the original and the copy for inspection + return original, copy diff --git a/tests/test_youtube_nodes.py b/tests/test_youtube_nodes.py new file mode 100644 index 0000000..c2cf1cd --- /dev/null +++ b/tests/test_youtube_nodes.py @@ -0,0 +1,17 @@ +import os + +os.environ["UNIT_TEST"] = "True" +from custom_nodes.EternalKernelLiteGraphNodes.youtube_nodes import YoutubeUploadNode + + +def test_upload(): + node = YoutubeUploadNode() + output = node.upload_video( + file_path=os.path.join("tests", "example.mp4"), + category=22, + title="Test Video", + description="This is a test video for the Youtube Upload Node", + keywords="test, youtube, node", + privacy_status="unlisted" + ) + print (output) diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..81e5881 --- /dev/null +++ b/utils.py @@ -0,0 +1,25 @@ +import copy + +from collections import OrderedDict + + +def etk_deep_copy(x, verbose=False): + if type(x).__module__ == "numpy": + return x.copy() + elif type(x).__module__ == "torch": + return x.clone() + elif isinstance(x, list): + return [etk_deep_copy(item) for item in x] + elif isinstance(x, tuple): + return tuple(etk_deep_copy(item) for item in x) + elif isinstance(x, dict): + return {k: etk_deep_copy(v) for k, v in x.items()} + elif isinstance(x, set): + return set(etk_deep_copy(item) for item in x) + elif isinstance(x, OrderedDict): + return OrderedDict((k, etk_deep_copy(v)) for k, v in x.items()) + else: + if verbose: + print(f"copying {x}, of type {type(x)}") + + return copy.deepcopy(x) diff --git a/video.py b/video.py index 7a38975..46ae2e2 100644 --- a/video.py +++ b/video.py @@ -352,7 +352,6 @@ class ImageStackToVideoFile(metaclass=WidgetMetaclass): except OSError: self.CODE = "Unable to get source code" - self.FUNC = self.handler try: @@ -541,22 +540,21 @@ class TemporalOpticalFlowSmoothing(metaclass=WidgetMetaclass): # Create intermediary frames frames = [] for i in range(k): - factor = i / k - #factor = factor * 2 - 1 + factor = i / k + # factor = factor * 2 - 1 - #normalize flow + # normalize flow n_flow = flow.clone() print(n_flow.min(), n_flow.max(), n_flow.dtype) - #n_flow[:,0,:,:] = n_flow[:,0,:,:] / (n_flow[:,0,:,:].max() - n_flow[:,0,:,:].min()) - #n_flow[:,1,:,:] = n_flow[:,1,:,:] / (n_flow[:,1,:,:].max() - n_flow[:,1,:,:].min()) - #print(n_flow.min(), n_flow.max(), n_flow.dtype) - #n_flow[:, 0, :, :] = n_flow[:, 0, :, :] - n_flow[:, 0, :, :].min() - #n_flow[:, 1, :, :] = n_flow[:, 1, :, :] - n_flow[:, 1, :, :].min() - #print (n_flow.min(), n_flow.max(),n_flow.dtype) - n_flow= n_flow/n_flow.shape[2] + # n_flow[:,0,:,:] = n_flow[:,0,:,:] / (n_flow[:,0,:,:].max() - n_flow[:,0,:,:].min()) + # n_flow[:,1,:,:] = n_flow[:,1,:,:] / (n_flow[:,1,:,:].max() - n_flow[:,1,:,:].min()) + # print(n_flow.min(), n_flow.max(), n_flow.dtype) + # n_flow[:, 0, :, :] = n_flow[:, 0, :, :] - n_flow[:, 0, :, :].min() + # n_flow[:, 1, :, :] = n_flow[:, 1, :, :] - n_flow[:, 1, :, :].min() + # print (n_flow.min(), n_flow.max(),n_flow.dtype) + n_flow = n_flow / n_flow.shape[2] print(n_flow.min(), n_flow.max(), n_flow.dtype) - flow_interpolated = n_flow * factor grid = TemporalOpticalFlowSmoothing.make_grid(flow_interpolated) grid = grid.permute(0, 2, 3, 1) @@ -567,7 +565,7 @@ class TemporalOpticalFlowSmoothing(metaclass=WidgetMetaclass): return frames @staticmethod - def calculate_optical_flow(image_stack, device,k): + def calculate_optical_flow(image_stack, device, k): # from torch.multiprocessing import Pool, set_start_method # set_start_method('spawn') # set the start method to spawn @@ -594,8 +592,7 @@ class TemporalOpticalFlowSmoothing(metaclass=WidgetMetaclass): flows.append(keys[0].cpu()) for f in flow: flows.append(f[0].cpu()) - #flows.append(keys[1].cpu()) - + # flows.append(keys[1].cpu()) return flows @@ -608,7 +605,6 @@ class TemporalOpticalFlowSmoothing(metaclass=WidgetMetaclass): image_stack = image_stack.permute(0, 2, 1, 3) - # make sure image_stack is a torch tensor and has the right dimensions if not torch.is_tensor(image_stack): image_stack = torch.tensor(image_stack, dtype=torch.float32) @@ -620,17 +616,140 @@ class TemporalOpticalFlowSmoothing(metaclass=WidgetMetaclass): image_stack = image_stack.to(device) # Calculate optical flow - flows = self.calculate_optical_flow(image_stack, device,k) + flows = self.calculate_optical_flow(image_stack, device, k) stack = torch.stack(flows, dim=0) # Convert optical flows to images for visualization # flow_images = [flow_to_image(flow) for flow in flows] - #change back to (b or t,w,h,c) + # change back to (b or t,w,h,c) stack = stack.permute(0, 2, 1, 3) return (stack,) -if __name__ == "__main__": - import doctest +# TW: Inherit for code reuse, consistency +import torch +from torchvision.utils import draw_keypoints - doctest.testmod() + +class MovingCircle(metaclass=WidgetMetaclass): + # VideoWidget class definition + + def calculate_start(self, x1, y1, width, height): + # Convert to 0-1 range + start_x = x1 * width + start_y = y1 * height + return (start_x, start_y) + + def calculate_end(self, x2, y2, width, height): + # Convert to 0-1 range + end_x = x2 * width + end_y = y2 * height + return (end_x, end_y) + + def generate_frames(self, num_frames, width, height): + # Match tensor shape from example + frames = torch.zeros(num_frames, width, height, 3) + return frames + + def draw_circle_or_oval(self, kp, width, height, radius, aspect_ratio, color=(255, 255, 255)): + import torch.nn.functional as F + # Initialize black background + img = torch.zeros((3, height, width), dtype=torch.uint8) + + # Draw circle as keypoints + keypoints = torch.tensor([[kp[0], kp[1]]]).unsqueeze(0) + img = draw_keypoints(img, keypoints, colors=color, radius=radius) + + # convert to (w,h,c) + #img = img.permute(1, 2, 0).unsqueeze(0) + img=img.unsqueeze(0) + + # Calculate new size + if aspect_ratio < 1: + new_size = (height, round(width / aspect_ratio)) + else: + new_size = (round(height * aspect_ratio), width) + + # Stretch the image + img_stretched = F.interpolate(img, size=new_size, mode='bilinear', align_corners=False) + + # Shrink it back + img_shrunk = F.interpolate(img_stretched, size=(height, width), mode='bilinear', align_corners=False) + + # Convert back to float values + + img_shrunk = img_shrunk.float() / 255 + + img_shrunk = img_shrunk.permute(0, 2, 3, 1) + return img_shrunk + + def draw_circle(self, kp, width, height, radius, color=(255, 255, 255)): + # Initialize black background + img = torch.zeros((3, height, width), dtype=torch.uint8) + # Draw circle as keypoints + keypoints = torch.tensor([[kp[0], kp[1]]]).unsqueeze(0) + img = draw_keypoints(img, keypoints, colors=color, radius=radius) + # convert to (w,h,c) + img = img.permute(1, 2, 0).unsqueeze(0) + # convert back to float values + img = img.float() / 255 + + return img + + # Metadata for documentation + RETURN_TYPES = ("IMAGE",) + CATEGORY = "ETK/video" + FUNCTION = "handler" + + # Match parent class method + @classmethod + def INPUT_TYPES(cls): + # Validate types, ranges + return {"required": + { + "x1": ("FLOAT", {"min": 0, "max": 1, "step": 0.1}), + "y1": ("FLOAT", {"min": 0, "max": 1, "step": 0.1}), + "x2": ("FLOAT", {"min": 0, "max": 1, "step": 0.1}), + "y2": ("FLOAT", {"min": 0, "max": 1, "step": 0.1}), + "num_frames": ("INT", {"min": 1}), + "width": ("INT", {"min": 1}), + "height": ("INT", {"min": 1}), + "radius": ("INT", {"min": 1}), + "color": ("STRING", {"default": "255 255 255"}), + "aspect_ratio": ("FLOAT", {"min": 0.1, "max": 10}) + + }} + + def handler(self, x1, y1, x2, y2, num_frames, width, height, radius, color="255 255 255", aspect_ratio=1): + import torch + # Avoid errors on bad values + assert 0 <= x1 <= 1 + assert 0 <= y1 <= 1 + + # Split calculations + start = torch.tensor(self.calculate_start(x1, y1, width, height)) + + # Split calculations + end = torch.tensor(self.calculate_end(x2, y2, width, height)) + + # color is supposed to be a tuple + color = tuple(map(lambda x:int(x), color.split())) + # Initialize for compositing + + frames = [] + for i in range(num_frames): + # Leverage PyTorch + cur_pos = start + (end - start) * i / num_frames + circle = self.draw_circle_or_oval(cur_pos, width, height, radius, aspect_ratio, color) + # Use broadcasting to overlay + frames.append(circle) + + # cat + frames = torch.cat(frames, dim=0) + + return (frames,) + + if __name__ == "__main__": + import doctest + + doctest.testmod() diff --git a/youtube_nodes.py b/youtube_nodes.py new file mode 100644 index 0000000..e6f1d9a --- /dev/null +++ b/youtube_nodes.py @@ -0,0 +1,73 @@ +import subprocess +import os + +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} + + +def ETK_youtube_base(cls): + cls.CATEGORY = "ETK/Youtube" + NODE_CLASS_MAPPINGS[cls.__name__] = cls + return cls + + +@ETK_youtube_base +class YoutubeUploadNode: + VALID_PRIVACY_STATUSES = ["public", "private", "unlisted"] + VALID_CATEGORIES = [1, 2, 10, 15, 17, 19, 20, 22, 23, 24, 25, 26, 27, 28, 29] + CATEGORY_NAMES = ['Film & Animation', 'Autos & Vehicles', 'Music', 'Pets & Animals', 'Sports', 'Travel & Events', + 'Gaming', 'People & Blogs', 'Comedy', 'Entertainment', 'News & Politics', 'Howto & Style', + 'Education', 'Science & Technology', 'Nonprofits & Activism'] + + OUTPUT_NODE = True + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "file_path": ("STRING", {"description": "Path to the video file to upload"}), + "category": (cls.CATEGORY_NAMES,), + "title": ("STRING", {"default": "My video", "description": "Title of the video"}), + "description": ("STRING", {"default": "My video description", "multiline": True, + "description": "Description of the video"}), + "keywords": ("STRING", {"default": "tag1, tag2, tag3", "description": "Keywords for the video"}), + "privacy_status": (cls.VALID_PRIVACY_STATUSES,) + } + } + + RETURN_TYPES = ('STRING',) + FUNCTION = "upload_video" + CATEGORY = "ETK/Youtube" + + def upload_video(self, file_path, category, title, description, keywords, privacy_status): + import sys + my_python_exe = sys.executable + + category = self.CATEGORY_NAMES.index(category) + + args = ["--file", file_path, "--category", str(category), "--title", title, "--description", description, + "--keywords", keywords, "--privacyStatus", privacy_status] + + current_dir = os.getcwd() + os.chdir(os.path.dirname(os.path.realpath(__file__))) + + # result = subprocess.run([my_python_exe, "./modules/youtube.py"] + args, + # capture_output=True, + # text=True) + + # the above does not capture the output, so we use the below instead + result = subprocess.run([my_python_exe, "./modules/youtube.py"] + args, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True) + + os.chdir(current_dir) + error = result.stderr + if error: + raise Exception(error) + + return (result.stdout,) + + +if __name__ == "__main__": + pass