Stable Diffusion via REST

backend/src/nodes/impl/external_stable_diffusion.py

@@ -0,0 +1,96 @@
+import requests
+import base64
+import cv2
+import io
+import numpy as np
+import os
+from PIL import Image
+from enum import Enum
+from typing import Dict, Union
+from sanic.log import logger
+
+from .image_utils import normalize
+from ..utils.utils import get_h_w_c
+
+STABLE_DIFFUSION_HOST = os.environ.get("STABLE_DIFFUSION_HOST", "127.0.0.1")
+STABLE_DIFFUSION_PORT = os.environ.get("STABLE_DIFFUSION_PORT", "7860")
+
+STABLE_DIFFUSION_TEXT2IMG_URL = (
+    f"http://{STABLE_DIFFUSION_HOST}:{STABLE_DIFFUSION_PORT}/sdapi/v1/txt2img"
+)
+STABLE_DIFFUSION_IMG2IMG_URL = (
+    f"http://{STABLE_DIFFUSION_HOST}:{STABLE_DIFFUSION_PORT}/sdapi/v1/img2img"
+)
+STABLE_DIFFUSION_INTERROGATE_URL = (
+    f"http://{STABLE_DIFFUSION_HOST}:{STABLE_DIFFUSION_PORT}/sdapi/v1/interrogate"
+)
+STABLE_DIFFUSION_OPTIONS_URL = (
+    f"http://{STABLE_DIFFUSION_HOST}:{STABLE_DIFFUSION_PORT}/sdapi/v1/options"
+)
+
+
+def get(url) -> Dict:
+    response = requests.get(url)
+    return response.json()
+
+
+def post(url, json_data: Dict) -> Dict:
+    response = requests.post(url, json=json_data)
+    return response.json()
+
+
+# Call the API at import time
+# If this fails (because the API isn't up) it will raise an exception and the
+# nodes importing this file won't be registered
+STABLE_DIFFUSION_OPTIONS = get(STABLE_DIFFUSION_OPTIONS_URL)
+
+
+def decode_base64_image(image_bytes: Union[bytes, str]) -> np.ndarray:
+    image = Image.open(io.BytesIO(base64.b64decode(image_bytes)))
+    image_nparray = np.array(image)
+    _, _, c = get_h_w_c(image_nparray)
+    if c == 3:
+        image_nparray = cv2.cvtColor(image_nparray, cv2.COLOR_RGB2BGR)
+    elif c == 4:
+        image_nparray = cv2.cvtColor(image_nparray, cv2.COLOR_RGBA2BGRA)
+    return normalize(image_nparray)
+
+
+def encode_base64_image(image_nparray: np.ndarray) -> str:
+    image_nparray = (np.clip(image_nparray, 0, 1) * 255).round().astype("uint8")
+    _, _, c = get_h_w_c(image_nparray)
+    if c == 1:
+        # PIL supports grayscale images just fine, so we don't need to do any conversion
+        pass
+    elif c == 3:
+        image_nparray = cv2.cvtColor(image_nparray, cv2.COLOR_BGR2RGB)
+    elif c == 4:
+        image_nparray = cv2.cvtColor(image_nparray, cv2.COLOR_BGRA2RGBA)
+    else:
+        raise RuntimeError
+    with io.BytesIO() as buffer:
+        with Image.fromarray(image_nparray) as image:
+            image.save(buffer, format="PNG")
+        return base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+
+class SamplerName(Enum):
+    EULER = "Euler"
+    EULER_A = "Euler a"
+    LMS = "LMS"
+    HEUN = "Heun"
+    DPM2 = "DPM2"
+    DPM2_A = "DPM2 a"
+    DPMpp_2S_A = "DPM++ 2S a"
+    DPMpp_2M = "DPM++ 2M"
+    DPMpp_SDE = "DPM++ SDE"
+    DPM_FAST = "DPM fast"
+    DPM_A = "DPM adaptive"
+    LMS_KARRAS = "LMS Karras"
+    DPM2_KARRAS = "DPM2 Karras"
+    DPM2_A_KARRAS = "DPM2 a Karras"
+    DPMpp_2S_A_KARRAS = "DPM++ 2S a Karras"
+    DPMpp_2M_KARRAS = "DPM++ 2M Karras"
+    DPMpp_SDE_KARRAS = "DPM++ SDE Karras"
+    DDIM = "DDIM"
+    PLMS = "PLMS"

backend/src/nodes/nodes/builtin_categories.py

@@ -4,6 +4,7 @@
 from .image_dimension import category as ImageDimensionCategory
 from .image_channel import category as ImageChannelCategory
 from .image_utility import category as ImageUtilityCategory
