Add OpenAiGpt

__init__.py

@@ -11,6 +11,8 @@ from .negi.depth_estimation_by_marigold import DepthEstimationByMarigold
 from .negi.stereo_image_generator import StereoImageGenerator
 from .negi.image_reader_writer import RandomImageLoader, SaveImageToDirectory
 from .negi.detect_face_rotation_for_inpainting import DetectFaceRotationForInpainting
+from .negi.openai_gpt4v import OpenAiGpt4v
+from .negi.openai_gpt import OpenAiGpt
 
 NODE_CLASS_MAPPINGS = {
     "NegiTools_OpenAiDalle3": OpenAiDalle3,
@@ -28,6 +30,8 @@ NODE_CLASS_MAPPINGS = {
     "NegiTools_RandomImageLoader": RandomImageLoader,
     "NegiTools_SaveImageToDirectory": SaveImageToDirectory,
     "NegiTools_DetectFaceRotationForInpainting": DetectFaceRotationForInpainting,
+    "NegiTools_OpenAiGpt4v": OpenAiGpt4v,
+    "NegiTools_OpenAiGpt": OpenAiGpt,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -46,4 +50,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "NegiTools_RandomImageLoader": "Random Image Loader 🧅",
     "NegiTools_SaveImageToDirectory": "Save Image to Directory 🧅",
     "NegiTools_DetectFaceRotationForInpainting": "Detect Face Rotation for Inpainting 🧅",
+    "NegiTools_OpenAiGpt4v": "OpenAI GPT4V 🧅",
+    "NegiTools_OpenAiGpt": "OpenAI GPT 🧅",
 }

negi/image_reader_writer.py

@@ -1,18 +1,12 @@
 import glob
 import os
-import re
 
 import torch
 from PIL import Image
 import torchvision
 from torchvision.transforms import functional as TF
 
-
-def _get_directory(directory):
-    base_path = os.path.abspath(__file__)
-    for _ in range(4):
-        base_path = os.path.dirname(base_path)
-    return os.path.abspath(os.path.join(base_path, directory))
+from . import utils
 
 
 class RandomImageLoader:
@@ -34,7 +28,7 @@ class RandomImageLoader:
     CATEGORY = "utils"
 
     def doit(self, directory, seed):
-        directory = _get_directory(directory)
+        directory = utils.get_directory(directory)
         print("RandomImageLoader: directory = %s" % directory)
 
         files = (glob.glob(os.path.join(directory, "*.png")) +
@@ -76,18 +70,10 @@ class SaveImageToDirectory:
     CATEGORY = "utils"
 
     def doit(self, directory, image):
-        directory = _get_directory(directory)
-        os.makedirs(directory, exist_ok=True)
+        directory = utils.get_directory(directory)
         print("SaveImageToDirectory: directory = %s" % directory)
 
-        next_index = 0
-        files = glob.glob(os.path.join(directory, "out.??????.png"))
-        for file in files:
-            r = re.match(r"out\.(\d{6})\.png", os.path.basename(file))
-            if r is None:
-                continue
-            next_index = max(next_index, int(r.group(1)) + 1)
-
+        next_index = utils.find_next_index(directory)
         file_name = os.path.join(directory, "out.%06d.png" % next_index)
         print("SaveImageToDirectory: save to %s" % file_name)
 

negi/openai_dalle3.py

@@ -1,11 +1,15 @@
 import openai
 import base64
 import io
+import os
+import json
 
 import torch
 from PIL import Image
 from torchvision.transforms import functional as TF
 
+from . import utils
+
 
 class OpenAiDalle3:
     def __init__(self):
@@ -25,7 +29,12 @@ class OpenAiDalle3:
                 "prompt": ("STRING", {
                     "multiline": True,
                     "default": "great picture"
-                })
+                }),
+                "quality": (["HD", "Standard"],),
+                "style": (["vivid", "natural"],),
+                "retry": ("INT", {"default": 0, "min": 0, "max": 5}),
+                "auto_save": ("BOOLEAN", {"default": False}),
+                "auto_save_dir": ("STRING", {"multiline": False, "default": "./output_dalle3"}),
             }
         }
 
