ui_sheel_privet

README.md

@@ -17,6 +17,7 @@ To install, either:
 - Input Image
 - Input Float
 - Input Integer
+- Input Video
 
 Each input node supports setting a default value and additional configuration options.
 
@@ -25,6 +26,9 @@ Each input node supports setting a default value and additional configuration op
 - Save Image
 - Save Images
 - Save Video - VHS
+- Output Text
+- Output Float
+- Output Integer
 
 ### Convert Widgets to ShellAgent Inputs
 

comfy-nodes/input_audio.py

@@ -0,0 +1,200 @@
+import folder_paths
+import node_helpers
+
+from PIL import Image, ImageOps, ImageSequence, ImageFile
+import numpy as np
+import torch
+import os
+import uuid
+import tqdm
+import torchaudio
+import hashlib
+from comfy_extras.nodes_audio import SaveAudio
+
+
+class LoadAudio:
+    @classmethod
+    def INPUT_TYPES(s):
+        input_dir = folder_paths.get_input_directory()
+        files = folder_paths.filter_files_content_types(
+            os.listdir(input_dir), ["audio", "video"])
+        return {"required": {"audio": (sorted(files), {"audio_upload": True})}}
+
+    CATEGORY = "audio"
+
+    RETURN_TYPES = ("AUDIO", )
+    FUNCTION = "load"
+
+    def load(self, audio):
+        audio_path = folder_paths.get_annotated_filepath(audio)
+        waveform, sample_rate = torchaudio.load(audio_path)
+        audio = {"waveform": waveform.unsqueeze(0), "sample_rate": sample_rate}
+        return (audio, )
+
+    @classmethod
+    def IS_CHANGED(s, audio):
+        image_path = folder_paths.get_annotated_filepath(audio)
+        m = hashlib.sha256()
+        with open(image_path, 'rb') as f:
+            m.update(f.read())
+        return m.digest().hex()
+
+    @classmethod
+    def VALIDATE_INPUTS(s, audio):
+        if not folder_paths.exists_annotated_filepath(audio):
+            return "Invalid audio file: {}".format(audio)
+        return True
+
+
+class ShellAgentPluginInputAudio:
+    @classmethod
+    def INPUT_TYPES(s):
+        input_dir = folder_paths.get_input_directory()
+        files = folder_paths.filter_files_content_types(
+            os.listdir(input_dir), ["audio", "video"])
+        return {
+            "required": {
+                "input_name": (
+                    "STRING",
+                    {"multiline": False, "default": "input_audio", "forceInput": False},
+                ),
+                "default_value": (
+                    sorted(files), {"audio_upload": True, "forceInput": False}
+                ),
+            },
+            "optional": {
+                "description": (
+                    "STRING",
+                    {"multiline": True, "default": "", "forceInput": False},
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("AUDIO", )
+    FUNCTION = "load"
+
+    CATEGORY = "shellagent"
+
+    @classmethod
+    def validate(cls, **kwargs):
+        schema = {
+            "title": kwargs["input_name"],
+            "type": "string",
+            "default": kwargs["default_value"],
+            "description": kwargs.get("description", ""),
+            "url_type": "audio"
+        }
+        return schema
+
+    @classmethod
+    def VALIDATE_INPUTS(s, audio):
+        if not folder_paths.exists_annotated_filepath(audio):
+            return "Invalid audio file: {}".format(audio)
+        return True
+
+    @classmethod
+    def VALIDATE_INPUTS(s, input_name, default_value, description=""):
+        audio = default_value
+        if audio.startswith("http"):
+            return True
+
+        if not folder_paths.exists_annotated_filepath(audio):
+            return "Invalid audio file: {}".format(audio)
+        return True
+
+    def load(self, input_name, default_value=None, display_name=None, description=None):
+        input_dir = folder_paths.get_input_directory()
+        audio_path = default_value
+        try:
+            if audio_path.startswith('http'):
+                import requests
+                from io import BytesIO
+                print("Fetching audio from url: ", audio_path)
+                response = requests.get(audio_path)
+                response.raise_for_status()
+                audio_file = BytesIO(response.content)
+                waveform, sample_rate = torchaudio.load(audio_file)
+            else:
+                if not os.path.isfile(audio_path):  # abs path
+                    # local path
+                    audio_path = os.path.join(input_dir, audio_path)
+                waveform, sample_rate = torchaudio.load(audio_path)
+
+            audio = {"waveform": waveform.unsqueeze(
+                0), "sample_rate": sample_rate}
+            return (audio, )
+            # image = ImageOps.exif_transpose(image)
+            # image = image.convert("RGB")
+            # image = np.array(image).astype(np.float32) / 255.0
+            # image = torch.from_numpy(image)[None,]
+            # return [image]
+        except Exception as e:
+            raise e
+
+
+class ShellAgentSaveAudios(SaveAudio):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"audio": ("AUDIO", ),
+                "output_name": ("STRING", {"multiline": False, "default": "output_audio"},),
+                             "filename_prefix": ("STRING", {"default": "audio/ComfyUI"})},
+                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
+                }
+    # {
+    #         "required": {
+    #             "images": ("IMAGE", {"tooltip": "The audio to save."}),
+    #             "output_name": ("STRING", {"multiline": False, "default": "output_image"},),
+    #             "filename_prefix": ("STRING", {"default": "ComfyUI", "tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."})
+    #         },
+    #         "hidden": {
+    #             "prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"
+    #         },
+    #     }
+
+    CATEGORY = "shellagent"
+
+    @classmethod
+    def validate(cls, **kwargs):
+        schema = {
+            "title": kwargs["output_name"],
+            "type": "array",
+            "items": {
+                "type": "string",
+                "url_type": "audio",
+            }
+        }
+        return schema
+
+    def save_audio(self, audio, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None, **extra_kwargs):
+        results = super().save_audio(audio, filename_prefix, prompt, extra_pnginfo)
+        results["shellagent_kwargs"] = extra_kwargs
+        return results
+
+    def save_flac(self, audio, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None, **extra_kwargs):
+        results = super().save_flac(audio, filename_prefix, "flac", prompt, extra_pnginfo)
+        results["shellagent_kwargs"] = extra_kwargs
+        return results
+
+
+class ShellAgentSaveAudio(ShellAgentSaveAudios):
+    @classmethod
+    def validate(cls, **kwargs):
+        schema = {
+            "title": kwargs["output_name"],
+            "type": "string",
+            "url_type": "audio",
+        }
+        return schema
+
+
+NODE_CLASS_MAPPINGS = {
+    "ShellAgentPluginInputAudio": ShellAgentPluginInputAudio,
+    "ShellAgentPluginSaveAudios": ShellAgentSaveAudios,
+    "ShellAgentPluginSaveAudio": ShellAgentSaveAudio,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "ShellAgentPluginInputAudio": "Input Audio (ShellAgent Plugin)",
+    "ShellAgentPluginSaveAudios": "Save Audios (ShellAgent Plugin)",
+    "ShellAgentPluginSaveAudio": "Save Audio (ShellAgent Plugin)",
+}

comfy-nodes/input_image.py

@@ -1,11 +1,39 @@
 import folder_paths
-from PIL import Image, ImageOps
+import node_helpers
+
+from PIL import Image, ImageOps, ImageSequence, ImageFile
 import numpy as np
 import torch
 import os
 import uuid
 import tqdm
-
+from io import BytesIO
+import PIL
+import cv2
+from pillow_heif import register_heif_opener
+
+register_heif_opener()
+
+def safe_open_image(image_bytes):
+    try:
+        image_pil = Image.open(BytesIO(image_bytes))
+    except PIL.UnidentifiedImageError as e:
+        print(e)
+        # Convert response content (bytes) to a NumPy array
+        image_array = np.frombuffer(image_bytes, np.uint8)
+
+        # Decode the image from the NumPy array (OpenCV format: BGR)
+        image_cv = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
+
+        if image_cv is not None:
+            # Convert the BGR image to RGB
+            image_rgb = cv2.cvtColor(image_cv, cv2.COLOR_BGR2RGB)
+
+            # Convert the RGB NumPy array to a PIL Image
+            image_pil = Image.fromarray(image_rgb)
+        else:
+            raise ValueError("The image cannot be identified by neither PIL nor OpenCV")
+    return image_pil
 
 class ShellAgentPluginInputImage:
     @classmethod
@@ -20,20 +48,19 @@ def INPUT_TYPES(s):
                     {"multiline": False, "default": "input_image", "forceInput": False},
                 ),
                 "default_value": (
-                    # "STRING", {"image_upload": True, "default": files[0] if len(files) else ""},
                     sorted(files), {"image_upload": True, "forceInput": False}
                 ),
             },
             "optional": {
                 "description": (
                     "STRING",
-                    {"multiline": True, "default": "", "forceInput": False},
+                    {"multiline": False, "default": "", "forceInput": False},
                 ),
             }
         }
 
