Add several new nodes:

- DecodeTranscriptNode to parse transcript text
- YoutubeUploadNode to upload videos to YouTube
- Audio metadata and file list nodes
- Image stack blending and stripping alpha channel
- Sampler dataset creation node
- Utils for deep copying tensors and arrays

Also includes:

- Testing for new utils module
- Git nodes marked with base class
- Renamed nodes.py to avoid conflicts
- Import refactoring to handle naming

Specifics:

- descript.py: New transcript parsing node
- youtube_nodes.py: New YouTube upload node
- audio.py: New audio metadata and file list nodes
- utils.py: Deep copy utility functions
- image.py: New image stack blending and alpha stripping
- tests/test_utils.py: Tests for new utils
- git.py: Marked git nodes with base class
- __init__.py: Refactored imports and mappings
- requirements.txt: Added YouTube dependencies
- functional.py: Minor tweaks
- modules/youtube.py: Helper for upload node
- video.py: Minor cleanup
- tests/test_youtube_nodes.py: Test for YouTube upload

__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

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,)

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,)

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

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):

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()

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))

requirements.txt

@@ -0,0 +1,5 @@
+google-api-python-client
+google-auth-oauthlib
+google-auth-httplib2
+oauth2client
+httplib2

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

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)

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)

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()

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