@@ -35,19 +44,45 @@ class OpenAiDalle3:
     OUTPUT_NODE = True
     CATEGORY = "Generator"
 
-    def doit(self, resolution, dummy_seed, prompt):
+    def doit(self, resolution, dummy_seed, prompt, quality, style, retry, auto_save, auto_save_dir):
         if (self.__cache_image is None or
                 self.__previous_resolution != resolution or self.__previous_seed != dummy_seed or
                 self.__previous_prompt != prompt):
-            r0 = self.__client.images.generate(
-                model="dall-e-3",
-                prompt=prompt,
-                size=resolution,
-                quality="hd",  # "standard"
-                n=1,
-                response_format="b64_json"
-            )
+            r0 = None
+            for retry_count in range(retry + 1):
+                try:
+                    r0 = self.__client.images.generate(
+                        model="dall-e-3",
+                        prompt=prompt,
+                        size=resolution,
+                        quality="hd" if quality == "HD" else "standard",
+                        style="vivid" if style == "vivid" else "natural",
+                        n=1,
+                        response_format="b64_json"
+                    )
+                    break
+                except openai.BadRequestError as ex:
+                    if retry_count >= retry:
+                        raise ex
+                    print("OpenAiDalle3: received BadRequestError, retrying... #%d : %s" % (
+                        retry_count + 1, json.dumps(ex.response.json())))
+
             im0 = Image.open(io.BytesIO(base64.b64decode(r0.data[0].b64_json)))
+
+            if auto_save:
+                directory = utils.get_directory(auto_save_dir)
+                next_index = utils.find_next_index(directory)
+                image_file_name = os.path.join(directory, "out.%06d.png" % next_index)
+                state_file_name = os.path.join(directory, "out.%06d.json" % next_index)
+                im0.save(image_file_name)
+                with open(state_file_name, "wt") as f:
+                    f.write(json.dumps({
+                        "resolution": resolution,
+                        "prompt": prompt,
+                        "quality": quality,
+                        "style": style
+                    }, indent=2, ensure_ascii=False))
+
             im1 = TF.to_tensor(im0.convert("RGBA"))
             im1[:3, im1[3, :, :] == 0] = 0
             revised_prompt = r0.data[0].revised_prompt

negi/openai_gpt.py

@@ -0,0 +1,83 @@
+import openai
+import time
+import urllib.error
+
+
+class OpenAiGpt:
+    def __init__(self):
+        self.__client = openai.OpenAI()
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ([
+                    "gpt-4-1106-preview",
+                    "gpt-4-vision-preview",
+                    "gpt-4",
+                    "gpt-4-0314",
+                    "gpt-4-0613",
+                    "gpt-4-32k",
+                    "gpt-4-32k-0314",
+                    "gpt-4-32k-0613",
+                    "gpt-3.5-turbo",
+                    "gpt-3.5-turbo-16k",
+                    "gpt-3.5-turbo-0301",
+                    "gpt-3.5-turbo-0613",
+                    "gpt-3.5-turbo-1106",
+                    "gpt-3.5-turbo-16k-0613",
+                ], {"default": "gpt-4-0613"}),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+                "max_tokens": ("INT", {"default": 512, "min": 16, "max": 4096}),
+                "system_prompt": ("STRING", {
+                    "multiline": True,
+                    "default": "You are a helpful assistant."
+                })
+            },
+            "optional": {
+                "a_role": (["user", "assistant", "system"], {"default": "user"}),
+                "a": ("STRING", {"multiline": False, "default": ""}),
+                "b_role": (["user", "assistant", "system"], {"default": "assistant"}),
+                "b": ("STRING", {"multiline": False, "default": ""}),
+                "c_role": (["user", "assistant", "system"], {"default": "user"}),
+                "c": ("STRING", {"multiline": False, "default": ""}),
+            }
+        }
+
+    RETURN_TYPES = ("STRING",)
+    FUNCTION = "doit"
+    OUTPUT_NODE = False
+    CATEGORY = "Generator"
+
+    def doit(self, model, seed, max_tokens, system_prompt, a_role, a, b_role, b, c_role, c):
+        messages = [{"role": "system", "content": system_prompt}]
+        input_role = [a_role, b_role, c_role]
+        input_text = [a, b, c]
+        for i in range(3):
+            if input_text[i] is not None and len(input_text[i]) > 0:
+                messages.append({
+                    "role": input_role[i] if input_role[i] is not None else "user",
+                    "content": input_text[i]
+                })
+
+        try_count = 0
+        r0 = None
+        while True:
+            try_count += 1
+            try:
+                r0 = self.__client.chat.completions.create(
+                    model=model,
+                    max_tokens=max_tokens,
+                    seed=seed,
+                    messages=messages
+                )
+                break
+            except openai.AuthenticationError as ex:
+                raise ex
+            except (urllib.error.HTTPError, openai.OpenAIError) as ex:
+                if try_count >= 3:
+                    raise ex
+                time.sleep(5)
+                continue
+
+        return (r0.choices[0].message.content,)