-    RETURN_TYPES = ("IMAGE",)
-    RETURN_NAMES = ("image",)
+    RETURN_TYPES = ("IMAGE", "MASK")
+    # RETURN_NAMES = ("image",)
 
     FUNCTION = "run"
 
@@ -45,32 +72,80 @@ def validate(cls, **kwargs):
             "title": kwargs["input_name"],
             "type": "string",
             "default": kwargs["default_value"],
-            "description": kwargs["description"],
+            "description": kwargs.get("description", ""),
             "url_type": "image"
         }
         return schema
 
     @classmethod
     def VALIDATE_INPUTS(s, input_name, default_value, description=""):
         image = default_value
+
         if image.startswith("http"):
             return True
 
+        if image == "":
+            return "Invalid image file: please check if the image is empty or invalid"
+
+        if os.path.isfile(image):
+            return True
+
         if not folder_paths.exists_annotated_filepath(image):
             return "Invalid image file: {}".format(image)
 
         return True
+
+    def convert_image_mask(self, img):
+        output_images = []
+        output_masks = []
+        w, h = None, None
+
+        excluded_formats = ['MPO']
+
+        for i in ImageSequence.Iterator(img):
+            i = node_helpers.pillow(ImageOps.exif_transpose, i)
+
+            if i.mode == 'I':
+                i = i.point(lambda i: i * (1 / 255))
+            image = i.convert("RGB")
+
+            if len(output_images) == 0:
+                w = image.size[0]
+                h = image.size[1]
+
+            if image.size[0] != w or image.size[1] != h:
+                continue
+
+            image = np.array(image).astype(np.float32) / 255.0
+            image = torch.from_numpy(image)[None,]
+            if 'A' in i.getbands():
+                mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
+                mask = 1. - torch.from_numpy(mask)
+            else:
+                mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
+            output_images.append(image)
+            output_masks.append(mask.unsqueeze(0))
+
+        if len(output_images) > 1 and img.format not in excluded_formats:
+            output_image = torch.cat(output_images, dim=0)
+            output_mask = torch.cat(output_masks, dim=0)
+        else:
+            output_image = output_images[0]
+            output_mask = output_masks[0]
+
+        return (output_image, output_mask)
+
 
     def run(self, input_name, default_value=None, display_name=None, description=None):
