Merge branch 'develop'

README.md

@@ -81,6 +81,15 @@ Each component of the output image is scaled in the range of 0.0 to 1.0.
 Depth estimation using Marigold.
 
 
+### Generator/Stereo Image Generator
+
+Generates stereo image; 
+This custom node calls the image transformation algorithm contained in the following extension for A1111
+(automatically cloned).
+https://github.com/thygate/stable-diffusion-webui-depthmap-script
+
+
+
 ### utils/OpenAI Translate to English
 
 Translates text written in any language into English using GPT-4.

__init__.py

@@ -8,6 +8,8 @@ from .negi.noise_image_generator import NoiseImageGenerator
 from .negi.open_pose_to_point_list import OpenPoseToPointList
 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
 
 NODE_CLASS_MAPPINGS = {
     "NegiTools_OpenAiDalle3": OpenAiDalle3,
@@ -21,6 +23,9 @@ NODE_CLASS_MAPPINGS = {
     "NegiTools_OpenPoseToPointList": OpenPoseToPointList,
     "NegiTools_PointListToMask": PointListToMask,
     "NegiTools_DepthEstimationByMarigold": DepthEstimationByMarigold,
+    "NegiTools_StereoImageGenerator": StereoImageGenerator,
+    "NegiTools_RandomImageLoader": RandomImageLoader,
+    "NegiTools_SaveImageToDirectory": SaveImageToDirectory,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -35,4 +40,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "NegiTools_OpenPoseToPointList": "OpenPose to Point List 🧅",
     "NegiTools_PointListToMask": "Point List to Mask 🧅",
     "NegiTools_DepthEstimationByMarigold": "Depth Estimation by Marigold (experimental) 🧅",
+    "NegiTools_StereoImageGenerator": "Stereo Image Generator 🧅",
+    "NegiTools_RandomImageLoader": "Random Image Loader 🧅",
+    "NegiTools_SaveImageToDirectory": "Save Image to Directory 🧅",
 }

negi/depth_estimation_by_marigold.py

@@ -1,6 +1,7 @@
 import os
 import sys
 import subprocess
+import gc
 
 import numpy as np
 import torch
@@ -76,6 +77,10 @@ class DepthEstimationByMarigold:
                    "bare (recommended)",
                    "venv (if \"bare\" doesn't work)",
                 ],),
+                "depth_exponent": ("FLOAT", {
+                    "default": 1.0, "min": 0.1, "max": 3.0, "step": 0.1, "round": 0.01, "display": "slider"
+                }),
+                "invert": ("BOOLEAN", {"default": False}),
             }
         }
 
@@ -85,7 +90,7 @@ class DepthEstimationByMarigold:
     OUTPUT_NODE = False
     CATEGORY = "Generator"
 
-    def doit(self, image, infer_passes, denoise_steps, seed, runtime):
+    def doit(self, image, infer_passes, denoise_steps, seed, runtime, depth_exponent, invert):
         use_venv = runtime.startswith("venv")
         self.__check_environment(use_venv)
 
@@ -110,6 +115,9 @@ class DepthEstimationByMarigold:
                 "output_dir_name": os.path.abspath(output_dir)
             })
 
+        # TODO It seems to be very bad idea, but it works; Try in-process execution
+        gc.collect()
+        torch.cuda.empty_cache()
         subprocess.run(["bash", os.path.join("work", "infer.sh")], cwd=self.rep_dir)
 
         im1 = np.load(os.path.join(output_dir, "depth_npy", "image_pred.npy")).astype(np.float32)
@@ -118,5 +126,8 @@ class DepthEstimationByMarigold:
         if im1_min == im1_max:
             im1_max = im1_min + 1.0
         im1 = (im1 - im1_min) * (1.0 / (im1_max - im1_min))
+        im1 = np.power(im1, depth_exponent)
+        if invert:
+            im1 = 1.0 - im1
 
         return (torch.from_numpy(np.expand_dims(np.stack([im1, im1, im1], axis=-1), axis=0)),)

negi/image_reader_writer.py

