106 lines
2.4 KiB
Python
106 lines
2.4 KiB
Python
|
|
#import custom_nodes.Derfuu_ComfyUI_ModdedNodes.components.fields as field
|
|
import PIL.Image
|
|
import torch
|
|
import torchvision
|
|
import PIL.Image as Image
|
|
|
|
import doctest
|
|
# make sure to import ABC
|
|
from abc import ABC, abstractmethod
|
|
|
|
import os
|
|
import openai
|
|
|
|
|
|
class Config:
|
|
api_key = os.environ.get('OPENAI_API_KEY')
|
|
organization = os.environ.get('OPENAI_ORG')
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {}
|
|
|
|
openai.organization = Config.organization
|
|
openai.api_key = Config.api_key
|
|
|
|
|
|
class SimpleTextWidget:
|
|
"""base class for simple construction"""
|
|
|
|
def __init__(self, func):
|
|
self.func = func
|
|
|
|
|
|
@classmethod
|
|
def register(self, name):
|
|
global NODE_CLASS_MAPPINGS
|
|
|
|
# register the widget
|
|
NODE_CLASS_MAPPINGS[name] = self
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required":
|
|
{
|
|
"text": ("STRING", {"multiline": True}),
|
|
"Display": ("SERVER_STRING", {"multiline": True}),
|
|
},
|
|
|
|
}
|
|
|
|
CATEGORY = "text"
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "handler"
|
|
|
|
def handler(self, text,Display):
|
|
ret = self.func(text)
|
|
return (ret,)
|
|
|
|
|
|
try:
|
|
import folder_paths
|
|
except:
|
|
pass
|
|
|
|
import os
|
|
import numpy as np
|
|
import hashlib
|
|
|
|
### info for code completion AI ###
|
|
"""
|
|
all of these classes are plugins for comfyui and follow the same pattern
|
|
all of the images are torch tensors and it is unknown and unimportant if they are on the cpu or gpu
|
|
all image inputs are (B,W,H,C)
|
|
|
|
avoid numpy and PIL as much as possible
|
|
"""
|
|
|
|
|
|
class LLMCompletion(SimpleTextWidget):
|
|
"""uses the input text to call the specified LLM model and returns the output string"""
|
|
|
|
def OAI_completion(self, text):
|
|
out_message = [{"role": "user", "content": text}]
|
|
|
|
response = openai.ChatCompletion.create(
|
|
model="gpt-3.5-turbo",
|
|
messages=out_message
|
|
)
|
|
gpt_message = response["choices"][0]["message"]["content"].strip()
|
|
self.server_strings[self.SSID] = gpt_message
|
|
return gpt_message
|
|
|
|
def __init__(self):
|
|
self.func = self.OAI_completion
|
|
super().__init__(self.func)
|
|
import uuid
|
|
self.SSID = str(uuid.uuid4())
|
|
from main import server_obj_holder
|
|
self.server_strings = server_obj_holder[0].server_strings
|
|
self.server_strings["self.SSID"] = "empty"
|
|
|
|
|
|
|
|
|
|
LLMCompletion.register("LLM_Completion")
|