120 lines
4.9 KiB
Python
120 lines
4.9 KiB
Python
from __future__ import annotations
|
||
|
||
import copy
|
||
import json
|
||
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)),
|
||
"save": (None, None),
|
||
},
|
||
}
|
||
|
||
LATENT_CLASSES = {"EmptySD3LatentImage", "EmptyLatentImage", "EmptyLatentImagePresets"}
|
||
|
||
|
||
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,
|
||
) -> 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
|
||
return workflow
|