From 5781a0efe0f2fb27e5df50e851900fa4166368f2 Mon Sep 17 00:00:00 2001 From: Victor Naranjo Date: Thu, 14 Dec 2023 15:18:16 -0800 Subject: [PATCH] Initial Commit --- README.md | 34 +++++++++ __init__.py | 16 ++++ chat_gpt_integration.py | 161 ++++++++++++++++++++++++++++++++++++++++ roles.json | 14 ++++ 4 files changed, 225 insertions(+) create mode 100644 README.md create mode 100644 __init__.py create mode 100644 chat_gpt_integration.py create mode 100644 roles.json diff --git a/README.md b/README.md new file mode 100644 index 0000000..29efa46 --- /dev/null +++ b/README.md @@ -0,0 +1,34 @@ +# Credits to Omar92. This is based heavily off of his code +https://github.com/omar92/ComfyUI-QualityOfLifeSuit_Omar92 +I rewrote the core logic because this code was not compaitible with the new version of the OpenAi api + +# ComfyUI-ChatGPTIntegration +Single node to prompt ChatGPT and will return an input for your CLip TEXT Encode Prompt + +## ComfyUI +ComfyUI is an advanced node-based UI that utilizes Stable Diffusion, allowing you to create customized workflows such as image post-processing or conversions. + +## How to install +Download the zip file. +Extract to ..\ComfyUI\custom_nodes. +Restart ComfyUI if it was running (reloading the web is not enough). +You will find my nodes under the new group O/.... + +## How to update +- No Auto Update +The file looks like this : + +{ +"openAI_API_Key": "sk-#################################" +} + +## ChatGPTPrompt +This node harnesses the power of chatGPT, an advanced language model that can generate detailed image descriptions from a small input. +- you need to have OpenAI API key , which you can find at https://beta.openai.com/docs/developer-apis/overview +- Once you have your API key, add it to the `config.json` file +- I have made it a separate file, so that the API key doesn't get embedded in the generated images. + + +## Contact +### Discord: vienteck#6218 +### GitHub: vienteck (https://github.com/vienteck) \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..36aa1e3 --- /dev/null +++ b/__init__.py @@ -0,0 +1,16 @@ +import json +import os +from .chat_gpt_integration import NODE_CLASS_MAPPINGS +from .chat_gpt_integration import NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] + +#check if config file exists +if not os.path.isfile(os.path.join(os.path.dirname(os.path.realpath(__file__)),"config.json")): + #create config file + print(f'config file not found. Creating one for you') + config = { + "openAI_API_Key": "sk-#########################################" + } + with open(os.path.join(os.path.dirname(os.path.realpath(__file__)),"config.json"), "w") as f: + json.dump(config, f, indent=4) diff --git a/chat_gpt_integration.py b/chat_gpt_integration.py new file mode 100644 index 0000000..d69eb57 --- /dev/null +++ b/chat_gpt_integration.py @@ -0,0 +1,161 @@ +import importlib +import json +import os +import random +from openai import OpenAI + +openAI_models = None +NODE_FILE = os.path.abspath(__file__) + +CONFIG_FILE = f"{os.path.dirname(NODE_FILE)}\\config.json" +ROLES_FILE = f"{os.path.dirname(NODE_FILE)}\\roles.json" + +def get_roles(): + try: + # open config file + roles_list = [] + print(f"Roles File Location: {ROLES_FILE}") + + with open(ROLES_FILE, "r") as f: # Open the file and read the API key + roles_content = json.load(f) + for role in roles_content["creativeThinkers"]: + roles_list.append(role["name"]) + except: + print("Error: OpenAI API key file not found OpenAI features wont work for you") + return "" + return roles_list + +def get_api_key() -> str: + # Helper function to get the API key from the file + try: + # open config file + print(f"Config File Location: {CONFIG_FILE}") + with open(CONFIG_FILE, "r") as f: # Open the file and read the API key + config = json.load(f) + api_key = config["openAI_API_Key"] + except: + print("Error: OpenAI API key file not found OpenAI features wont work for you") + return "" + return api_key # Return the API key + +def get_openAI_models() -> list[str]: + print('getting ai models') + global openAI_models + if openAI_models != None: + return openAI_models + + install_openai() + import openai + + # Set the API key for the OpenAI module + openai.api_key = get_api_key() + + try: + models = openai.models.list() # Get the list of models + except: + print("Error: OpenAI API key is invalid OpenAI features wont work for you") + return [] + + openAI_models = [] # Create a list for the chat models + for model in models.data: # Loop through the models + openAI_models.append(model.id) # Add the model to the list + + return openAI_models # Return the list of chat models + +openAI_gpt_models = None + +def get_gpt_models() -> list[str]: + print('getting models') + global openAI_gpt_models + if openAI_gpt_models != None: + return openAI_gpt_models + models = get_openAI_models() + openAI_gpt_models = [] # Create a list for the chat models + for model in models: # Loop through the models + if "gpt" in model.lower(): + openAI_gpt_models.append(model) + return openAI_gpt_models # Return the list of chat models + +def install_openai(): + # Helper function to install the OpenAI module if not already installed + try: + importlib.import_module("openai") + except ImportError: + import pip + + pip.main(["install", "openai"]) + +def get_gpt_roles(): + roles = ["AI Assistant", "Alejandro Jodorowsky", "H.R. Giger"] + return roles + +def GetPrompt(role, prompt, input_model, max_words,custom_formatting) -> str: + client = OpenAI( + api_key=get_api_key(), + ) + print("submitting prompt to OpenAi") + sb = prompt.strip() + if len(custom_formatting) > 0 : + sb += f' in the format: {custom_formatting}' + sb += f' using at most {max_words} words' + + completion = client.chat.completions.create( + model=input_model, + messages=[ + { + "role": "system", + "content": f"You are {role}", + }, + { + "role": "user", + "content": f"{prompt}", + }, + ], + ) + print(f'recieved {len(completion.choices[0].message.content)} from OpenAi') + return completion.choices[0].message.content + +class ChatGptPrompt: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + # Multiline string input for the prompt + "prompt": ("STRING", {"multiline": True}), + "model": (get_gpt_models(), {"default": "gpt-3.5-turbo"}), + "role": (get_roles(), {"default": "AI Assistant"}), + "max_words": ("INT", {"default": 400,"min":1,"max":1000}) + }, + "optional": { + "seed": ("INT", {"default": random.randint(1,0xFFFFFFFFFFFFFFFF), "min": 0, "max": 0xFFFFFFFFFFFFFFFF}), + "custom_formatting": ("STRING", {"multiline": True}), + + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("Prompt",) + + FUNCTION = "process" + + OUTPUT_NODE = True + + CATEGORY = "OpenAI" # Define the category for the node + + def process(self, prompt,model,seed,role,max_words,custom_formatting): + GetPrompt(role,prompt,model,max_words,custom_formatting) + return "This is a submission test" + +# A dictionary that contains all nodes you want to export with their names +# NOTE: names should be globally unique +NODE_CLASS_MAPPINGS = { + "ChatGptPrompt": ChatGptPrompt +} + +# A dictionary that contains the friendly/humanly readable titles for the nodes +NODE_DISPLAY_NAME_MAPPINGS = { + "ChatGptPrompt": "ChatGPT Prompt Node" +} diff --git a/roles.json b/roles.json new file mode 100644 index 0000000..ce74afa --- /dev/null +++ b/roles.json @@ -0,0 +1,14 @@ +{ + "creativeThinkers": [ + {"name": "Salvador Dalí", "field": "Painter", "style": "Surrealist"}, + {"name": "Frida Kahlo", "field": "Painter", "style": "Symbolic"}, + {"name": "David Lynch", "field": "Filmmaker", "style": "Surreal, Unsettling"}, + {"name": "Hayao Miyazaki", "field": "Animator/Filmmaker", "style": "Enchanting, Animated"}, + {"name": "Leonora Carrington", "field": "Painter/Writer", "style": "Magical, Mythical"}, + {"name": "Haruki Murakami", "field": "Author", "style": "Surreal, Dreamlike"}, + {"name": "Octavia E. Butler", "field": "Author", "style": "Innovative, Thought-Provoking"}, + {"name": "Hieronymus Bosch", "field": "Painter", "style": "Fantastical, Symbolic"}, + {"name": "Philip K. Dick", "field": "Author", "style": "Exploratory, Sci-Fi"}, + {"name": "Terry Gilliam", "field": "Filmmaker", "style": "Visually Imaginative, Eccentric"} + ] +}