@@ -0,0 +1,97 @@
+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))
+
+
+class RandomImageLoader:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "directory": ("STRING", {"multiline": False, "default": "./input"}),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "doit"
+    OUTPUT_NODE = False
+    CATEGORY = "utils"
+
+    def doit(self, directory, seed):
+        directory = _get_directory(directory)
+        print("RandomImageLoader: directory = %s" % directory)
+
+        files = (glob.glob(os.path.join(directory, "*.png")) +
+                 glob.glob(os.path.join(directory, "*.jpg")) +
+                 glob.glob(os.path.join(directory, "*.jpeg")))
+
+        if len(files) == 0:
+            raise ValueError("Specified directory does not contain any image files")
+
+        file = files[seed % len(files)]
+        print("RandomImageLoader: load %s; in %d files" % (file, len(files)))
+
+        im0 = Image.open(file)
+        im1 = TF.to_tensor(im0.convert("RGBA"))
+        im1[:3, im1[3, :, :] == 0] = 0
+
+        images = torch.stack([im1])
+        images = images.permute(0, 2, 3, 1)
+        images = images[:, :, :, :3]
+        return (images,)
+
+
+class SaveImageToDirectory:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "directory": ("STRING", {"multiline": False, "default": "./output"}),
+                "image": ("IMAGE",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "doit"
+    OUTPUT_NODE = True
+    CATEGORY = "utils"
+
+    def doit(self, directory, image):
+        directory = _get_directory(directory)
+        os.makedirs(directory, exist_ok=True)
+        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)
+
+        file_name = os.path.join(directory, "out.%06d.png" % next_index)
+        print("SaveImageToDirectory: save to %s" % file_name)
+
+        im0 = torchvision.transforms.functional.to_pil_image(torch.permute(image[0], (2, 0, 1)))
+        im0.save(file_name)
+
+        return (image,)

negi/open_pose_to_point_list.py

@@ -1,8 +1,11 @@
 import json
 
 import numpy as np
+import cv2
+import torch
 from controlnet_aux import OpenposeDetector
