agent_scheduler/task_helpers.py

import os
import io
import zlib
import base64
import inspect
import requests
import numpy as np
from typing import Union
from enum import Enum
from PIL import Image, ImageOps, ImageChops, ImageEnhance, ImageFilter

from modules import sd_samplers, scripts
from modules.generation_parameters_copypaste import create_override_settings_dict
from modules.sd_models import CheckpointInfo, get_closet_checkpoint_match
from modules.txt2img import txt2img
from modules.img2img import img2img
from modules.api.models import (
    StableDiffusionTxt2ImgProcessingAPI,
    StableDiffusionImg2ImgProcessingAPI,
)

from .helpers import log, get_dict_attribute

img2img_image_args_by_mode: dict[int, list[list[str]]] = {
    0: [["init_img"]],
    1: [["sketch"]],
    2: [["init_img_with_mask", "image"], ["init_img_with_mask", "mask"]],
    3: [["inpaint_color_sketch"], ["inpaint_color_sketch_orig"]],
    4: [["init_img_inpaint"], ["init_mask_inpaint"]],
}


def get_script_by_name(
    script_name: str, is_img2img: bool = False, is_always_on: bool = False
) -> scripts.Script:
    script_runner = scripts.scripts_img2img if is_img2img else scripts.scripts_txt2img
    available_scripts = (
        script_runner.alwayson_scripts
        if is_always_on
        else script_runner.selectable_scripts
    )

    return next(
        (s for s in available_scripts if s.title().lower() == script_name.lower()),
        None,
    )


def load_image_from_url(url: str):
    try:
        response = requests.get(url)
        buffer = io.BytesIO(response.content)
        return Image.open(buffer)
    except Exception as e:
        log.error(f"[AgentScheduler] Error downloading image from url: {e}")
        return None


def encode_image_to_base64(image):
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image.astype("uint8"))
    elif isinstance(image, str):
        if image.startswith("http://") or image.startswith("https://"):
            image = load_image_from_url(image)

    if not isinstance(image, Image.Image):
        return image

    with io.BytesIO() as output_bytes:
        image.save(output_bytes, format="PNG")
        bytes_data = output_bytes.getvalue()
        return "data:image/png;base64," + base64.b64encode(bytes_data).decode("utf-8")


def serialize_image(image):
    if isinstance(image, np.ndarray):
        shape = image.shape
        data = base64.b64encode(zlib.compress(image.tobytes())).decode()
        return {"shape": shape, "data": data, "cls": "ndarray"}
    elif isinstance(image, Image.Image):
        size = image.size
        mode = image.mode
        data = base64.b64encode(zlib.compress(image.tobytes())).decode()
        return {
            "size": size,
            "mode": mode,
            "data": data,
            "cls": "Image",
        }
    else:
        return image


def deserialize_image(image_str):
    if isinstance(image_str, dict) and image_str.get("cls", None):
        cls = image_str["cls"]
        data = zlib.decompress(base64.b64decode(image_str["data"]))

        if cls == "ndarray":
            shape = tuple(image_str["shape"])
            image = np.frombuffer(data, dtype=np.uint8)
            return image.reshape(shape)
        else:
            size = tuple(image_str["size"])
            mode = image_str["mode"]
            return Image.frombytes(mode, size, data)
    else:
        return image_str


def serialize_img2img_image_args(args: dict):
    for mode, image_args in img2img_image_args_by_mode.items():
        for keys in image_args:
            if mode != args["mode"]:
                # set None to unused image args to save space
                args[keys[0]] = None
            elif len(keys) == 1:
                image = args.get(keys[0], None)
                args[keys[0]] = serialize_image(image)
            else:
                value = args.get(keys[0], {})
                image = value.get(keys[1], None)
                value[keys[1]] = serialize_image(image)
                args[keys[0]] = value


def deserialize_img2img_image_args(args: dict):
    for mode, image_args in img2img_image_args_by_mode.items():
        if mode != args["mode"]:
            continue

        for keys in image_args:
            if len(keys) == 1:
                image = args.get(keys[0], None)
                args[keys[0]] = deserialize_image(image)
            else:
                value = args.get(keys[0], {})
                image = value.get(keys[1], None)
                value[keys[1]] = deserialize_image(image)
                args[keys[0]] = value


def serialize_controlnet_args(cnet_unit):
    args: dict = cnet_unit.__dict__
    args["is_cnet"] = True
    for k, v in args.items():
        if k == "image" and v is not None:
            args[k] = {
                "image": serialize_image(v["image"]),
                "mask": serialize_image(v["mask"])
                if v.get("mask", None) is not None
                else None,
            }
        if isinstance(v, Enum):
            args[k] = v.value

    return args


def deserialize_controlnet_args(args: dict):
    for k, v in args.items():
        if k == "image" and v is not None:
            args[k] = {
                "image": deserialize_image(v["image"]),
                "mask": deserialize_image(v["mask"])
                if v.get("mask", None) is not None
                else None,
            }

    return args


