Merge branch 'master' into develop

README.md

@@ -1,5 +1,10 @@
 # ComfyUI-NegiTools
 
+> [!IMPORTANT]
+> "Depth Estimation by Marigold (experimental)" module is not maintained and will be discontinued in the future; 
+> if you would like to continue using Marigold, please consider using this alternative choice.
+> https://github.com/kijai/ComfyUI-Marigold
+
 ## Installation
 
 - Install dependencies: pip install -r requirements.txt

__init__.py

@@ -10,6 +10,9 @@ from .negi.point_list_to_mask import PointListToMask
 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,
@@ -26,6 +29,9 @@ NODE_CLASS_MAPPINGS = {
     "NegiTools_StereoImageGenerator": StereoImageGenerator,
     "NegiTools_RandomImageLoader": RandomImageLoader,
     "NegiTools_SaveImageToDirectory": SaveImageToDirectory,
+    "NegiTools_DetectFaceRotationForInpainting": DetectFaceRotationForInpainting,
+    "NegiTools_OpenAiGpt4v": OpenAiGpt4v,
+    "NegiTools_OpenAiGpt": OpenAiGpt,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -43,4 +49,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "NegiTools_StereoImageGenerator": "Stereo Image Generator 🧅",
     "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/detect_face_rotation_for_inpainting.py

@@ -0,0 +1,118 @@
+import json
+
+import numpy as np
+import torch
+
+
+class DetectFaceRotationForInpainting:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "parts": ("STRING", {"multiline": False, "default": ""}),
+                "image": ("IMAGE",),
+                "radius_scale": ("FLOAT", {
+                    "default": 1.0,
+                    "min": 0.1,
+                    "max": 5.0,
+                    "step": 0.01,
+                    "round": 0.001,
+                    "display": "number"
+                }),
+                "overwrite_rotation": (["None", "0", "90", "180", "270"],),
+            }
+        }
+
+    RETURN_TYPES = ("INT", "MASK", "INT")
+    RETURN_NAMES = ("ROTATION_INV", "MASK", "ROTATION")
+    FUNCTION = "doit"
+    OUTPUT_NODE = False
+    CATEGORY = "utils"
+
+    @staticmethod
+    def get_face(xw, yw, radius_scale, parts):
+        x = 0.0
+        y = 0.0
+        n = 0
+        radius = 0
+        rot = 0
+
+        for name in ["Nose", "REye", "LEye", "REar", "LEar"]:
+            if name in parts:
+                x += parts[name]["x"]
+                y += parts[name]["y"]
+                n += 1
+
+        if n != 0:
+            x = x / n
+            y = y / n
+            for name in ["Nose", "REye", "LEye", "REar", "LEar"]:
+                if name in parts:
+                    x0 = x * xw
+                    y0 = y * yw
+                    x1 = parts[name]["x"] * xw
+                    y1 = parts[name]["y"] * yw
+                    radius = max(radius, int(np.sqrt((x1 - x0) * (x1 - x0) + (y1 - y0) * (y1 - y0)) * radius_scale))
+
+        if n != 0 and "Neck" in parts:
+            x0 = x * xw
+            y0 = y * yw
+            x1 = parts["Neck"]["x"] * xw
+            y1 = parts["Neck"]["y"] * yw
+            if abs(x1 - x0) < abs(y1 - y0):
+                rot = 0 if y0 < y1 else 180
+            else:
+                rot = 90 if x0 < x1 else 270
+
+        return x, y, radius, rot
+
+    @staticmethod
+    def rotate(rot, x, y):
+        if rot == 0:
+            return x, y
+        if rot == 90:
+            return y, 1 - x
+        if rot == 180:
+            return 1 - x, 1 - y
+        if rot == 270:
+            return 1 - y, x
+
+    def doit(self, parts, image, radius_scale, overwrite_rotation):
+        parts_list = json.loads(parts)
+        xw = image.shape[2]
+        yw = image.shape[1]
+
+        x = 0.0
+        y = 0.0
+        radius = 0
+        rot = 0
+        for parts in parts_list:
+            t_x, t_y, t_radius, t_rot = self.get_face(xw, yw, radius_scale, parts)
+            if t_radius > radius:
+                x = t_x
+                y = t_y
+                radius = t_radius
+                rot = t_rot
+
+        if overwrite_rotation != "None":
+            rot = int(overwrite_rotation)
+
+        rot_inv = (0 if rot == 0 else 360 - rot)
+        x_r, y_r = self.rotate(rot_inv, x, y)
+
+        xw_r = (xw if rot == 0 or rot == 180 else yw)
+        yw_r = (yw if rot == 0 or rot == 180 else xw)
+
+        if radius == 0:
+            return rot_inv, torch.from_numpy(np.zeros((1, yw_r, xw_r), dtype=np.float32)), rot
+
+        px = (np.reshape(np.arange(xw_r, dtype=np.float32), (1, -1))
+              * np.ones((yw_r, 1), dtype=np.float32))
+        py = (np.reshape(np.arange(yw_r, dtype=np.float32), (-1, 1))
+              * np.ones((1, xw_r), dtype=np.float32))
+        d2 = np.power(px - x_r * xw_r, 2.0) + np.power(py - y_r * yw_r, 2.0)
+        mask = torch.from_numpy(np.reshape(d2 <= radius * radius, (1, yw_r, xw_r)).astype(np.float32))
+        return rot_inv, mask, rot

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