+from .external_stable_diffusion import category as ExternalStableDiffusionCategory
 from .utility import category as UtilityCategory
 from .pytorch import category as PyTorchCategory
 from .ncnn import category as NCNNCategory
@@ -21,5 +22,6 @@
     PyTorchCategory,
     NCNNCategory,
     ONNXCategory,
+    ExternalStableDiffusionCategory,
 ]
 category_order = [x.name for x in builtin_categories]

backend/src/nodes/nodes/external_stable_diffusion/__init__.py

@@ -0,0 +1,8 @@
+from ...category import Category
+
+category = Category(
+    name="Stable Diffusion (External)",
+    description="Interact with am external Stable Diffusion API",
+    icon="BsFillImageFill",
+    color="#C53030",
+)

backend/src/nodes/nodes/external_stable_diffusion/img2img.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import numpy as np
+
+from . import category as ExternalStableDiffusionCategory
+from ...impl.external_stable_diffusion import (
+    decode_base64_image,
+    SamplerName,
+    STABLE_DIFFUSION_IMG2IMG_URL,
+    post,
+    encode_base64_image,
+)
+from ...node_base import NodeBase, group
+from ...node_factory import NodeFactory
+from ...properties.inputs import (
+    TextInput,
+    NumberInput,
+    SliderInput,
+    EnumInput,
+    ImageInput,
+)
+from ...properties.outputs import ImageOutput
+from typing import Optional
+
+
+@NodeFactory.register("chainner:external_stable_diffusion:img2img")
+class Img2Img(NodeBase):
+    def __init__(self):
+        super().__init__()
+        self.description = "Modify an image using Automatic1111"
+        self.inputs = [
+            ImageInput(),
+            TextInput("Prompt", default="an astronaut riding a horse"),
+            TextInput("Negative Prompt").make_optional(),
+            SliderInput(
+                "Denoising Strength",
+                minimum=0,
+                default=0.75,
+                maximum=1,
+                slider_step=0.01,
+                controls_step=0.1,
+                precision=2,
+            ),
+            group("seed")(
+                NumberInput("Seed", minimum=0, default=42, maximum=4294967296)
+            ),
+            SliderInput("Steps", minimum=1, default=20, maximum=150),
+            EnumInput(SamplerName, default_value=SamplerName.EULER),
+            SliderInput(
+                "CFG Scale",
+                minimum=1,
+                default=7,
+                maximum=20,
+                controls_step=0.1,
+                precision=1,
+            ),
+            SliderInput("Width", minimum=64, default=512, maximum=2048).with_id(8),
+            SliderInput("Height", minimum=64, default=512, maximum=2048).with_id(9),
+            TextInput("Model Checkpoint Override").make_optional(),
+        ]
+        self.outputs = [
+            ImageOutput(
+                image_type="Image {width: Input8, height: Input9, channels: 3}"
+            ),
+        ]
+
+        self.category = ExternalStableDiffusionCategory
+        self.name = "Image to Image"
+        self.icon = "MdChangeCircle"
+        self.sub = "Automatic1111"
+
+    def run(
+        self,
+        image: np.ndarray,
+        prompt: str,
+        negative_prompt: Optional[str],
+        denoising_strength: float,
+        seed: int,
+        steps: int,
+        sampler_name: SamplerName,
+        cfg_scale: float,
+        width: int,
+        height: int,
+        sd_model_checkpoint: Optional[str],
+    ) -> np.ndarray:
+        request_data = {
+            "init_images": [encode_base64_image(image)],
+            "prompt": prompt,
+            "negative_prompt": negative_prompt or "",
+            "denoising_strength": denoising_strength,
+            "seed": seed,
+            "steps": steps,
+            "sampler_name": sampler_name.value,
+            "cfg_scale": cfg_scale,
+            "width": width,
+            "height": height,
+            "override_settings": {},
+        }
+        if sd_model_checkpoint:
+            request_data["override_settings"][
+                "sd_model_checkpoint"
+            ] = sd_model_checkpoint
+        response = post(url=STABLE_DIFFUSION_IMG2IMG_URL, json_data=request_data)
+        return decode_base64_image(response["images"][0])