def map_controlnet_args_to_api_task_args(args: dict):
    if type(args).__name__ == "UiControlNetUnit":
        args = args.__dict__

    for k, v in args.items():
        if k == "image" and v is not None:
            args[k] = {
                "image": encode_image_to_base64(v["image"]),
                "mask": encode_image_to_base64(v["mask"])
                if v.get("mask", None) is not None
                else None,
            }
        if isinstance(v, Enum):
            args[k] = v.value

    return args


def map_ui_task_args_list_to_named_args(
    args: list, is_img2img: bool, checkpoint: str = None
):
    args_name = []
    if is_img2img:
        args_name = inspect.getfullargspec(img2img).args
    else:
        args_name = inspect.getfullargspec(txt2img).args

    named_args = dict(zip(args_name, args[0 : len(args_name)]))
    script_args = args[len(args_name) :]
    if checkpoint is not None:
        override_settings_texts = named_args.get("override_settings_texts", [])
        override_settings_texts.append("Model hash: " + checkpoint)
        named_args["override_settings_texts"] = override_settings_texts

    sampler_index = named_args.get("sampler_index", None)
    if sampler_index is not None:
        available_samplers = (
            sd_samplers.samplers_for_img2img if is_img2img else sd_samplers.samplers
        )
        sampler_name = available_samplers[named_args["sampler_index"]].name
        named_args["sampler_name"] = sampler_name
        log.debug(f"serialize sampler index: {str(sampler_index)} as {sampler_name}")

    return (
        named_args,
        script_args,
    )


def map_named_args_to_ui_task_args_list(
    named_args: dict, script_args: list, is_img2img: bool
):
    args_name = []
    if is_img2img:
        args_name = inspect.getfullargspec(img2img).args
    else:
        args_name = inspect.getfullargspec(txt2img).args

    sampler_name = named_args.get("sampler_name", None)
    if sampler_name is not None:
        available_samplers = (
            sd_samplers.samplers_for_img2img if is_img2img else sd_samplers.samplers
        )
        sampler_index = next(
            (i for i, x in enumerate(available_samplers) if x.name == sampler_name), 0
        )
        named_args["sampler_index"] = sampler_index

    args = [named_args.get(name, None) for name in args_name]
    args.extend(script_args)

    return args


def map_script_args_list_to_named(script: scripts.Script, args: list):
    script_name = script.title().lower()
    print("script", script_name, "is alwayson", script.alwayson)

    if script_name == "controlnet":
        for i, cnet_args in enumerate(args):
            args[i] = map_controlnet_args_to_api_task_args(cnet_args)

        return args

    fn = script.process if script.alwayson else script.run
    inspection = inspect.getfullargspec(fn)
    arg_names = inspection.args[2:]
    named_script_args = dict(zip(arg_names, args[: len(arg_names)]))
    if inspection.varargs is not None:
        named_script_args[inspection.varargs] = args[len(arg_names) :]

    return named_script_args


def map_named_script_args_to_list(
    script: scripts.Script, named_args: Union[dict, list]
):
    script_name = script.title().lower()

    if isinstance(named_args, dict):
        fn = script.process if script.alwayson else script.run
        inspection = inspect.getfullargspec(fn)
        arg_names = inspection.args[2:]
        args = [named_args.get(name, None) for name in arg_names]
        if inspection.varargs is not None:
            args.extend(named_args.get(inspection.varargs, []))

        return args

    if isinstance(named_args, list):
        if script_name == "controlnet":
            for i, cnet_args in enumerate(named_args):
                named_args[i] = map_controlnet_args_to_api_task_args(cnet_args)

        return named_args