-from controlnet_aux.util import HWC3, resize_image
+from controlnet_aux.util import HWC3
+from controlnet_aux.open_pose import draw_poses
 
 
 _names = [
@@ -17,6 +20,19 @@ _names = [
 _name_to_index = {name: i for i, name in enumerate(_names)}
 
 
+def _resize_image(input_image, resolution):
+    H, W, C = input_image.shape
+    H = float(H)
+    W = float(W)
+    k = float(resolution) / min(H, W)
+    H *= k
+    W *= k
+    H = int(np.round(H / 64.0)) * 64
+    W = int(np.round(W / 64.0)) * 64
+    img = cv2.resize(input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
+    return img, H, W
+
+
 class OpenPoseToPointList:
     def __init__(self):
         self.open_pose = OpenposeDetector.from_pretrained("lllyasviel/Annotators")
@@ -35,8 +51,8 @@ class OpenPoseToPointList:
             },
         }
 
-    RETURN_TYPES = ("STRING",)
-    RETURN_NAMES = ("POINT_LIST",)
+    RETURN_TYPES = ("STRING", "IMAGE")
+    RETURN_NAMES = ("POINT_LIST", "IMAGE")
     FUNCTION = "doit"
     OUTPUT_NODE = False
     CATEGORY = "utils"
@@ -44,9 +60,11 @@ class OpenPoseToPointList:
     def doit(self, image, detect_resolution, method):
         input_image = (np.fmax(0.0, np.fmin(1.0, image.to('cpu').detach().numpy()[0])) * 255.0).astype(np.uint8)
         input_image = HWC3(input_image)
-        input_image = resize_image(input_image, detect_resolution)
+        input_image, H, W = _resize_image(input_image, detect_resolution)
 
         poses = self.open_pose.detect_poses(input_image, include_hand=False, include_face=False)
+        img = draw_poses(poses, H, W, draw_hand=False, draw_face=False)
+        img = torch.from_numpy(np.expand_dims(HWC3(img) * (1.0 / 255), axis=0))
 
         if method == "face":
             ret = []
@@ -83,4 +101,4 @@ class OpenPoseToPointList:
         else:
             raise ValueError()
 
-        return (json.dumps(ret, indent=2),)
+        return (json.dumps(ret, indent=2), img)

negi/stereo_image_generator.py

@@ -0,0 +1,102 @@
+import os
+import subprocess
+import importlib
+
+import numpy as np
+import torch
+import torchvision
+from PIL import Image
+
+
+_dependency_dir = "dependencies"
+_repository_name = "stable-diffusion-webui-depthmap-script"
+
+
+class StereoImageGenerator:
+    def __check_environment(self):
+        if not os.path.isdir(os.path.join(self.dep_dir, _repository_name)):
+            r0 = subprocess.run([
+                "git", "clone", "https://github.com/thygate/stable-diffusion-webui-depthmap-script.git"
+            ], cwd=self.dep_dir)
+            if r0.returncode != 0:
+                subprocess.run(["rm", "-rf", _repository_name], cwd=self.dep_dir)
+                raise RuntimeError("Marigold repository not found or connection error")
+
+    def __init__(self):
+        self.dep_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), _dependency_dir)
+        os.makedirs(self.dep_dir, exist_ok=True)
+        self.rep_dir = os.path.join(self.dep_dir, _repository_name)
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "depth_image": ("IMAGE",),
+                "divergence": ("FLOAT", {
+                    "default": 5.0, "min": 0.05, "max": 10.0, "step": 0.01, "round": 0.001, "display": "slider"
+                }),
+                "stereo_offset_exponent": ("FLOAT", {
+                    "default": 1.0, "min": 0.1, "max": 3.0, "step": 0.1, "round": 0.01, "display": "slider"
+                }),
+                "fill_technique": ([
+                    "polylines_sharp", "polylines_soft", "naive", "naive_interpolating", "none"
+                ],),
+                "output_mode": ([
+                    "L-R", "R-L", "L-R-L",
+                ],),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE")
+    RETURN_NAMES = ("STEREO_IMAGE", "IMAGE_L", "IMAGE_R")
+    FUNCTION = "doit"
+    OUTPUT_NODE = False
+    CATEGORY = "Generator"
+
+    from .noise_image_generator import NoiseImageGenerator
+
+    def doit(self, image, depth_image, divergence, stereo_offset_exponent, fill_technique, output_mode):
+        self.__check_environment()
+        m = importlib.import_module(
+            "." + ".".join([_dependency_dir, _repository_name, "src", "stereoimage_generation"]),
+            ".".join(__name__.split(".")[:-2]))
+
+        xw = image.shape[2]
+        yw = image.shape[1]
+        image = torchvision.transforms.functional.to_pil_image(torch.permute(image[0], (2, 0, 1)))
+
+        depth_map = depth_image.to('cpu').detach().numpy()[0, :, :, 0]
+        depth_min = np.min(depth_map)
+        depth_max = np.max(depth_map)
+        if depth_max == depth_min:
+            depth_max = depth_min + 1.0
+        depth_map = (depth_map - depth_min) * (1.0 / (depth_max - depth_min))
+        depth_map = 1.0 - depth_map
+
+        modes = ["left-only", "only-right"]
+        if output_mode == "L-R":
+            modes.append("left-right")
+        elif output_mode == "R-L":
+            modes.append("right-left")
+        elif output_mode == "L-R-L":
+            pass
+        else:
+            raise ValueError()
+
+        images = m.create_stereoimages(
+            image, depth_map, divergence, modes=modes,
+            stereo_offset_exponent=stereo_offset_exponent, fill_technique=fill_technique)
+
+        if output_mode == "L-R-L":
+            out_image = Image.new("RGB", (xw * 3, yw))
+            out_image.paste(images[0], (0, 0))
+            out_image.paste(images[1], (xw, 0))
+            out_image.paste(images[0], (xw * 2, 0))
+            images.append(out_image)
+
+        return (
+            torch.from_numpy(np.expand_dims(np.array(images[2]) * (1.0 / 255), axis=0)),
+            torch.from_numpy(np.expand_dims(np.array(images[0]) * (1.0 / 255), axis=0)),
+            torch.from_numpy(np.expand_dims(np.array(images[1]) * (1.0 / 255), axis=0)),
+        )

requirements.txt

@@ -1,2 +1,3 @@
 openai >= 1.3.0
 controlnet-aux >= 0.0.7
+numba >= 0.58.1