negi/openai_gpt4v.py

@@ -0,0 +1,85 @@
+import os
+import base64
+
+import openai
+import requests
+
+import torch
+import torchvision
+
+
+_api_key = os.environ.get("OPENAI_API_KEY")
+_tmp_file = "gpt4v_tmp.jpg"
+
+
+class OpenAiGpt4v:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+                "detail": (["auto", "low", "high"],),
+                "max_tokens": ("INT", {"default": 512, "min": 16, "max": 4096}),
+                "prompt": ("STRING", {
+                    "multiline": True,
+                    "default": "What’s in this image?"
+                }),
+            }
+        }
+
+    RETURN_TYPES = ("STRING",)
+    FUNCTION = "doit"
+    OUTPUT_NODE = False
+    CATEGORY = "Generator"
+
+    def doit(self, image, seed, detail, max_tokens, prompt):
+        _ = seed
+
+        im0 = torchvision.transforms.functional.to_pil_image(torch.permute(image[0], (2, 0, 1)))
+        im0.save(_tmp_file)
+        with open(_tmp_file, "rb") as f:
+            encoded_image = base64.b64encode(f.read()).decode("utf-8")
+
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {_api_key}"
+        }
+
+        payload = {
+            "model": "gpt-4-vision-preview",
+            "messages": [
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": prompt
+                        },
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": f"data:image/jpeg;base64,{encoded_image}",
+                                "detail": detail
+                            }
+                        }
+                    ]
+                }
+            ],
+            "max_tokens": max_tokens
+        }
+
+        r0 = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
+        if r0.status_code != 200:
+            raise openai.BadRequestError("Server returned an error", body=None, response=r0)
+
+        r1 = r0.json()
+        if "choices" not in r1 or len(r1["choices"]) < 1:
+            raise openai.BadRequestError("Empty results returned", body=None, response=r0)
+        r2 = r1["choices"][0]
+        if "finish_reason" not in r2 or r2["finish_reason"] != "stop":
+            raise openai.BadRequestError("Request was not completed correctly", body=None, response=r0)
+        return (r2["message"]["content"],)

negi/utils.py

@@ -0,0 +1,23 @@
+import os
+import glob
+import re
+
+
+def get_directory(directory: str) -> str:
+    base_path = os.path.abspath(__file__)
+    for _ in range(4):
+        base_path = os.path.dirname(base_path)
+    abs_path = os.path.abspath(os.path.join(base_path, directory))
+    os.makedirs(abs_path, exist_ok=True)
+    return abs_path
+
+
+def find_next_index(directory: str, glob_pattern="out.??????.png", re_pattern=r"out\.(\d{6})\.png") -> int:
+    next_index = 0
+    files = glob.glob(os.path.join(directory, glob_pattern))
+    for file in files:
+        r = re.match(re_pattern, os.path.basename(file))
+        if r is None:
+            continue
+        next_index = max(next_index, int(r.group(1)) + 1)
+    return next_index