def map_ui_task_args_to_api_task_args(
    named_args: dict, script_args: list, is_img2img: bool
):
    api_task_args: dict = named_args.copy()

    prompt_styles = api_task_args.pop("prompt_styles", [])
    api_task_args["styles"] = prompt_styles

    sampler_index = api_task_args.pop("sampler_index", 0)
    api_task_args["sampler_name"] = sd_samplers.samplers[sampler_index].name

    override_settings_texts = api_task_args.pop("override_settings_texts", [])
    api_task_args["override_settings"] = create_override_settings_dict(
        override_settings_texts
    )

    if is_img2img:
        mode = api_task_args.pop("mode", 0)
        for arg_mode, image_args in img2img_image_args_by_mode.items():
            if mode != arg_mode:
                for keys in image_args:
                    api_task_args.pop(keys[0], None)

        # the logic below is copied from modules/img2img.py
        if mode == 0:
            image = api_task_args.pop("init_img")
            image = image.convert("RGB") if image else None
            mask = None
        elif mode == 1:
            image = api_task_args.pop("sketch")
            image = image.convert("RGB") if image else None
            mask = None
        elif mode == 2:
            init_img_with_mask: dict = api_task_args.pop("init_img_with_mask") or {}
            image = init_img_with_mask.get("image", None)
            image = image.convert("RGB") if image else None
            mask = init_img_with_mask.get("mask", None)
            if mask:
                alpha_mask = (
                    ImageOps.invert(image.split()[-1])
                    .convert("L")
                    .point(lambda x: 255 if x > 0 else 0, mode="1")
                )
                mask = ImageChops.lighter(alpha_mask, mask.convert("L")).convert("L")
        elif mode == 3:
            image = api_task_args.pop("inpaint_color_sketch")
            orig = api_task_args.pop("inpaint_color_sketch_orig") or image
            if image is not None:
                mask_alpha = api_task_args.pop("mask_alpha", 0)
                mask_blur = api_task_args.get("mask_blur", 4)
                pred = np.any(np.array(image) != np.array(orig), axis=-1)
                mask = Image.fromarray(pred.astype(np.uint8) * 255, "L")
                mask = ImageEnhance.Brightness(mask).enhance(1 - mask_alpha / 100)
                blur = ImageFilter.GaussianBlur(mask_blur)
                image = Image.composite(image.filter(blur), orig, mask.filter(blur))
                image = image.convert("RGB")
        elif mode == 4:
            image = api_task_args.pop("init_img_inpaint")
            mask = api_task_args.pop("init_mask_inpaint")
        else:
            raise Exception(f"Batch mode is not supported yet")

        image = ImageOps.exif_transpose(image) if image else None
        api_task_args["init_images"] = [encode_image_to_base64(image)] if image else []
        api_task_args["mask"] = encode_image_to_base64(mask) if mask else None

        selected_scale_tab = api_task_args.pop("selected_scale_tab", 0)
        scale_by = api_task_args.get("scale_by", 1)
        if selected_scale_tab == 1 and image:
            api_task_args["width"] = int(image.width * scale_by)
            api_task_args["height"] = int(image.height * scale_by)
    else:
        hr_sampler_index = api_task_args.pop("hr_sampler_index", 0)
        api_task_args["hr_sampler_name"] = (
            sd_samplers.samplers_for_img2img[hr_sampler_index - 1].name
            if hr_sampler_index != 0
            else None
        )

    # script
    script_runner = scripts.scripts_img2img if is_img2img else scripts.scripts_txt2img
    script_id = script_args[0]
    if script_id == 0:
        api_task_args["script_name"] = None
        api_task_args["script_args"] = []
    else:
        script: scripts.Script = script_runner.selectable_scripts[script_id - 1]
        api_task_args["script_name"] = script.title().lower()
        current_script_args = script_args[script.args_from : script.args_to]
        api_task_args["script_args"] = map_script_args_list_to_named(
            script, current_script_args
        )

    # alwayson scripts
    alwayson_scripts = api_task_args.get("alwayson_scripts", None)
    if not alwayson_scripts:
        api_task_args["alwayson_scripts"] = {}
        alwayson_scripts = api_task_args["alwayson_scripts"]

    for script in script_runner.alwayson_scripts:
        alwayson_script_args = script_args[script.args_from : script.args_to]
        script_name = script.title().lower()
        if script_name != "agent scheduler":
            named_script_args = map_script_args_list_to_named(
                script, alwayson_script_args
            )
            alwayson_scripts[script_name] = {"args": named_script_args}

    return api_task_args


def serialize_api_task_args(
    params: dict,
    is_img2img: bool,
    checkpoint: str = None,
):
    # handle named script args
    script_name = params.get("script_name", None)
    if script_name is not None:
        script = get_script_by_name(script_name, is_img2img)
        if script is None:
            raise Exception(f"Not found script {script_name}")

        script_args = params.get("script_args", {})
        params["script_args"] = map_named_script_args_to_list(script, script_args)

    # handle named alwayson script args
    alwayson_scripts = get_dict_attribute(params, "alwayson_scripts", {})
    valid_alwayson_scripts = {}
    script_runner = scripts.scripts_img2img if is_img2img else scripts.scripts_txt2img
    for script in script_runner.alwayson_scripts:
        script_name = script.title().lower()
        if script_name == "agent scheduler":
            continue

        script_args = get_dict_attribute(alwayson_scripts, f"{script_name}.args", None)
        if script_args:
            arg_list = map_named_script_args_to_list(script, script_args)
            valid_alwayson_scripts[script_name] = {"args": arg_list}

    params["alwayson_scripts"] = valid_alwayson_scripts

    args = (
        StableDiffusionImg2ImgProcessingAPI(**params)
        if is_img2img
        else StableDiffusionTxt2ImgProcessingAPI(**params)
    )

    if args.override_settings is None:
        args.override_settings = {}

    if checkpoint is not None:
        checkpoint_info: CheckpointInfo = get_closet_checkpoint_match(checkpoint)
        if not checkpoint_info:
            raise Exception(f"No checkpoint found for model hash {checkpoint}")
        args.override_settings["sd_model_checkpoint"] = checkpoint_info.title

    # load images from url or file if needed
    if is_img2img:
        init_images = args.init_images
        if len(init_images) == 0:
            raise Exception("At least one init image is required")

        for i, image in enumerate(init_images):
            init_images[i] = encode_image_to_base64(image)

        args.mask = encode_image_to_base64(args.mask)
        args.batch_size = len(init_images)

    return args.dict()