-        input_dir = folder_paths.get_input_directory()
         image_path = default_value
+        input_dir = folder_paths.get_input_directory()
         try:
             if image_path.startswith('http'):
                 import requests
                 from io import BytesIO
                 print("Fetching image from url: ", image_path)
                 response = requests.get(image_path)
-                image = Image.open(BytesIO(response.content))
+                image = safe_open_image(response.content)
             elif image_path.startswith('data:image/png;base64,') or image_path.startswith('data:image/jpeg;base64,') or image_path.startswith('data:image/jpg;base64,'):
                 import base64
                 from io import BytesIO
@@ -82,13 +157,14 @@ def run(self, input_name, default_value=None, display_name=None, description=Non
                 if not os.path.isfile(image_path): # abs path
                     # local path
                     image_path = os.path.join(input_dir, image_path)
-                image = Image.open(image_path).convert("RGB")
-
-            image = ImageOps.exif_transpose(image)
-            image = image.convert("RGB")
-            image = np.array(image).astype(np.float32) / 255.0
-            image = torch.from_numpy(image)[None,]
-            return [image]
+                image = node_helpers.pillow(Image.open, image_path)
+
+            return self.convert_image_mask(image)
+            # image = ImageOps.exif_transpose(image)
+            # image = image.convert("RGB")
+            # image = np.array(image).astype(np.float32) / 255.0
+            # image = torch.from_numpy(image)[None,]
+            # return [image]
         except Exception as e:
             raise e