backend/src/nodes/nodes/external_stable_diffusion/interrogate.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+import numpy as np
+
+from . import category as ExternalStableDiffusionCategory
+from ...impl.external_stable_diffusion import (
+    STABLE_DIFFUSION_INTERROGATE_URL,
+    post,
+    encode_base64_image,
+)
+from ...node_base import NodeBase
+from ...node_factory import NodeFactory
+from ...properties.inputs import ImageInput
+from ...properties.outputs import TextOutput
+
+
+@NodeFactory.register("chainner:external_stable_diffusion:interrograte")
+class Interrogate(NodeBase):
+    def __init__(self):
+        super().__init__()
+        self.description = "Use Automatic1111 to get a description of an image"
+        self.inputs = [
+            ImageInput(),
+        ]
+        self.outputs = [
+            TextOutput("Text"),
+        ]
+
+        self.category = ExternalStableDiffusionCategory
+        self.name = "CLIP Interrogate"
+        self.icon = "MdTextFields"
+        self.sub = "Automatic1111"
+
+    def run(self, image: np.ndarray) -> str:
+        request_data = {
+            "image": encode_base64_image(image),
+        }
+        response = post(url=STABLE_DIFFUSION_INTERROGATE_URL, json_data=request_data)
+        return response["caption"]

backend/src/nodes/nodes/external_stable_diffusion/txt2img.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import numpy as np
+
+from . import category as ExternalStableDiffusionCategory
+from ...impl.external_stable_diffusion import (
+    decode_base64_image,
+    SamplerName,
+    STABLE_DIFFUSION_TEXT2IMG_URL,
+    post,
+)
+from ...node_base import NodeBase, group
+from ...node_factory import NodeFactory
+from ...properties.inputs import (
+    TextInput,
+    NumberInput,
+    SliderInput,
+    EnumInput,
+)
+from ...properties.outputs import ImageOutput
+from typing import Optional
+
+
+@NodeFactory.register("chainner:external_stable_diffusion:txt2img")
+class Txt2Img(NodeBase):
+    def __init__(self):
+        super().__init__()
+        self.description = "Generate an image using Automatic1111"
+        self.inputs = [
+            TextInput("Prompt", default="an astronaut riding a horse"),
+            TextInput("Negative Prompt").make_optional(),
+            group("seed")(
+                NumberInput("Seed", minimum=0, default=42, maximum=4294967296)
+            ),
+            SliderInput("Steps", minimum=1, default=20, maximum=150),
+            EnumInput(SamplerName, default_value=SamplerName.EULER),
+            SliderInput(
+                "CFG Scale",
+                minimum=1,
+                default=7,
+                maximum=20,
+                controls_step=0.1,
+                precision=1,
+            ),
+            SliderInput("Width", minimum=64, default=512, maximum=2048).with_id(6),
+            SliderInput("Height", minimum=64, default=512, maximum=2048).with_id(7),
+            TextInput("Model Checkpoint Override").make_optional(),
+        ]
+        self.outputs = [
+            ImageOutput(
+                image_type="Image {width: Input6, height: Input7, channels: 3}"
+            ),
+        ]
+
+        self.category = ExternalStableDiffusionCategory
+        self.name = "Text to Image"
+        self.icon = "BsFillImageFill"
+        self.sub = "Automatic1111"
+
+    def run(
+        self,
+        prompt: str,
+        negative_prompt: Optional[str],
+        seed: int,
+        steps: int,
+        sampler_name: SamplerName,
+        cfg_scale: float,
+        width: int,
+        height: int,
+        sd_model_checkpoint: Optional[str],
+    ) -> np.ndarray:
+        request_data = {
+            "prompt": prompt,
+            "negative_prompt": negative_prompt or "",
+            "seed": seed,
+            "steps": steps,
+            "sampler_name": sampler_name.value,
+            "cfg_scale": cfg_scale,
+            "width": width,
+            "height": height,
+            "override_settings": {},
+        }
+        if sd_model_checkpoint:
+            request_data["override_settings"][
+                "sd_model_checkpoint"
+            ] = sd_model_checkpoint
+        response = post(url=STABLE_DIFFUSION_TEXT2IMG_URL, json_data=request_data)
+        return decode_base64_image(response["images"][0])

requirements.txt

@@ -17,6 +17,7 @@ pylint
 pyright==1.1.280
 pytest
 pywin32 ; sys_platform=="win32"
+requests
 sanic
 sanic_cors
 # torch

src/common/dependencies.ts

@@ -125,6 +125,10 @@ export const requiredDependencies: Dependency[] = [
         name: 'FFMPEG',
         packages: [{ packageName: 'ffmpeg-python', version: '0.2.0', sizeEstimate: 25 * KB }],
     },
+    {
+        name: 'Requests',
+        packages: [{ packageName: 'requests', version: '2.28.2', sizeEstimate: 452 * KB }],
+    },
 ];
 
 if (isMac && !isM1) {