MpakoImageAIbot/workflows.py

from __future__ import annotations

import copy
import json
import math
from collections.abc import Callable
from pathlib import Path
from typing import Any

RoleField = tuple[str | None, Callable[[Any], bool] | None]

ROLE_FIELDS: dict[str, dict[str, RoleField]] = {
    "txt2img": {
        "positive": ("text", lambda value: isinstance(value, str)),
        "negative": ("text", lambda value: isinstance(value, str)),
        "seed": ("seed", lambda value: isinstance(value, int)),
        "latent": (None, None),
        "save": (None, None),
    },
    "img2img": {
        "positive": ("prompt", lambda value: isinstance(value, str)),
        "negative": ("text", lambda value: isinstance(value, str)),
        "seed": ("noise_seed", lambda value: isinstance(value, int)),
        "input_image": ("image", lambda value: isinstance(value, str)),
        "latent": (None, None),
        "save": (None, None),
    },
}

LATENT_CLASSES = {"EmptySD3LatentImage", "EmptyLatentImage", "EmptyLatentImagePresets"}


def fit_latent_size(
    width: int,
    height: int,
    *,
    target_pixels: int = 1024 * 1024,
    multiple: int = 16,
    max_side: int = 1536,
) -> tuple[int, int]:
    if width <= 0 or height <= 0:
        return 1024, 1024

    ratio = math.sqrt(target_pixels / (width * height))
    new_width = width * ratio
    new_height = height * ratio
    longest = max(new_width, new_height)
    if longest > max_side:
        scale = max_side / longest
        new_width *= scale
        new_height *= scale

    new_width = max(multiple, round(new_width / multiple) * multiple)
    new_height = max(multiple, round(new_height / multiple) * multiple)
    return int(new_width), int(new_height)


class WorkflowError(RuntimeError):
    pass


class WorkflowBuilder:
    def __init__(self, t2i_path: Path, i2i_path: Path, mapping_path: Path) -> None:
        self._paths = {"txt2img": t2i_path, "img2img": i2i_path}
        self._templates = {mode: self._load_json(path) for mode, path in self._paths.items()}
        self._mapping = self._load_json(mapping_path)
        self._validate()

    @staticmethod
    def _load_json(path: Path) -> dict[str, Any]:
        try:
            data = json.loads(path.read_text(encoding="utf-8"))
        except FileNotFoundError as exc:
            raise WorkflowError(f"Не найден workflow: {path}") from exc
        except json.JSONDecodeError as exc:
            raise WorkflowError(f"Некорректный JSON в {path}: {exc}") from exc
        if not isinstance(data, dict) or not data:
            raise WorkflowError(f"Workflow {path} должен быть непустым объектом")
        return data

    def _validate(self) -> None:
        if not isinstance(self._mapping, dict):
            raise WorkflowError("mapping.json должен быть объектом")

        for mode, roles in ROLE_FIELDS.items():
            mode_mapping = self._mapping.get(mode)
            if not isinstance(mode_mapping, dict):
                raise WorkflowError(f"В mapping.json отсутствует раздел {mode!r}")
            template = self._templates[mode]

            for role, (field_name, check) in roles.items():
                node_id = mode_mapping.get(role)
                if not isinstance(node_id, str) or node_id not in template:
                    raise WorkflowError(
                        f"{mode}: узел для роли {role!r} не найден в workflow"
                    )
                node = template[node_id]
                inputs = node.get("inputs")
                if not isinstance(inputs, dict):
                    raise WorkflowError(f"{mode}: у узла {node_id} нет inputs")

                if role == "latent":
                    if node.get("class_type") not in LATENT_CLASSES:
                        raise WorkflowError(
                            f"{mode}: узел latent должен создавать пустой латент"
                        )
                    if "width" not in inputs or "height" not in inputs:
                        raise WorkflowError(f"{mode}: у latent нет width/height")
                    continue

                if field_name is None:
                    continue
                if field_name not in inputs:
                    raise WorkflowError(
                        f"{mode}: у узла {node_id} нет поля {field_name!r} "
                        f"для роли {role!r}"
                    )
                if check is not None and not check(inputs[field_name]):
                    raise WorkflowError(
                        f"{mode}: поле {field_name!r} у узла {node_id} "
                        f"имеет неожиданный тип"
                    )

    def build_txt2img(self, *, prompt: str, negative: str, seed: int) -> dict[str, Any]:
        workflow = copy.deepcopy(self._templates["txt2img"])
        mapping = self._mapping["txt2img"]
        workflow[mapping["positive"]]["inputs"]["text"] = prompt
        workflow[mapping["negative"]]["inputs"]["text"] = negative
        workflow[mapping["seed"]]["inputs"]["seed"] = seed
        return workflow

    def build_img2img(
        self,
        *,
        prompt: str,
        negative: str,
        seed: int,
        image_name: str,
        width: int = 1024,
        height: int = 1024,
    ) -> dict[str, Any]:
        workflow = copy.deepcopy(self._templates["img2img"])
        mapping = self._mapping["img2img"]
        workflow[mapping["positive"]]["inputs"]["prompt"] = prompt
        workflow[mapping["negative"]]["inputs"]["text"] = negative
        workflow[mapping["seed"]]["inputs"]["noise_seed"] = seed
        workflow[mapping["input_image"]]["inputs"]["image"] = image_name
        latent = workflow[mapping["latent"]]["inputs"]
        latent["width"] = width
        latent["height"] = height
        return workflow