input output nodes v1

Author: wl-zhao

.gitignore

@@ -0,0 +1 @@
+__pycache__

__init__.py

@@ -0,0 +1,54 @@
+"""
+@author: MyShell
+@title: comfyui-shellagent-plugin
+@description: 
+"""
+import os
+import sys
+
+sys.path.append(os.path.join(os.path.dirname(__file__)))
+
+import inspect
+import sys
+import importlib
+import subprocess
+import requests
+import folder_paths
+from folder_paths import add_model_folder_path, get_filename_list, get_folder_paths
+from tqdm import tqdm
+
+# from . import custom_routes
+# import routes
+
+ag_path = os.path.join(os.path.dirname(__file__))
+
+def get_python_files(path):
+    return [f[:-3] for f in os.listdir(path) if f.endswith(".py")]
+
+def append_to_sys_path(path):
+    if path not in sys.path:
+        sys.path.append(path)
+
+paths = ["comfy-nodes"]
+files = []
+
+for path in paths:
+    full_path = os.path.join(ag_path, path)
+    append_to_sys_path(full_path)
+    files.extend(get_python_files(full_path))
+
+NODE_CLASS_MAPPINGS = {}
+NODE_DISPLAY_NAME_MAPPINGS = {}
+
+# Import all the modules and append their mappings
+for file in files:
+    module = importlib.import_module(file)
+
+    if hasattr(module, "NODE_CLASS_MAPPINGS"):
+        NODE_CLASS_MAPPINGS.update(module.NODE_CLASS_MAPPINGS)
+    if hasattr(module, "NODE_DISPLAY_NAME_MAPPINGS"):
+        NODE_DISPLAY_NAME_MAPPINGS.update(module.NODE_DISPLAY_NAME_MAPPINGS)
+
+WEB_DIRECTORY = "web-plugin"
+print(NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS)
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

comfy-nodes/input_image.py

@@ -0,0 +1,151 @@
+import folder_paths
+from PIL import Image, ImageOps
+import numpy as np
+import torch
+import os
+import uuid
+import tqdm
+
+
+class ShellAgentPluginInputImage:
+    @classmethod
+    def INPUT_TYPES(s):
+        input_dir = folder_paths.get_input_directory()
+        files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
+        files = sorted(files)
+        return {
+            "required": {
+                "input_name": (
+                    "STRING",
+                    {"multiline": False, "default": "input_image"},
+                ),
+                "default_value": (
+                    "STRING", {"image_upload": True, "default": files[0] if len(files) else ""},
+                ),
+            },
+            "optional": {
+                "description": (
+                    "STRING",
+                    {"multiline": True, "default": ""},
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("image",)
+
+    FUNCTION = "run"
+
+    CATEGORY = "shellagent"
+
+    def run(self, input_name, default_value=None, display_name=None, description=None):
+        input_dir = folder_paths.get_input_directory()
+        image_path = default_value
+        try:
+            if image_path.startswith('http'):
+                import requests
+                from io import BytesIO
+                print("Fetching image from url: ", image)
+                response = requests.get(image)
+                image = Image.open(BytesIO(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
+                print("Decoding base64 image")
+                base64_image = image_path[image_path.find(",")+1:]
+                decoded_image = base64.b64decode(base64_image)
+                image = Image.open(BytesIO(decoded_image))
+            else:
+                # 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]
+        except Exception as e:
+            raise e
+
+video_extensions = ["webm", "mp4", "mkv", "gif"]
+
+class ShellAgentPluginInputVideo:
+    @classmethod
+    def INPUT_TYPES(s):
+        input_dir = folder_paths.get_input_directory()
+        files = []
+        for f in os.listdir(input_dir):
+            if os.path.isfile(os.path.join(input_dir, f)):
+                file_parts = f.split(".")
+                if len(file_parts) > 1 and (file_parts[-1] in video_extensions):
+                    files.append(f)
+                
+        return {
+            "required": {
+                "input_name": (
+                    "STRING",
+                    {"multiline": False, "default": "input_video"},
+                ),
+                "default_value": (
+                    "STRING", {"video_upload": True, "default": files[0] if len(files) else ""},
+                ),
+            },
+            "optional": {
+                "description": (
+                    "STRING",
+                    {"multiline": True, "default": ""},
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("video",)
+
+    FUNCTION = "run"
+
+    CATEGORY = "shellagent"
+
+    def run(self, input_name, default_value=None, description=None):
+        input_dir = folder_paths.get_input_directory()
+        if default_value.startswith("http"):
+            import requests
+
+            print("Fetching video from URL: ", default_value)
+            response = requests.get(default_value, stream=True)
+            file_size = int(response.headers.get("Content-Length", 0))
+            file_extension = default_value.split(".")[-1].split("?")[
+                0
+            ]  # Extract extension and handle URLs with parameters
+            if file_extension not in video_extensions:
+                file_extension = ".mp4"
+
+            unique_filename = str(uuid.uuid4()) + "." + file_extension
+            video_path = os.path.join(input_dir, unique_filename)
+            chunk_size = 1024  # 1 Kibibyte
+
+            num_bars = int(file_size / chunk_size)
+
+            with open(video_path, "wb") as out_file:
+                for chunk in tqdm(
+                    response.iter_content(chunk_size=chunk_size),
+                    total=num_bars,
+                    unit="KB",
+                    desc="Downloading",
+                    leave=True,
+                ):
+                    out_file.write(chunk)
+        else:
+            video_path = os.path.abspath(os.path.join(input_dir, default_value))
+
+        return (video_path,)
+
+
+NODE_CLASS_MAPPINGS = {
+    "ShellAgentPluginInputImage": ShellAgentPluginInputImage,
+    # "ShellAgentPluginInputVideo": ShellAgentPluginInputVideo,
+}
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "ShellAgentPluginInputImage": "Input Image (ShellAgent Plugin)",
+    # "ShellAgentPluginInputVideo": "Input Video (ShellAgent Plugin)"
+}

comfy-nodes/input_text.py

@@ -0,0 +1,44 @@
+import folder_paths
+from PIL import Image, ImageOps
+import numpy as np
+import torch
+
+class ShellAgentPluginInputText:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "input_name": (
+                    "STRING",
+                    {"multiline": False, "default": "input_text"},
+                ),
+            },
+            "optional": {
+                "default_value": (
+                    "STRING",
+                    {"multiline": True, "default": ""},
+                ),
+                "description": (
+                    "STRING",
+                    {"multiline": True, "default": ""},
+                ),
+                "choices": (
+                    "STRING",
+                    {"multiline": False, "default": ""},
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("text",)
+
+    FUNCTION = "run"
+
+    CATEGORY = "shellagent"
+
+    def run(self, input_name, default_value=None, display_name=None, description=None, choices=None):
+        return [default_value]
+
+
+NODE_CLASS_MAPPINGS = {"ShellAgentPluginInputText": ShellAgentPluginInputText}
+NODE_DISPLAY_NAME_MAPPINGS = {"ShellAgentPluginInputText": "Input Text (ShellAgent Plugin)"}