add ControlNet extension support

setup.py

@@ -11,7 +11,7 @@
 
 setup(
     name="webuiapi",
-    version="0.3.3",
+    version="0.4.0",
     description="Python API client for AUTOMATIC1111/stable-diffusion-webui",
     url="https://github.com/mix1009/sdwebuiapi",
     author="ChunKoo Park",

webuiapi/__init__.py

@@ -1,6 +1,6 @@
-from .webuiapi import WebUIApi, WebUIApiResult, Upscaler, HiResUpscaler, b64_img, ModelKeywordResult, ModelKeywordInterface, InstructPix2PixInterface
+from .webuiapi import WebUIApi, WebUIApiResult, Upscaler, HiResUpscaler, b64_img, raw_b64_img, ModelKeywordResult, ModelKeywordInterface, InstructPix2PixInterface, ControlNetInterface
 
-__version__ = "0.3.3"
+__version__ = "0.4.0"
 
 __all__ = [
     "__version__",
@@ -11,6 +11,7 @@
     "b64_img",
     "ModelKeywordResult",
     "ModelKeywordInterface",
-    "InstructPix2PixInterface",   
+    "InstructPix2PixInterface",
+    "ControlNetInterface",
 ]
 

webuiapi/webuiapi.py

@@ -5,6 +5,7 @@
 from PIL import Image
 from dataclasses import dataclass
 from enum import Enum
+from typing import List, Dict, Any
 
 class Upscaler(str, Enum):
     none = 'None'
@@ -51,6 +52,13 @@ def b64_img(image: Image):
     img_base64 = 'data:image/png;base64,' + str(base64.b64encode(buffered.getvalue()), 'utf-8')
     return img_base64
 
+def raw_b64_img(x):
+    # XXX controlnet only accepts RAW base64 without headers
+    buffered = io.BytesIO()
+    image.save(buffered, format="PNG")
+    img_base64 = str(base64.b64encode(buffered.getvalue()), 'utf-8')
+    return img_base64
+
 class WebUIApi:
     def __init__(self,
                  host='127.0.0.1',
@@ -555,3 +563,167 @@ def img2img(self,
         }
         return self.api.custom_post('instruct-pix2pix/img2img', payload=payload)
 
+
+# https://github.com/Mikubill/sd-webui-controlnet
+class ControlNetInterface:
+    def __init__(self, webuiapi):
+        self.api = webuiapi
+
+    def txt2img(self,
+        prompt: str = "",
+        negative_prompt: str = "",
+        controlnet_input_image: [] = [],
+        controlnet_mask: [] = [],
+        controlnet_module: str = "",
+        controlnet_model: str = "",
+        controlnet_weight: float = 0.5,
+        controlnet_resize_mode: str = "Scale to Fit (Inner Fit)",
+        controlnet_lowvram: bool = False,
+        controlnet_processor_res: int = 512,
+        controlnet_threshold_a: int = 64,
+        controlnet_threshold_b: int = 64,
+        controlnet_guidance: float = 1.0,
+        enable_hr: bool = False, # hiresfix
+        denoising_strength: float = 0.5,
+        hr_scale: float = 1.5,
+        hr_upscale: str = "Latent",
+        guess_mode: bool = True,
+        seed: int = -1,
+        subseed: int = -1,
+        subseed_strength: int = -1,
+        sampler_index: str = "Euler a",
+        batch_size: int = 1,
+        n_iter: int = 1, # Iteration
+        steps: int = 20,
+        cfg_scale: float = 7,
+        width: int = 512,
+        height: int = 512,
+        restore_faces: bool = False,
+        override_settings: Dict[str, Any] = None, 
+        override_settings_restore_afterwards: bool = True):
+
+        controlnet_input_image_b64 = [raw_b64_img(x) for x in controlnet_input_image]
+        controlnet_mask_b64 = [raw_b64_img(x) for x in controlnet_mask]
+
+        payload = {
+            "prompt": prompt,
+            "negative_prompt": negative_prompt,
+            "controlnet_input_image": controlnet_input_image_b64,
+            "controlnet_mask": controlnet_mask_b64,
+            "controlnet_module": controlnet_module,
+            "controlnet_model": controlnet_model,
+            "controlnet_weight": controlnet_weight,
+            "controlnet_resize_mode": controlnet_resize_mode,
+            "controlnet_lowvram": controlnet_lowvram,
+            "controlnet_processor_res": controlnet_processor_res,
+            "controlnet_threshold_a": controlnet_threshold_a,
+            "controlnet_threshold_b": controlnet_threshold_b,
+            "controlnet_guidance": controlnet_guidance,
+            "enable_hr": enable_hr,
+            "denoising_strength": denoising_strength,
+            "hr_scale": hr_scale,
+            "hr_upscale": hr_upscale,
+            "guess_mode": guess_mode,
+            "seed": seed,
+            "subseed": subseed,
+            "subseed_strength": subseed_strength,
+            "sampler_index": sampler_index,
+            "batch_size": batch_size,
+            "n_iter": n_iter,
+            "steps": steps,
+            "cfg_scale": cfg_scale,
+            "width": width,
+            "height": height,
+            "restore_faces": restore_faces,
+            "override_settings": override_settings,
+            "override_settings_restore_afterwards": override_settings_restore_afterwards,
+        }
+        return self.api.custom_post('controlnet/txt2img', payload=payload)
+
+    def img2img(self,
+        init_images: [] = [],
+        mask: str = None,
+        mask_blur: int = 30,
+        inpainting_fill: int = 0,
+        inpaint_full_res: bool = True,
+        inpaint_full_res_padding: int = 1,
+        inpainting_mask_invert: int = 1,
+        resize_mode: int = 0,
+        denoising_strength: float = 0.7,
+        prompt: str = "",
+        negative_prompt: str = "",
+        controlnet_input_image: [] = [],
+        controlnet_mask: [] = [],
+        controlnet_module: str = "",
+        controlnet_model: str = "",
+        controlnet_weight: float = 1.0,
+        controlnet_resize_mode: str = "Scale to Fit (Inner Fit)",
+        controlnet_lowvram: bool = False,
+        controlnet_processor_res: int = 512,
+        controlnet_threshold_a: int = 64,
+        controlnet_threshold_b: int = 64,
+        controlnet_guidance: float = 1.0,
+        guess_mode: bool = True,
+        seed: int = -1,
+        subseed: int = -1,
+        subseed_strength: int = -1,
+        sampler_index: str = "",
+        batch_size: int = 1,
+        n_iter: int = 1, # Iteration
+        steps: int = 20,
+        cfg_scale: float = 7,
+        width: int = 512,
+        height: int = 512,
+        restore_faces: bool = False,
+        include_init_images: bool = True,
+        override_settings: Dict[str, Any] = None, 
+        override_settings_restore_afterwards: bool = True):
+
+        init_images_b64 = [raw_b64_img(x) for x in init_images]
+        controlnet_input_image_b64 = [raw_b64_img(x) for x in controlnet_input_image]
+        controlnet_mask_b64 = [raw_b64_img(x) for x in controlnet_mask]
+
+        payload = {
+            "init_images": init_images_b64,
+            "mask": raw_b64_img(mask) if mask else None,
+            "mask_blur": mask_blur,
+            "inpainting_fill": inpainting_fill,
+            "inpaint_full_res": inpaint_full_res,
+            "inpaint_full_res_padding": inpaint_full_res_padding,
+            "inpainting_mask_invert": inpainting_mask_invert,
+            "resize_mode": resize_mode,
+            "denoising_strength": denoising_strength,
+            "prompt": prompt,
+            "negative_prompt": negative_prompt,
+            "controlnet_input_image": controlnet_input_image_b64,
+            "controlnet_mask": controlnet_mask_b64,
+            "controlnet_module": controlnet_module,
+            "controlnet_model": controlnet_model,
+            "controlnet_weight": controlnet_weight,
+            "controlnet_resize_mode": controlnet_resize_mode,
+            "controlnet_lowvram": controlnet_lowvram,
+            "controlnet_processor_res": controlnet_processor_res,
+            "controlnet_threshold_a": controlnet_threshold_a,
+            "controlnet_threshold_b": controlnet_threshold_b,
+            "controlnet_guidance": controlnet_guidance,
+            "guess_mode": guess_mode,
+            "seed": seed,
+            "subseed": subseed,
+            "subseed_strength": subseed_strength,
+            "sampler_index": sampler_index,
+            "batch_size": batch_size,
+            "n_iter": n_iter,
+            "steps": steps,
+            "cfg_scale": cfg_scale,
+            "width": width,
+            "height": height,
+            "restore_faces": restore_faces,
+            "include_init_images": include_init_images,
+            "override_settings": override_settings,
+            "override_settings_restore_afterwards": override_settings_restore_afterwards,
+        }
+        return self.api.custom_post('controlnet/img2img', payload=payload)
+
+    def model_list(self):
+        r = self.api.custom_get('controlnet/model_list')
+        return r['model_list']