diff --git a/nodes/any.py b/nodes/any.py index 93e32ea..a18201c 100644 --- a/nodes/any.py +++ b/nodes/any.py @@ -1,30 +1,28 @@ """ -AnyNode v0.1 +@author: newsbubbles +@title: AnyNode v0.1 +@nickname: AnyNode +@description: This single node uses an LLM to generate a functionality based on your request. You can make the node do anything. author: newsbubbles -inspired by: rgthree AnySwitch -shouts and thanks: MachineLearners Discord """ +# inspired by: rgthree AnySwitch, shouts and thanks to MachineLearners Discord # TODO: finish inputs section which shows python types of inputs for the final prompt # - Store the json in some sort of hidden file which stores in the saved workflow # - How do I use just the "internal variables" since it seems scope on AnyNode class params are global (not good for storing script) -import os, json, random, string, sys, math, datetime, collections, itertools, functools, urllib, shutil, re, numpy, torch +import os, json, random, string, sys, math, datetime, collections, itertools, functools, urllib, shutil, re, torch +import numpy as np +import collections.abc import traceback import os import openai from openai import OpenAI -from .context_utils import is_context_empty, _create_context_data -from .constants import get_category, get_name -from .utils import any_type - -def is_none(value): - """Checks if a value is none. Pulled out in case we want to expand what 'None' means.""" - if value is not None: - if isinstance(value, dict) and 'model' in value and 'clip' in value: - return is_context_empty(value) - return value is None +# from .context_utils import is_context_empty, _create_context_data +# from .constants import get_category, get_name +from .utils import any_type, is_none, variable_info +# The template for the system message sent to ChatCompletions SYSTEM_TEMPLATE = """ # Coding a Python Function You are an expert python coder who specializes in writing custom nodes for ComfyUI. @@ -33,6 +31,10 @@ You are an expert python coder who specializes in writing custom nodes for Comfy It is not required to use any of these libraries, but if you do use any import in your code it must be on this list: [[IMPORTS]] +## Input Data +Here is some important information about the input data: +- input_data: [[INPUT]] + ## Coding Instructions - Your job is to code the user's requested node given the input and desired output type. - Code only the contents of the function itself. @@ -47,13 +49,14 @@ class AnyNode: NAME = "AnyNode" CATEGORY = "utils" - - script = None - last_prompt = None - imports = [] - + ALLOWED_IMPORTS = {"os", "re", "json", "random", "string", "sys", "math", "datetime", "collections", "itertools", "functools", "urllib", "shutil", "numpy", "openai", "traceback", "torch"} + def __init__(self): + self.script = None + self.last_prompt = None + self.imports = [] + @classmethod def INPUT_TYPES(cls): # pylint: disable = invalid-name, missing-function-docstring return { @@ -75,16 +78,17 @@ class AnyNode: RETURN_NAMES = ('any',) FUNCTION = "go" - def get_openai_response(self, prompt): + def get_openai_response(self, prompt:str, any=None) -> str: try: client = OpenAI( # This is the default and can be omitted api_key=os.environ.get("OPENAI_API_KEY"), ) + final_template = SYSTEM_TEMPLATE.replace('[[IMPORTS]]', ",".join(list(self.ALLOWED_IMPORTS))).replace('[[INPUT]]', variable_info(any)) response = client.chat.completions.create( model="gpt-4o", # Use the model of your choice, e.g., gpt-4 or gpt-3.5-turbo messages=[ - {"role": "system", "content": SYSTEM_TEMPLATE}, + {"role": "system", "content": final_template}, {"role": "user", "content": prompt} ] ) @@ -125,8 +129,7 @@ class AnyNode: if not is_none(any): if self.script is None or self.last_prompt != prompt: # Generate the function code using OpenAI - r = self.get_openai_response(prompt.replace('[[IMPORTS]]', ",".join(list(self.ALLOWED_IMPORTS)))) - print(f"Generated code:\n{r}") + r = self.get_openai_response(prompt, any=any) # Store the script for future use self.script = self.extract_imports(r) diff --git a/nodes/utils.py b/nodes/utils.py index 09c20e5..0f93b61 100644 --- a/nodes/utils.py +++ b/nodes/utils.py @@ -1,7 +1,9 @@ import json import os import re - +import numpy as np +import torch +from .context_utils import is_context_empty class AnyType(str): """A special class that is always equal in not equal comparisons. Credit to pythongosssss""" @@ -82,4 +84,61 @@ def path_exists(path): """Checks if a path exists, accepting None type.""" if path is not None: return os.path.exists(path) - return False \ No newline at end of file + return False + +def is_none(value): + """Checks if a value is none. Pulled out in case we want to expand what 'None' means.""" + if value is not None: + if isinstance(value, dict) and 'model' in value and 'clip' in value: + return is_context_empty(value) + return value is None + +def get_variable_info(variable): + info = {} + + # Get the exact type of the variable + info['type'] = type(variable).__name__ + + # Check for common types and add relevant information + if isinstance(variable, (np.ndarray, torch.Tensor)): + info['shape'] = tuple(variable.shape) + info['dtype'] = str(variable.dtype) + elif isinstance(variable, (list, tuple, set)): + info['length'] = len(variable) + if len(variable) > 128: + info['structure'] = f"{info['type']} of length {info['length']}" + else: + info['structure'] = [type(v).__name__ for v in variable] + elif isinstance(variable, dict): + info['length'] = len(variable) + if len(variable) > 128: + info['structure'] = f"{info['type']} of length {info['length']}" + else: + info['keys'] = list(variable.keys()) + info['values'] = [type(variable[k]).__name__ for k in variable] + elif isinstance(variable, str): + info['length'] = len(variable) + elif hasattr(variable, '__dict__'): + # For user-defined classes, show their attributes + info['attributes'] = vars(variable) + + return info + +def variable_info(variable): + info = get_variable_info(variable) + info_str = f"Type: {info['type']}" + + if 'shape' in info: + info_str += f", Shape: {info['shape']}, Dtype: {info['dtype']}" + if 'length' in info: + info_str += f", Length: {info['length']}" + if 'structure' in info: + info_str += f", Structure: {info['structure']}" + if 'keys' in info: + info_str += f", Keys: {info['keys']}" + if 'values' in info: + info_str += f", Values: {info['values']}" + if 'attributes' in info: + info_str += f", Attributes: {info['attributes']}" + + return info_str