diff --git a/README.md b/README.md index a11ab45..63e920b 100644 --- a/README.md +++ b/README.md @@ -208,10 +208,34 @@ node/ ### Backend Selection - `backend_mode = auto | sdxl | flux` -- `auto` 会根据 capability metadata 自动路由，找不到专业后端时直接报错，不会静默降级到旧的 Simple repair 逻辑 +- `auto` 现在使用 profile-driven routing：先解析模型 profile，再按 `backend_type` 自动路由 +- 找不到专业后端时直接报错，不会静默降级到旧的 Simple repair 逻辑 - `sdxl` 和 `flux` 会强制指定后端，但仍然会验证当前模型是否兼容 -### Capability Requirements +### Profile-Driven Routing + +LLS Simple Checkpoint Loader writes the resolved model profile into runtime metadata. + +- `LLS Pro Image Edit Prepare` routes by `backend_type` +- `LLS Pro KSampler Bridge` routes by `sampler_strategy` + +重要 profile 字段： + +- `profile_id` +- `backend_type` +- `sampler_strategy` + +示例： + +- `sdxl_inpaint` -> `backend_type = sdxl_native` +- `flux_edit` -> `backend_type = flux_edit` + +Base profile 不会再模糊地进入 Pro 链： + +- `SDXL` base 仍然保持 base profile，除非显式 override +- `FLUX base models should remain on the Simple workflow unless they resolve to flux_edit` + +### Compatibility Metadata - `model_role` - `supports_inpaint_native` @@ -223,13 +247,14 @@ node/ - `model_info` - `LLS Simple Checkpoint Loader` 写入到对象上的 `_lls_*` 元信息 - `utils/model_info.py` 里的 family / name 推断逻辑 +- profile resolver 对旧 capability tags 的兼容升级 ### Adding new professional edit models -1. 在 `utils/model_info.py` 里补充或修正角色与 capability 推断规则 -2. 确认 `LLS Simple Checkpoint Loader` 会把正确的 `_lls_*` capability tags 写到 `model` / `clip` / `vae` -3. 优先用 `backend_mode=auto` 验证 capability metadata 是否足以自动识别 -4. 调试阶段可临时改用 `backend_mode=sdxl` 或 `backend_mode=flux`，快速确认显式路由是否匹配预期 +1. 在 `model_profiles/` 里补充或修正 profile rule +2. 确认 `LLS Simple Checkpoint Loader` 会把正确的 `_lls_*` profile tags 写到 `model` / `clip` / `vae` +3. 优先用 `backend_mode=auto` 验证 resolved profile 是否足以自动识别 +4. 调试阶段可通过 `model_info` 显式覆盖 `profile_id` / `backend_type` / `sampler_strategy` ### Notes diff --git a/docs/superpowers/plans/2026-05-21-pro-model-profile-routing.md b/docs/superpowers/plans/2026-05-21-pro-model-profile-routing.md new file mode 100644 index 0000000..a473ca4 --- /dev/null +++ b/docs/superpowers/plans/2026-05-21-pro-model-profile-routing.md @@ -0,0 +1,1255 @@ +# LLS Pro Model Profile Routing Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Refactor the current Pro image edit chain so loader, backend routing, and sampler routing all use a single explicit model-profile contract while keeping the existing Pro node names and workflow wiring stable. + +**Architecture:** Add a new internal `model_profiles/` package that resolves a normalized `ModelProfile` from loader inputs, overrides, and legacy tags. Update `utils/model_info.py`, `model_loader/nodes.py`, `pro_edit/backends/registry.py`, `pro_edit/pro_edit_prepare.py`, and `pro_edit/pro_edit_bridge.py` so `backend_type` and `sampler_strategy` come from that profile instead of repeated filename heuristics. + +**Tech Stack:** Python, unittest, existing ComfyUI node classes, JSON-ish metadata parsing, fake model/image/mask helpers, `python3 -m compileall` + +--- + +## File Structure + +- Create: `model_profiles/__init__.py` + - Export the profile resolution helpers. +- Create: `model_profiles/base.py` + - Define the normalized `ModelProfile` record and helper conversions. +- Create: `model_profiles/rules.py` + - Define built-in profile rules for `sd15_base`, `sdxl_base`, `sdxl_inpaint`, `sdxl_edit`, `flux_base`, `flux_edit`. +- Create: `model_profiles/registry.py` + - Resolve explicit overrides, legacy tags, and fallback rule matches into a final profile. +- Create: `tests/test_model_profiles.py` + - Verify profile resolution from family, name patterns, overrides, and legacy tags. +- Create: `tests/test_model_profile_loader.py` + - Verify loader writes `_lls_profile_*` metadata and preserves compatibility tags. +- Create: `tests/test_pro_edit_profile_prepare.py` + - Verify Pro prepare fails cleanly for `backend_type = none` and writes profile metadata into `edit_info`. +- Modify: `tests/test_pro_edit_helpers.py` + - Extend `FakeModel` to carry explicit profile fields. +- Modify: `tests/test_pro_edit_capabilities.py` + - Add profile-aware expectations to existing model-info compatibility tests. +- Modify: `tests/test_pro_edit_registry.py` + - Route by `backend_type` and reject unsupported base profiles. +- Modify: `tests/test_pro_edit_prepare_sdxl.py` + - Assert profile metadata appears in `edit_info`. +- Modify: `tests/test_pro_edit_prepare_flux.py` + - Assert profile metadata appears in `edit_info`. +- Modify: `tests/test_pro_edit_bridge.py` + - Assert `sample_info` uses `profile_id` and `sampler_strategy`. +- Modify: `tests/test_pro_edit_docs.py` + - Assert README documents the profile-driven routing model. +- Modify: `utils/model_info.py` + - Wrap legacy helpers around the new profile registry. +- Modify: `model_loader/nodes.py` + - Resolve and tag complete profiles during checkpoint loading. +- Modify: `pro_edit/backends/base.py` + - Carry profile-aware routing records and backend validation helpers. +- Modify: `pro_edit/backends/registry.py` + - Route by `backend_type` using resolved profiles. +- Modify: `pro_edit/backends/sdxl.py` + - Validate and report `sdxl_native`. +- Modify: `pro_edit/backends/flux.py` + - Validate and report `flux_edit`. +- Modify: `pro_edit/pro_edit_prepare.py` + - Resolve the effective profile and reject unsupported base profiles. +- Modify: `pro_edit/pro_edit_bridge.py` + - Resolve the effective profile and dispatch by `sampler_strategy`. +- Modify: `pro_edit/pro_edit_finish.py` + - Preserve `profile_id`, `backend_type`, and `sampler_strategy` from `edit_info`. +- Modify: `pro_edit/pro_edit_utils.py` + - Normalize and serialize the new profile fields inside `edit_info`. +- Modify: `README.md` + - Document the profile-driven routing model and override strategy. + +### Task 1: Add The Model Profile Core + +**Files:** +- Create: `model_profiles/__init__.py` +- Create: `model_profiles/base.py` +- Create: `model_profiles/rules.py` +- Create: `model_profiles/registry.py` +- Create: `tests/test_model_profiles.py` +- Modify: `tests/test_pro_edit_helpers.py` + +- [ ] **Step 1: Write the failing tests** + +```python +# tests/test_model_profiles.py +import unittest + +try: + from .test_pro_edit_helpers import FakeModel, import_plugin_submodule, load_plugin_package +except ImportError: + from test_pro_edit_helpers import FakeModel, import_plugin_submodule, load_plugin_package + + +class TestModelProfiles(unittest.TestCase): + def setUp(self): + self.plugin = load_plugin_package() + self.registry = import_plugin_submodule(self.plugin, "model_profiles.registry") + + def test_resolves_sdxl_base_profile_from_family(self): + profile = self.registry.resolve_model_profile( + model=None, + model_info={"checkpoint_name": "plain-sdxl-base.safetensors", "family": "SDXL"}, + ) + + self.assertEqual(profile["profile_id"], "sdxl_base") + self.assertEqual(profile["backend_type"], "none") + self.assertEqual(profile["sampler_strategy"], "standard_k") + + def test_resolves_flux_kontext_profile_from_name(self): + profile = self.registry.resolve_model_profile( + model=None, + model_info={"checkpoint_name": "demo-flux-kontext-dev.safetensors", "family": "FLUX_DEV"}, + ) + + self.assertEqual(profile["profile_id"], "flux_edit") + self.assertEqual(profile["role"], "edit") + self.assertEqual(profile["backend_type"], "flux_edit") + self.assertEqual(profile["sampler_strategy"], "flux_guided") + + def test_model_info_override_has_highest_priority(self): + profile = self.registry.resolve_model_profile( + model=None, + model_info={ + "checkpoint_name": "plain-sdxl-base.safetensors", + "family": "SDXL", + "profile_id": "sdxl_edit", + "role": "edit", + "backend_type": "sdxl_native", + "sampler_strategy": "standard_k", + "supports_inpaint_native": True, + "supports_image_edit_native": True, + "preferred_edit_backend": "sdxl", + }, + ) + + self.assertEqual(profile["profile_id"], "sdxl_edit") + self.assertEqual(profile["backend_type"], "sdxl_native") + self.assertTrue(profile["supports_image_edit_native"]) + + def test_legacy_tags_upgrade_to_matching_profile(self): + model = FakeModel( + family="FLUX_DEV", + model_role="edit", + supports_inpaint_native=False, + supports_image_edit_native=True, + preferred_edit_backend="flux", + profile_id="", + backend_type="", + sampler_strategy="", + loader_strategy="", + model_name="legacy-flux-edit.safetensors", + ) + + profile = self.registry.resolve_model_profile(model=model, model_info=None) + + self.assertEqual(profile["profile_id"], "flux_edit") + self.assertEqual(profile["backend_type"], "flux_edit") + self.assertEqual(profile["sampler_strategy"], "flux_guided") + + +if __name__ == "__main__": + unittest.main() +``` + +```python +# tests/test_pro_edit_helpers.py +class FakeModel: + def __init__( + self, + family="SDXL", + model_role="base", + supports_inpaint_native=False, + supports_image_edit_native=False, + preferred_edit_backend=None, + model_name="demo-model.safetensors", + profile_id="", + backend_type="", + sampler_strategy="", + loader_strategy="", + ): + self._lls_family = family + self._lls_model_role = model_role + self._lls_supports_inpaint_native = supports_inpaint_native + self._lls_supports_image_edit_native = supports_image_edit_native + self._lls_preferred_edit_backend = preferred_edit_backend + self._lls_model_name = model_name + self._lls_profile_id = profile_id + self._lls_backend_type = backend_type + self._lls_sampler_strategy = sampler_strategy + self._lls_loader_strategy = loader_strategy +``` + +- [ ] **Step 2: Run the tests to verify they fail** + +Run: `env PYTHONDONTWRITEBYTECODE=1 python3 -m unittest discover -s tests -p 'test_model_profiles.py' -v` + +Expected: FAIL because `model_profiles` does not exist yet and `FakeModel` does not carry profile fields. + +- [ ] **Step 3: Implement the model profile core** + +```python +# model_profiles/base.py +from __future__ import annotations + +from dataclasses import asdict, dataclass +from typing import Any + + +@dataclass(frozen=True) +class ModelProfile: + profile_id: str + family: str + role: str + backend_type: str + sampler_strategy: str + loader_strategy: str + supports_inpaint_native: bool + supports_image_edit_native: bool + preferred_edit_backend: str | None + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +def normalize_profile_dict(value: dict[str, Any]) -> dict[str, Any]: + return { + "profile_id": str(value.get("profile_id") or ""), + "family": str(value.get("family") or value.get("model_family") or "SD1.5"), + "role": str(value.get("role") or value.get("model_role") or "base"), + "backend_type": str(value.get("backend_type") or "none"), + "sampler_strategy": str(value.get("sampler_strategy") or "standard_k"), + "loader_strategy": str(value.get("loader_strategy") or ""), + "supports_inpaint_native": bool(value.get("supports_inpaint_native", False)), + "supports_image_edit_native": bool(value.get("supports_image_edit_native", False)), + "preferred_edit_backend": value.get("preferred_edit_backend"), + } + + +def build_model_profile(**kwargs: Any) -> ModelProfile: + normalized = normalize_profile_dict(kwargs) + return ModelProfile(**normalized) +``` + +```python +# model_profiles/rules.py +from __future__ import annotations + +from ..utils.model_info import canonicalize_family, infer_family_from_name, infer_model_role_from_name + + +_ROLE_TO_BACKEND = { + ("SDXL", "inpaint"): ("sdxl_inpaint", "sdxl_native", "standard_k", "sdxl_checkpoint", True, False, "sdxl"), + ("SDXL", "edit"): ("sdxl_edit", "sdxl_native", "standard_k", "sdxl_checkpoint", True, True, "sdxl"), + ("SDXL", "fill"): ("sdxl_edit", "sdxl_native", "standard_k", "sdxl_checkpoint", True, True, "sdxl"), + ("FLUX_DEV", "edit"): ("flux_edit", "flux_edit", "flux_guided", "flux_split_or_bundle", False, True, "flux"), + ("FLUX_DEV", "fill"): ("flux_edit", "flux_edit", "flux_guided", "flux_split_or_bundle", False, True, "flux"), + ("FLUX_DEV", "inpaint"): ("flux_edit", "flux_edit", "flux_guided", "flux_split_or_bundle", True, True, "flux"), + ("FLUX_SCHNELL", "edit"): ("flux_edit", "flux_edit", "flux_guided", "flux_split_or_bundle", False, True, "flux"), + ("FLUX_SCHNELL", "fill"): ("flux_edit", "flux_edit", "flux_guided", "flux_split_or_bundle", False, True, "flux"), + ("FLUX_SCHNELL", "inpaint"): ("flux_edit", "flux_edit", "flux_guided", "flux_split_or_bundle", True, True, "flux"), +} + + +def _base_profile_for_family(family: str) -> dict[str, object]: + if family == "SDXL": + return { + "profile_id": "sdxl_base", + "family": family, + "role": "base", + "backend_type": "none", + "sampler_strategy": "standard_k", + "loader_strategy": "sdxl_checkpoint", + "supports_inpaint_native": False, + "supports_image_edit_native": False, + "preferred_edit_backend": None, + } + if family in {"FLUX_DEV", "FLUX_SCHNELL"}: + return { + "profile_id": "flux_base", + "family": family, + "role": "base", + "backend_type": "none", + "sampler_strategy": "flux_guided", + "loader_strategy": "flux_split_or_bundle", + "supports_inpaint_native": False, + "supports_image_edit_native": False, + "preferred_edit_backend": None, + } + return { + "profile_id": "sd15_base", + "family": "SD1.5", + "role": "base", + "backend_type": "none", + "sampler_strategy": "standard_k", + "loader_strategy": "sd15_checkpoint", + "supports_inpaint_native": False, + "supports_image_edit_native": False, + "preferred_edit_backend": None, + } + + +def match_builtin_profile(model_name: str | None, family: str | None, role: str | None = None) -> dict[str, object]: + resolved_family = canonicalize_family(family or infer_family_from_name(model_name, "SD1.5")) + resolved_role = str(role or infer_model_role_from_name(model_name, resolved_family)) + key = (resolved_family, resolved_role) + matched = _ROLE_TO_BACKEND.get(key) + if matched is None: + return _base_profile_for_family(resolved_family) + profile_id, backend_type, sampler_strategy, loader_strategy, supports_inpaint_native, supports_image_edit_native, preferred_backend = matched + return { + "profile_id": profile_id, + "family": resolved_family, + "role": resolved_role, + "backend_type": backend_type, + "sampler_strategy": sampler_strategy, + "loader_strategy": loader_strategy, + "supports_inpaint_native": supports_inpaint_native, + "supports_image_edit_native": supports_image_edit_native, + "preferred_edit_backend": preferred_backend, + } +``` + +```python +# model_profiles/registry.py +from __future__ import annotations + +from .base import build_model_profile +from .rules import match_builtin_profile +from ..utils.model_info import get_lls_attr, parse_jsonish_info + + +def _extract_profile_overrides(model=None, model_info=None, extra_info=None) -> dict[str, object]: + raw = parse_jsonish_info(model_info) + extra = parse_jsonish_info(extra_info) + + def _pick(*values): + for value in values: + if value not in (None, ""): + return value + return None + + return { + "profile_id": _pick(extra.get("profile_id"), raw.get("profile_id"), get_lls_attr(model, "profile_id", None)), + "family": _pick(extra.get("family"), extra.get("model_family"), raw.get("family"), raw.get("model_family"), get_lls_attr(model, "family", None)), + "role": _pick(extra.get("role"), extra.get("model_role"), raw.get("role"), raw.get("model_role"), get_lls_attr(model, "model_role", None)), + "backend_type": _pick(extra.get("backend_type"), raw.get("backend_type"), get_lls_attr(model, "backend_type", None)), + "sampler_strategy": _pick(extra.get("sampler_strategy"), raw.get("sampler_strategy"), get_lls_attr(model, "sampler_strategy", None)), + "loader_strategy": _pick(extra.get("loader_strategy"), raw.get("loader_strategy"), get_lls_attr(model, "loader_strategy", None)), + "supports_inpaint_native": ( + extra["supports_inpaint_native"] if "supports_inpaint_native" in extra + else raw["supports_inpaint_native"] if "supports_inpaint_native" in raw + else get_lls_attr(model, "supports_inpaint_native", None) + ), + "supports_image_edit_native": ( + extra["supports_image_edit_native"] if "supports_image_edit_native" in extra + else raw["supports_image_edit_native"] if "supports_image_edit_native" in raw + else get_lls_attr(model, "supports_image_edit_native", None) + ), + "preferred_edit_backend": _pick(extra.get("preferred_edit_backend"), raw.get("preferred_edit_backend"), get_lls_attr(model, "preferred_edit_backend", None)), + "model_name": _pick(extra.get("checkpoint_name"), extra.get("model_name"), raw.get("checkpoint_name"), raw.get("model_name"), get_lls_attr(model, "model_name", None)), + } + + +def resolve_model_profile(model=None, model_info=None, extra_info=None, checkpoint_name: str | None = None, family: str | None = None) -> dict[str, object]: + overrides = _extract_profile_overrides(model=model, model_info=model_info, extra_info=extra_info) + model_name = str(checkpoint_name or overrides.get("model_name") or "") + base = match_builtin_profile(model_name, family or overrides.get("family"), overrides.get("role")) + merged = dict(base) + for key in ( + "profile_id", + "family", + "role", + "backend_type", + "sampler_strategy", + "loader_strategy", + "preferred_edit_backend", + ): + if overrides.get(key): + merged[key] = overrides[key] + for key in ("supports_inpaint_native", "supports_image_edit_native"): + if overrides.get(key) is not None: + merged[key] = bool(overrides[key]) + return build_model_profile(**merged).to_dict() +``` + +```python +# model_profiles/__init__.py +from .base import ModelProfile, build_model_profile, normalize_profile_dict +from .registry import resolve_model_profile + +__all__ = [ + "ModelProfile", + "build_model_profile", + "normalize_profile_dict", + "resolve_model_profile", +] +``` + +- [ ] **Step 4: Run the tests to verify they pass** + +Run: `env PYTHONDONTWRITEBYTECODE=1 python3 -m unittest discover -s tests -p 'test_model_profiles.py' -v` + +Expected: PASS with four tests. + +- [ ] **Step 5: Commit** + +```bash +git add model_profiles tests/test_model_profiles.py tests/test_pro_edit_helpers.py +git commit -m "feat: add model profile core" +``` + +### Task 2: Make `utils/model_info.py` Profile-Aware Without Breaking Compatibility + +**Files:** +- Modify: `utils/model_info.py` +- Modify: `tests/test_pro_edit_capabilities.py` + +- [ ] **Step 1: Write the failing tests** + +```python +# tests/test_pro_edit_capabilities.py + def test_parse_model_info_exposes_profile_fields(self): + plugin = load_plugin_package() + model_info = import_plugin_submodule(plugin, "utils.model_info") + + info = model_info.parse_model_info( + { + "checkpoint_name": "demo-flux-kontext-dev.safetensors", + "family": "FLUX_DEV", + } + ) + + self.assertEqual(info["profile_id"], "flux_edit") + self.assertEqual(info["backend_type"], "flux_edit") + self.assertEqual(info["sampler_strategy"], "flux_guided") + + def test_resolve_edit_capabilities_reads_profile_tags(self): + plugin = load_plugin_package() + model_info = import_plugin_submodule(plugin, "utils.model_info") + model = FakeModel( + family="SDXL", + model_role="edit", + supports_inpaint_native=True, + supports_image_edit_native=True, + preferred_edit_backend="sdxl", + profile_id="sdxl_edit", + backend_type="sdxl_native", + sampler_strategy="standard_k", + loader_strategy="sdxl_checkpoint", + model_name="demo-sdxl-edit.safetensors", + ) + + capabilities = model_info.resolve_edit_capabilities(model=model, model_info=None) + + self.assertEqual(capabilities["profile_id"], "sdxl_edit") + self.assertEqual(capabilities["backend_type"], "sdxl_native") + self.assertEqual(capabilities["sampler_strategy"], "standard_k") +``` + +- [ ] **Step 2: Run the tests to verify they fail** + +Run: `env PYTHONDONTWRITEBYTECODE=1 python3 -m unittest discover -s tests -p 'test_pro_edit_capabilities.py' -v` + +Expected: FAIL because `parse_model_info()` and `resolve_edit_capabilities()` do not yet expose profile fields. + +- [ ] **Step 3: Implement the profile-aware compatibility layer** + +```python +# utils/model_info.py +def parse_model_info(model_info: dict[str, Any] | str | None) -> dict[str, Any]: + from ..model_profiles.registry import resolve_model_profile + + raw = parse_jsonish_info(model_info) + raw_model_name = _get_model_name_alias_value(raw) + family = canonicalize_family( + raw.get("family") + or raw.get("model_family") + or infer_family_from_name(raw_model_name, "SD1.5") + ) + defaults = get_family_defaults(family) + profile = resolve_model_profile(model=None, model_info=raw, checkpoint_name=raw_model_name, family=family) + + info: dict[str, Any] = {**defaults} + info.update(raw) + info["family"] = profile["family"] + info["model_family"] = profile["family"] + info["model_role"] = profile["role"] + info["profile_id"] = profile["profile_id"] + info["backend_type"] = profile["backend_type"] + info["sampler_strategy"] = profile["sampler_strategy"] + info["loader_strategy"] = profile["loader_strategy"] + info["supports_inpaint_native"] = bool(profile["supports_inpaint_native"]) + info["supports_image_edit_native"] = bool(profile["supports_image_edit_native"]) + info["preferred_edit_backend"] = profile["preferred_edit_backend"] + info.setdefault("checkpoint_name", raw_model_name or "") + info.setdefault("model_name", info["checkpoint_name"]) + return info + + +def resolve_edit_capabilities(model=None, model_info: dict[str, Any] | str | None = None) -> dict[str, Any]: + from ..model_profiles.registry import resolve_model_profile + + info = parse_model_info(model_info) + profile = resolve_model_profile( + model=model, + model_info=model_info, + checkpoint_name=str( + info.get("checkpoint_name") + or info.get("model_name") + or get_lls_attr(model, "model_name", "") + or "" + ), + family=info.get("family") or info.get("model_family"), + ) + return { + "model_family": profile["family"], + "model_name": str( + get_lls_attr(model, "model_name", None) + or info.get("checkpoint_name") + or info.get("model_name") + or "" + ), + "model_role": profile["role"], + "profile_id": profile["profile_id"], + "backend_type": profile["backend_type"], + "sampler_strategy": profile["sampler_strategy"], + "loader_strategy": profile["loader_strategy"], + "supports_inpaint_native": bool(profile["supports_inpaint_native"]), + "supports_image_edit_native": bool(profile["supports_image_edit_native"]), + "preferred_edit_backend": profile["preferred_edit_backend"], + } +``` + +- [ ] **Step 4: Run the tests to verify they pass** + +Run: `env PYTHONDONTWRITEBYTECODE=1 python3 -m unittest discover -s tests -p 'test_pro_edit_capabilities.py' -v` + +Expected: PASS with the new profile-aware assertions and the existing capability regression coverage. + +- [ ] **Step 5: Commit** + +```bash +git add utils/model_info.py tests/test_pro_edit_capabilities.py +git commit -m "feat: expose model profiles through model info helpers" +``` + +### Task 3: Stamp Resolved Profiles In The Loader + +**Files:** +- Create: `tests/test_model_profile_loader.py` +- Modify: `model_loader/nodes.py` + +- [ ] **Step 1: Write the failing tests** + +```python +# tests/test_model_profile_loader.py +import unittest +from types import SimpleNamespace +from unittest import mock + +try: + from .test_pro_edit_helpers import import_plugin_submodule, load_plugin_package +except ImportError: + from test_pro_edit_helpers import import_plugin_submodule, load_plugin_package + + +class TestModelProfileLoader(unittest.TestCase): + def test_loader_writes_profile_fields_to_outputs(self): + plugin = load_plugin_package() + loader_module = import_plugin_submodule(plugin, "model_loader.nodes") + loader = loader_module.LLSSimpleCheckpointLoader() + model = SimpleNamespace() + clip = SimpleNamespace() + vae = SimpleNamespace() + + with mock.patch.object(loader_module, "folder_paths", object()), \ + mock.patch.object(loader_module, "comfy_sd", object()), \ + mock.patch.object(loader_module, "_resolve_model_path", return_value=("diffusion_models", "/fake/demo-flux-kontext-dev.safetensors")), \ + mock.patch.object(loader_module, "_load_model", return_value=(model, clip, vae)), \ + mock.patch.object(loader_module, "_resolve_text_encoder", return_value=(clip, "embedded", None, None)), \ + mock.patch.object(loader_module, "_resolve_vae", return_value=(vae, "embedded", None)): + loaded_model, loaded_clip, loaded_vae, text_encoder = loader.load_checkpoint( + ckpt_name="demo-flux-kontext-dev.safetensors", + model_family="FLUX_DEV", + load_mode="simple", + vae_source="auto", + text_encoder_source="auto", + external_vae_name=loader_module.AUTO_PLACEHOLDER, + external_text_encoder_1=loader_module.AUTO_PLACEHOLDER, + external_text_encoder_2=loader_module.AUTO_PLACEHOLDER, + ) + + for obj in (loaded_model, loaded_clip, loaded_vae, text_encoder): + self.assertEqual(obj._lls_profile_id, "flux_edit") + self.assertEqual(obj._lls_backend_type, "flux_edit") + self.assertEqual(obj._lls_sampler_strategy, "flux_guided") + self.assertEqual(obj._lls_loader_strategy, "flux_split_or_bundle") + self.assertEqual(obj._lls_model_role, "edit") +``` + +- [ ] **Step 2: Run the tests to verify they fail** + +Run: `env PYTHONDONTWRITEBYTECODE=1 python3 -m unittest discover -s tests -p 'test_model_profile_loader.py' -v` + +Expected: FAIL because loader output tags do not yet include `_lls_profile_*` fields. + +- [ ] **Step 3: Update the loader to stamp full profiles** + +```python +# model_loader/nodes.py +from ..model_profiles.registry import resolve_model_profile + + +def _build_capability_tags(model_name: str, family: str) -> dict[str, object]: + profile = resolve_model_profile( + model=None, + model_info={"checkpoint_name": model_name, "family": family}, + checkpoint_name=model_name, + family=family, + ) + return { + "profile_id": profile["profile_id"], + "backend_type": profile["backend_type"], + "sampler_strategy": profile["sampler_strategy"], + "loader_strategy": profile["loader_strategy"], + "model_role": profile["role"], + "supports_inpaint_native": profile["supports_inpaint_native"], + "supports_image_edit_native": profile["supports_image_edit_native"], + "preferred_edit_backend": profile["preferred_edit_backend"], + } +``` + +```python +# model_loader/nodes.py + profile_tags = _build_capability_tags(ckpt_name, family) + + tag_lls_object( + model, + family=family, + model_name=ckpt_name, + checkpoint_name=ckpt_name, + load_mode=load_mode, + **profile_tags, + ) + tag_lls_object( + text_encoder, + family=family, + model_name=ckpt_name, + checkpoint_name=ckpt_name, + text_encoder_type=defaults["text_encoder_type"], + text_encoder_source=resolved_text_encoder_source, + text_encoder_name=resolved_text_encoder_name, + text_encoder_name_1=resolved_text_encoder_name_1, + text_encoder_name_2=resolved_text_encoder_name_2, + **profile_tags, + ) + tag_lls_object( + vae, + family=family, + model_name=ckpt_name, + checkpoint_name=ckpt_name, + vae_name=resolved_vae_label, + vae_source=resolved_vae_source, + **profile_tags, + ) +``` + +- [ ] **Step 4: Run the tests to verify they pass** + +Run: `env PYTHONDONTWRITEBYTECODE=1 python3 -m unittest discover -s tests -p 'test_model_profile_loader.py' -v` + +Expected: PASS with one loader-tagging test. + +- [ ] **Step 5: Commit** + +```bash +git add model_loader/nodes.py tests/test_model_profile_loader.py +git commit -m "feat: stamp resolved model profiles in loader" +``` + +### Task 4: Refactor Pro Backend Routing To Use `backend_type` + +**Files:** +- Modify: `pro_edit/backends/base.py` +- Modify: `pro_edit/backends/registry.py` +- Modify: `pro_edit/backends/sdxl.py` +- Modify: `pro_edit/backends/flux.py` +- Modify: `tests/test_pro_edit_registry.py` + +- [ ] **Step 1: Write the failing tests** + +```python +# tests/test_pro_edit_registry.py + def test_auto_routes_sdxl_profile_by_backend_type(self): + model = FakeModel( + family="SDXL", + model_role="edit", + supports_inpaint_native=True, + supports_image_edit_native=True, + preferred_edit_backend="sdxl", + profile_id="sdxl_edit", + backend_type="sdxl_native", + sampler_strategy="standard_k", + loader_strategy="sdxl_checkpoint", + ) + + backend, routing = self.registry.resolve_backend("auto", model=model) + + self.assertEqual(backend.backend_name, "sdxl") + self.assertEqual(routing.profile["profile_id"], "sdxl_edit") + self.assertEqual(routing.profile["backend_type"], "sdxl_native") + + def test_auto_rejects_base_profile_for_pro_chain(self): + model = FakeModel( + family="SDXL", + model_role="base", + supports_inpaint_native=False, + supports_image_edit_native=False, + preferred_edit_backend=None, + profile_id="sdxl_base", + backend_type="none", + sampler_strategy="standard_k", + loader_strategy="sdxl_checkpoint", + ) + + with self.assertRaisesRegex(RuntimeError, "Pro image edit is not available for profile 'sdxl_base'"): + self.registry.resolve_backend("auto", model=model) +``` + +- [ ] **Step 2: Run the tests to verify they fail** + +Run: `env PYTHONDONTWRITEBYTECODE=1 python3 -m unittest discover -s tests -p 'test_pro_edit_registry.py' -v` + +Expected: FAIL because routing records do not yet carry `profile`, and `backend_type = none` is not treated as a first-class unsupported profile state. + +- [ ] **Step 3: Route by `backend_type`** + +```python +# pro_edit/backends/base.py +from dataclasses import dataclass +from typing import Any + + +@dataclass(frozen=True) +class RoutingResult: + backend_name: str + routing_reason: str + capabilities: dict[str, Any] + profile: dict[str, Any] + + +def validate_manual_backend(backend_name: str, backend, profile: dict[str, Any]) -> None: + if backend.supports(profile): + return + raise RuntimeError( + f"[LLS] Pro edit backend '{backend_name}' is incompatible with profile " + f"'{profile.get('profile_id')}' (family '{profile.get('family')}', role '{profile.get('role')}')." + ) +``` + +```python +# pro_edit/backends/registry.py +from ...model_profiles.registry import resolve_model_profile +from ...utils.model_info import parse_jsonish_info + + +def resolve_backend(backend_mode: str, *, model=None, model_info=None, edit_info=None): + _ensure_builtin_backends() + mode = str(backend_mode or "auto").strip().lower() or "auto" + merged_extra_info = dict(parse_jsonish_info(model_info)) + merged_extra_info.update(parse_jsonish_info(edit_info)) + profile = resolve_model_profile(model=model, model_info=model_info, extra_info=merged_extra_info) + capabilities = { + "model_family": profile["family"], + "model_role": profile["role"], + "profile_id": profile["profile_id"], + "backend_type": profile["backend_type"], + "sampler_strategy": profile["sampler_strategy"], + "loader_strategy": profile["loader_strategy"], + "supports_inpaint_native": profile["supports_inpaint_native"], + "supports_image_edit_native": profile["supports_image_edit_native"], + "preferred_edit_backend": profile["preferred_edit_backend"], + } + + backend_type = str(profile.get("backend_type") or "none") + if backend_type == "none": + raise RuntimeError( + f"[LLS] Pro image edit is not available for profile '{profile['profile_id']}' " + f"(family '{profile['family']}', role '{profile['role']}'). " + "Use a supported native edit/inpaint profile or provide an explicit override." + ) + + backend_name = { + "sdxl_native": "sdxl", + "flux_edit": "flux", + }.get(backend_type) + if backend_name is None: + raise RuntimeError(f"[LLS] Unsupported pro edit backend_type '{backend_type}'.") + if mode != "auto": + backend = get_backend(mode) + validate_manual_backend(mode, backend, profile) + return backend, RoutingResult(backend.backend_name, "manual.override", capabilities, profile) + + backend = get_backend(backend_name) + validate_manual_backend(backend_name, backend, profile) + return backend, RoutingResult(backend.backend_name, "profile.backend_type", capabilities, profile) +``` + +```python +# pro_edit/backends/sdxl.py + def supports(self, profile): + return str(profile.get("backend_type") or "") == "sdxl_native" +``` + +```python +# pro_edit/backends/flux.py + def supports(self, profile): + return str(profile.get("backend_type") or "") == "flux_edit" +``` + +- [ ] **Step 4: Run the tests to verify they pass** + +Run: `env PYTHONDONTWRITEBYTECODE=1 python3 -m unittest discover -s tests -p 'test_pro_edit_registry.py' -v` + +Expected: PASS with the updated profile-based routing assertions plus the legacy regression tests adjusted to the new `routing_reason`. + +- [ ] **Step 5: Commit** + +```bash +git add pro_edit/backends tests/test_pro_edit_registry.py +git commit -m "feat: route pro backends by model profile" +``` + +### Task 5: Make Pro Prepare Profile-Driven + +**Files:** +- Create: `tests/test_pro_edit_profile_prepare.py` +- Modify: `pro_edit/pro_edit_prepare.py` +- Modify: `pro_edit/pro_edit_utils.py` +- Modify: `tests/test_pro_edit_prepare_sdxl.py` +- Modify: `tests/test_pro_edit_prepare_flux.py` + +- [ ] **Step 1: Write the failing tests** + +```python +# tests/test_pro_edit_profile_prepare.py +import unittest + +try: + from .test_pro_edit_helpers import FakeMask, FakeModel, FakeTensor, FakeVAE, load_plugin_package, make_conditioning +except ImportError: + from test_pro_edit_helpers import FakeMask, FakeModel, FakeTensor, FakeVAE, load_plugin_package, make_conditioning + + +class TestProEditProfilePrepare(unittest.TestCase): + def test_prepare_rejects_base_profile(self): + plugin = load_plugin_package() + node = plugin.NODE_CLASS_MAPPINGS["LLSProImageEditPrepare"]() + image = FakeTensor((1, 64, 64, 3), label="image") + mask = FakeMask((1, 64, 64), mask_bbox=(8, 8, 32, 32), mask_area_ratio=0.2) + vae = FakeVAE(latent_channels=4, downscale_ratio=8) + model = FakeModel( + family="SDXL", + model_role="base", + supports_inpaint_native=False, + supports_image_edit_native=False, + preferred_edit_backend=None, + profile_id="sdxl_base", + backend_type="none", + sampler_strategy="standard_k", + loader_strategy="sdxl_checkpoint", + ) + + with self.assertRaisesRegex(RuntimeError, "Pro image edit is not available for profile 'sdxl_base'"): + node.prepare( + image=image, + mask=mask, + vae=vae, + positive=make_conditioning("positive"), + negative=make_conditioning("negative"), + backend_mode="auto", + edit_scope="region", + mask_grow=0, + mask_blur=0.0, + mask_threshold=0.5, + invert_mask=False, + crop_context=32, + crop_context_factor=1.0, + min_size=128, + max_size=512, + resize_mode="fit", + expand_left=0, + expand_right=0, + expand_top=0, + expand_bottom=0, + canvas_fill="edge", + auto_recommend="enabled", + model=model, + model_info=None, + ) +``` + +```python +# tests/test_pro_edit_prepare_sdxl.py + self.assertEqual(edit_info["profile_id"], "sdxl_inpaint") + self.assertEqual(edit_info["backend_type"], "sdxl_native") + self.assertEqual(edit_info["sampler_strategy"], "standard_k") +``` + +```python +# tests/test_pro_edit_prepare_flux.py + self.assertEqual(edit_info["profile_id"], "flux_edit") + self.assertEqual(edit_info["backend_type"], "flux_edit") + self.assertEqual(edit_info["sampler_strategy"], "flux_guided") +``` + +- [ ] **Step 2: Run the tests to verify they fail** + +Run: `env PYTHONDONTWRITEBYTECODE=1 python3 -m unittest discover -s tests -p 'test_pro_edit_profile_prepare.py' -v` + +Expected: FAIL because prepare/edit_info do not yet preserve the resolved profile metadata. + +- [ ] **Step 3: Resolve and serialize profiles in prepare** + +```python +# pro_edit/pro_edit_utils.py +def build_edit_info(workspace: dict[str, Any], routing, *, backend_hints: dict[str, Any] | None = None) -> dict[str, Any]: + profile = routing.profile + info = { + "backend_name": routing.backend_name, + "routing_reason": routing.routing_reason, + "model_family": profile["family"], + "model_role": profile["role"], + "profile_id": profile["profile_id"], + "backend_type": profile["backend_type"], + "sampler_strategy": profile["sampler_strategy"], + "loader_strategy": profile["loader_strategy"], + "supports_inpaint_native": profile["supports_inpaint_native"], + "supports_image_edit_native": profile["supports_image_edit_native"], + "preferred_edit_backend": profile["preferred_edit_backend"], + "edit_scope": workspace["edit_scope"], + "original_size": list(workspace["original_size"]), + "work_size": list(workspace["work_size"]), + "crop_box": workspace["crop_box"], + "crop_scale": workspace["crop_scale"], + "canvas_expand": workspace["canvas_expand"], + "original_box_in_canvas": workspace["original_box_in_canvas"], + "recommended_denoise": float(workspace["recommended_denoise"]), + "edit_payload_version": "1.0", + } + if backend_hints: + info.update(backend_hints) + return info +``` + +```python +# pro_edit/pro_edit_prepare.py + backend, routing = resolve_backend( + backend_mode, + model=model, + model_info=model_info, + ) + prepared = backend.prepare( + model=model, + vae=vae, + work_image=workspace["work_image"], + work_mask=workspace["work_mask"], + positive=positive, + negative=negative, + workspace=workspace, + routing=routing, + ) +``` + +- [ ] **Step 4: Run the tests to verify they pass** + +Run: `env PYTHONDONTWRITEBYTECODE=1 python3 -m unittest discover -s tests -p 'test_pro_edit_profile_prepare.py' -v` + +Expected: PASS with the unsupported-base-profile failure. + +Run: `env PYTHONDONTWRITEBYTECODE=1 python3 -m unittest discover -s tests -p 'test_pro_edit_prepare_sdxl.py' -v` + +Expected: PASS with profile metadata assertions. + +Run: `env PYTHONDONTWRITEBYTECODE=1 python3 -m unittest discover -s tests -p 'test_pro_edit_prepare_flux.py' -v` + +Expected: PASS with profile metadata assertions. + +- [ ] **Step 5: Commit** + +```bash +git add pro_edit/pro_edit_prepare.py pro_edit/pro_edit_utils.py tests/test_pro_edit_profile_prepare.py tests/test_pro_edit_prepare_sdxl.py tests/test_pro_edit_prepare_flux.py +git commit -m "feat: make pro prepare profile driven" +``` + +### Task 6: Route The Bridge By `sampler_strategy` + +**Files:** +- Modify: `pro_edit/pro_edit_bridge.py` +- Modify: `pro_edit/pro_edit_finish.py` +- Modify: `tests/test_pro_edit_bridge.py` + +- [ ] **Step 1: Write the failing tests** + +```python +# tests/test_pro_edit_bridge.py + self.assertEqual(sample_info["profile_id"], "sdxl_inpaint") + self.assertEqual(sample_info["sampler_strategy"], "standard_k") +``` + +```python +# tests/test_pro_edit_bridge.py + self.assertEqual(sample_info["profile_id"], "flux_edit") + self.assertEqual(sample_info["sampler_strategy"], "flux_guided") +``` + +```python +# tests/test_pro_edit_bridge.py + def test_bridge_rejects_unknown_sampler_strategy(self): + model = FakeModel( + family="FLUX_DEV", + model_role="edit", + supports_inpaint_native=False, + supports_image_edit_native=True, + preferred_edit_backend="flux", + profile_id="flux_edit", + backend_type="flux_edit", + sampler_strategy="mystery_strategy", + loader_strategy="flux_split_or_bundle", + ) + latent = {"samples": FakeLatentTensor((1, 16, 64, 64)), "source": "pro_edit_prepare_region"} + + with self.assertRaisesRegex(RuntimeError, "Unsupported sampler_strategy 'mystery_strategy'"): + self.node.sample( + model=model, + positive=make_conditioning("positive"), + negative=make_conditioning("negative"), + latent_image=latent, + backend_mode="auto", + quality_preset="Manual", + seed=9, + steps=12, + cfg=1.0, + sampler_name="euler", + scheduler="simple", + denoise=0.8, + denoise_mode="manual", + flux_guidance=4.2, + model_family="Auto", + edit_info=None, + model_info=None, + ) +``` + +- [ ] **Step 2: Run the tests to verify they fail** + +Run: `env PYTHONDONTWRITEBYTECODE=1 python3 -m unittest discover -s tests -p 'test_pro_edit_bridge.py' -v` + +Expected: FAIL because `sample_info` does not yet expose `profile_id` / `sampler_strategy`, and bridge logic still branches directly on backend name. + +- [ ] **Step 3: Dispatch by `sampler_strategy`** + +```python +# pro_edit/pro_edit_bridge.py +from ..model_profiles.registry import resolve_model_profile +from ..utils.model_info import parse_jsonish_info + + +def _apply_sampler_strategy(strategy: str, *, positive, negative, flux_guidance, defaults): + if strategy == "standard_k": + return positive, negative, None + if strategy == "flux_guided": + guidance_value = _normalize_flux_guidance(flux_guidance, defaults.get("default_guidance")) + positive = set_conditioning_values(positive, {"guidance": guidance_value}) + negative = set_conditioning_values(negative, {"guidance": guidance_value}) + return positive, negative, guidance_value + raise RuntimeError(f"[LLS] Unsupported sampler_strategy '{strategy}'.") +``` + +```python +# pro_edit/pro_edit_bridge.py + merged_profile_info = dict(parse_jsonish_info(effective_model_info)) + merged_profile_info.update(parse_jsonish_info(edit_info)) + profile = resolve_model_profile(model=model, model_info=effective_model_info, extra_info=merged_profile_info) + backend, routing = resolve_backend( + backend_mode, + model=model, + model_info=effective_model_info, + edit_info=edit_info, + ) + defaults = get_family_defaults(profile["family"]) + positive, negative, guidance_value = _apply_sampler_strategy( + profile["sampler_strategy"], + positive=positive, + negative=negative, + flux_guidance=flux_guidance, + defaults=defaults, + ) +``` + +```python +# pro_edit/pro_edit_bridge.py + sample_info = info_to_json( + { + "backend_name": routing.backend_name, + "routing_reason": routing.routing_reason, + "family": profile["family"], + "model_role": profile["role"], + "profile_id": profile["profile_id"], + "backend_type": profile["backend_type"], + "sampler_strategy": profile["sampler_strategy"], + "seed": actual_seed, + "steps": int(steps), + "cfg": float(cfg), + "guidance": guidance_value, + "sampler_name": sampler_name, + "scheduler": scheduler, + "denoise": actual_denoise, + "denoise_mode": denoise_mode, + "quality_preset": quality_preset, + } + ) +``` + +```python +# pro_edit/pro_edit_finish.py + info = normalize_edit_info(edit_info) + info.setdefault("profile_id", "") + info.setdefault("backend_type", "none") + info.setdefault("sampler_strategy", "standard_k") +``` + +- [ ] **Step 4: Run the tests to verify they pass** + +Run: `env PYTHONDONTWRITEBYTECODE=1 python3 -m unittest discover -s tests -p 'test_pro_edit_bridge.py' -v` + +Expected: PASS with strategy-based routing and richer `sample_info`. + +- [ ] **Step 5: Commit** + +```bash +git add pro_edit/pro_edit_bridge.py pro_edit/pro_edit_finish.py tests/test_pro_edit_bridge.py +git commit -m "feat: dispatch pro bridge by sampler strategy" +``` + +### Task 7: Update Documentation And Run Full Verification + +**Files:** +- Modify: `README.md` +- Modify: `tests/test_pro_edit_docs.py` + +- [ ] **Step 1: Write the failing documentation tests** + +```python +# tests/test_pro_edit_docs.py + for needle in ( + "profile-driven routing", + "profile_id", + "backend_type", + "sampler_strategy", + "LLS Simple Checkpoint Loader writes the resolved model profile", + "FLUX base models should remain on the Simple workflow unless they resolve to flux_edit", + ): + self.assertIn(needle, text) +``` + +- [ ] **Step 2: Run the documentation tests to verify they fail** + +Run: `env PYTHONDONTWRITEBYTECODE=1 python3 -m unittest discover -s tests -p 'test_pro_edit_docs.py' -v` + +Expected: FAIL because README does not yet describe the profile-driven routing model. + +- [ ] **Step 3: Update the README** + +~~~markdown +## Pro Image Edit / Inpaint + +The Pro workflow now uses profile-driven routing. + +- `LLS Simple Checkpoint Loader` writes the resolved model profile into runtime metadata. +- `LLS Pro Image Edit Prepare` routes by `backend_type`. +- `LLS Pro KSampler Bridge` routes by `sampler_strategy`. + +Important profile fields: + +- `profile_id` +- `backend_type` +- `sampler_strategy` + +Examples: + +- `sdxl_inpaint` -> `backend_type = sdxl_native` +- `flux_edit` -> `backend_type = flux_edit` + +Base profiles do not silently use the Pro chain: + +- `SDXL base` remains a base profile unless explicitly overridden +- `FLUX base models should remain on the Simple workflow unless they resolve to flux_edit` + +Manual correction still happens through `model_info` overrides, for example: + +~~~json +{"profile_id":"flux_edit","backend_type":"flux_edit","sampler_strategy":"flux_guided","role":"edit"} +~~~ +~~~ + +- [ ] **Step 4: Run the focused docs test** + +Run: `env PYTHONDONTWRITEBYTECODE=1 python3 -m unittest discover -s tests -p 'test_pro_edit_docs.py' -v` + +Expected: PASS with one documentation test file. + +- [ ] **Step 5: Run full verification** + +Run: `env PYTHONDONTWRITEBYTECODE=1 python3 -m unittest discover -s tests -p 'test_model_profiles.py' -v` + +Expected: PASS. + +Run: `env PYTHONDONTWRITEBYTECODE=1 python3 -m unittest discover -s tests -p 'test_model_profile_loader.py' -v` + +Expected: PASS. + +Run: `env PYTHONDONTWRITEBYTECODE=1 python3 -m unittest discover -s tests -p 'test_pro_edit*.py' -v` + +Expected: PASS, with finish compositing tests allowed to skip when torch is unavailable in the local verification interpreter. + +Run: `env PYTHONDONTWRITEBYTECODE=1 python3 -m unittest discover -s tests -p 'test_model_info_inference.py' -v` + +Expected: PASS. + +Run: `env PYTHONDONTWRITEBYTECODE=1 python3 -m unittest discover -s tests -p 'test_repair*.py' -v` + +Expected: PASS and confirm the existing `LLSSimple*` repair chain remains unaffected. + +Run: `python3 -m compileall __init__.py model_profiles pro_edit utils model_loader` + +Expected: PASS with no syntax errors. + +- [ ] **Step 6: Commit** + +```bash +git add README.md tests/test_pro_edit_docs.py docs/superpowers/specs/2026-05-21-pro-model-profile-routing-design.md docs/superpowers/plans/2026-05-21-pro-model-profile-routing.md +git commit -m "docs: add profile driven pro routing design" +``` + +## Self-Review Checklist + +- Spec coverage: + - explicit `ModelProfile` contract is covered by Task 1 + - profile-aware `utils/model_info.py` compatibility is covered by Task 2 + - loader-side profile stamping is covered by Task 3 + - backend routing by `backend_type` is covered by Task 4 + - bridge routing by `sampler_strategy` is covered by Task 6 + - README and regression verification are covered by Task 7 +- Placeholder scan: + - no `TBD`, `TODO`, or “implement later” markers remain + - every task has exact file paths, exact commands, and concrete code snippets +- Type consistency: + - `profile_id`, `backend_type`, `sampler_strategy`, `loader_strategy` are used consistently across model profiles, loader tags, `edit_info`, and `sample_info` + - `sdxl_native` and `flux_edit` remain the only built-in Pro backend types in this iteration + - base profiles consistently use `backend_type = none` diff --git a/docs/superpowers/specs/2026-05-21-pro-model-profile-routing-design.md b/docs/superpowers/specs/2026-05-21-pro-model-profile-routing-design.md new file mode 100644 index 0000000..31c441e --- /dev/null +++ b/docs/superpowers/specs/2026-05-21-pro-model-profile-routing-design.md @@ -0,0 +1,402 @@ +# LLS Pro Model Profile Routing Design + +## Context + +`LLS-node` now has a working `pro_edit/` feature family with three nodes: + +- `LLS Pro Image Edit Prepare` +- `LLS Pro KSampler Bridge` +- `LLS Pro Image Edit Finish` + +The current implementation already separates `SDXL` and `FLUX` professional edit backends, but the routing contract is still too inference-heavy: + +- the loader infers a family and lightweight capability tags +- `pro_edit` resolves a backend from `family + role + capability flags` +- `role` often comes from checkpoint-name keyword guessing + +This works for a narrow set of obvious names such as `inpaint`, `edit`, `fill`, and now `kontext`, but it does not scale well to real-world model collections: + +- many `SD1.5` or `SDXL` derivatives share the same sampling core but differ in edit support +- many `FLUX`-family edit models use different naming conventions +- capability inference is duplicated across loader, routing, and bridge layers +- each node still partially reinterprets the model instead of reading a single authoritative profile + +The result is exactly the current user pain: + +- a model may clearly belong to `FLUX` or `SDXL`, but still fail the Pro workflow because `model_role` was inferred as `base` +- onboarding a new model often means adding more filename aliases +- routing logic becomes harder to trust as more families and edit variants are added + +## Goals + +- Keep the existing node names and primary ports: + - `LLS Pro Image Edit Prepare` + - `LLS Pro KSampler Bridge` + - `LLS Pro Image Edit Finish` +- Preserve existing workflow wiring as much as possible. +- Introduce a single explicit internal `ModelProfile` contract that becomes the authoritative source for: + - loader behavior + - backend routing + - sampler routing +- Stop treating checkpoint-name guessing as the main routing truth. +- Support large model families by adapting once per model type, not once per individual checkpoint. +- Keep the current `LLSSimple*` repair chain untouched and separate from the Pro chain. +- Keep auto-routing, but make it profile-driven instead of keyword-driven. +- Preserve manual override capability for debugging and exceptional models. + +## Non-Goals + +This refactor will not: + +- rename the existing Pro nodes +- force users to rebuild current Pro workflows from scratch +- replace the current `LLSSimple*` chain with the Pro chain +- add every future model backend in this iteration +- implement a full external YAML or JSON profile-management UI +- promise that arbitrary unknown edit models can work with zero metadata + +## Current Problem Statement + +The current Pro path has three architectural weaknesses: + +### 1. Inference Is Not Centralized + +`model_loader`, `utils/model_info.py`, `pro_edit/backends/registry.py`, and `pro_edit/pro_edit_bridge.py` all participate in deciding what the model is. This creates overlapping heuristics and weak ownership. + +### 2. Capability Flags Are Too Low-Level To Be The Main Contract + +Flags such as: + +- `supports_inpaint_native` +- `supports_image_edit_native` +- `preferred_edit_backend` + +are useful, but they are not enough to fully describe execution behavior. They do not explicitly tell the pipeline: + +- which backend implementation to use +- which sampler strategy to apply +- which loader assumptions are valid + +### 3. Model Family Is Not The Same As Execution Strategy + +A family such as `SDXL` or `FLUX_DEV` only tells part of the story. + +Examples: + +- `SDXL base` and `SDXL inpaint` can share core diffusion family defaults but should not share the same professional edit backend eligibility. +- `FLUX base` and `FLUX Kontext/edit` can share the same broad family while requiring different Pro-routing behavior. +- `SD1.5` belongs in the same unified profile system, but most `SD1.5` models should still route away from the current Pro backend set and remain on Simple workflows. + +## Chosen Architecture + +Introduce an explicit `ModelProfile` layer and make it the single authoritative routing contract. + +The new execution model is: + +1. The loader resolves a `ModelProfile`. +2. The loader writes the resolved profile into `_lls_*` metadata on `model`, `clip`, and `vae`. +3. `LLS Pro Image Edit Prepare` reads the resolved profile and dispatches by `backend_type`. +4. `LLS Pro KSampler Bridge` reads the same resolved profile and dispatches by `sampler_strategy`. +5. `LLS Pro Image Edit Finish` continues to composite primarily by `edit_scope`, but can also read `backend_type` for backend-specific finish behavior later. + +Checkpoint-name inference remains as a fallback source of evidence, but not as the primary truth source once a profile has been resolved. + +## Model Profile Contract + +Create a unified internal profile structure with these fields: + +- `profile_id` +- `family` +- `role` +- `backend_type` +- `sampler_strategy` +- `loader_strategy` +- `supports_inpaint_native` +- `supports_image_edit_native` +- `preferred_edit_backend` + +### Field Meanings + +- `profile_id` + - stable identifier for the resolved profile, such as `sdxl_base`, `sdxl_inpaint`, `flux_edit` +- `family` + - broad model family, for example `SD1.5`, `SDXL`, `FLUX_DEV` +- `role` + - semantic model role, such as `base`, `inpaint`, `edit`, `fill` +- `backend_type` + - Pro backend routing target, for example `none`, `sdxl_native`, `flux_edit` +- `sampler_strategy` + - bridge-time sampling strategy, for example `standard_k`, `flux_guided` +- `loader_strategy` + - loader-time resource strategy, for example `sd15_checkpoint`, `sdxl_checkpoint`, `flux_split_or_bundle` +- `supports_inpaint_native` + - whether the model can consume native inpaint-style conditioning +- `supports_image_edit_native` + - whether the model can consume native image-edit-style conditioning +- `preferred_edit_backend` + - compatibility hint for UI, migration, and legacy metadata interop + +## Built-In Profiles In This Iteration + +The first refactor will ship with built-in profiles for the currently supported scope: + +- `sd15_base` +- `sdxl_base` +- `sdxl_inpaint` +- `sdxl_edit` +- `flux_base` +- `flux_edit` + +### Intended Behavior + +- `sd15_base` + - stays outside the current Pro backend set + - keeps using normal/simple workflows unless future SD1.5 native edit support is added +- `sdxl_base` + - does not enter the current Pro backend set by default +- `sdxl_inpaint` + - maps to `backend_type = sdxl_native` + - maps to `sampler_strategy = standard_k` +- `sdxl_edit` + - maps to `backend_type = sdxl_native` + - maps to `sampler_strategy = standard_k` +- `flux_base` + - does not enter the current Pro backend set by default +- `flux_edit` + - maps to `backend_type = flux_edit` + - maps to `sampler_strategy = flux_guided` + +This structure ensures that many model variants can share one profile instead of requiring one-off code changes per checkpoint. + +## Profile Resolution Order + +Profile resolution must be deterministic and centralized. + +Resolution order: + +1. explicit `model_info` overrides +2. existing `_lls_*` profile metadata already attached to the runtime objects +3. built-in registry rules using: + - family + - known naming patterns + - capability hints +4. fallback base profile for the detected family + +### Manual Override Behavior + +Manual override remains supported through `model_info`, not through new required ports. + +The following keys should be accepted when present in `model_info`: + +- `profile_id` +- `family` +- `role` +- `backend_type` +- `sampler_strategy` +- `loader_strategy` +- `supports_inpaint_native` +- `supports_image_edit_native` +- `preferred_edit_backend` + +This keeps the current node surface stable while still allowing explicit correction when auto-resolution is wrong. + +## File Structure Changes + +Add a new package: + +- `model_profiles/__init__.py` +- `model_profiles/base.py` +- `model_profiles/registry.py` +- `model_profiles/rules.py` + +### Responsibilities + +- `model_profiles/base.py` + - define the profile record contract and normalization helpers +- `model_profiles/rules.py` + - define built-in family and pattern matching rules +- `model_profiles/registry.py` + - resolve the final `ModelProfile` from runtime inputs and overrides +- `model_profiles/__init__.py` + - export the public profile-resolution helpers + +## Existing File Changes + +### `utils/model_info.py` + +Keep this file as the low-level family and metadata helper layer, but reduce its authority: + +- keep family inference helpers +- keep JSON-like parsing helpers +- keep backwards-compatible accessors +- stop treating capability inference as the final routing truth +- add profile-aware wrappers where needed + +This file becomes a supporting utility for profile resolution instead of the main owner of routing semantics. + +### `model_loader/nodes.py` + +`LLSSimpleCheckpointLoader` becomes the authoritative profile stamping point. + +It must: + +- resolve family as it does today +- ask the new profile registry for the final `ModelProfile` +- tag `model`, `clip`, and `vae` with: + - `_lls_profile_id` + - `_lls_backend_type` + - `_lls_sampler_strategy` + - `_lls_loader_strategy` + - `_lls_model_role` + - `_lls_supports_inpaint_native` + - `_lls_supports_image_edit_native` + - `_lls_preferred_edit_backend` + +The loader remains family-aware for text encoder and VAE resolution, but those behaviors should now align with explicit `loader_strategy` values instead of only raw family checks. + +### `pro_edit/backends/registry.py` + +This file must stop doing broad capability inference. + +Its new responsibility is: + +- read resolved profile data +- validate `backend_mode` overrides against `backend_type` +- route to the correct backend implementation + +New routing logic should be driven by `backend_type`, not by repeated filename or role guessing. + +### `pro_edit/pro_edit_prepare.py` + +This node should: + +- resolve the effective profile +- reject profiles whose `backend_type` is `none` +- dispatch strictly by `backend_type` +- keep current ports and return values + +### `pro_edit/pro_edit_bridge.py` + +This node should: + +- resolve the same effective profile +- dispatch strictly by `sampler_strategy` +- keep current ports and return values +- continue using the shared low-level `_common_ksampler` where valid +- keep FLUX-specific guidance handling behind the `flux_guided` strategy instead of ad-hoc family checks + +### `pro_edit/pro_edit_finish.py` + +This node can remain primarily `edit_scope`-based in this iteration, but it should: + +- normalize and preserve `backend_type` in `edit_info` +- leave explicit extension points for future backend-specific finish behavior + +## Runtime Data Flow + +### Happy Path + +1. User loads a model through `LLS Simple Checkpoint Loader`. +2. Loader resolves family and final profile. +3. Loader writes profile metadata to `model`, `clip`, and `vae`. +4. `LLS Pro Image Edit Prepare` reads profile metadata and chooses the backend via `backend_type`. +5. Backend-specific prepare logic builds native edit conditioning. +6. `LLS Pro KSampler Bridge` reads the same profile and applies the right `sampler_strategy`. +7. `LLS Pro Image Edit Finish` composites the generated output according to `edit_scope`. + +### Override Path + +1. Auto profile resolution produces an incorrect result or a generic base profile. +2. User supplies `model_info` override fields. +3. Profile registry merges those overrides into the final profile. +4. Pro nodes use the overridden explicit profile instead of fallback inference. + +## Simple vs Pro Separation + +The separation must become stricter, not weaker. + +- `Simple` remains the compatibility-first path for broad masked resampling. +- `Pro` remains the native local edit path for true edit/inpaint-capable profiles. + +The Pro chain should not silently downgrade to the Simple chain when a profile is `base` or unsupported. It should fail with a clear error that explains: + +- the detected family +- the detected profile +- the missing backend eligibility + +This keeps workflow behavior predictable. + +## Error Handling + +The new profile-driven path should improve the error messages. + +### Unsupported Profile For Pro Chain + +If `backend_type = none`, raise an error such as: + +> `[LLS] Pro image edit is not available for profile '' (family '', role ''). Use a supported native edit/inpaint profile or provide an explicit override.` + +### Invalid Manual Override + +If the user forces `backend_mode = flux` but the resolved profile maps to `sdxl_native`, raise a clear mismatch error instead of trying to proceed. + +### Incomplete Override + +If the user provides a partial override that produces an invalid combination, the profile registry should normalize or reject it immediately instead of letting the mismatch leak deeper into the pipeline. + +## Compatibility Strategy + +The refactor must preserve current workflow usability: + +- existing node names stay unchanged +- existing ports remain valid +- `model_info` remains the optional override surface +- old `_lls_*` capability tags continue to be read during migration + +During migration, the profile resolver may derive `backend_type` and `sampler_strategy` from old capability fields when a full explicit profile is not yet present. + +This prevents a hard break between the current Pro implementation and the new profile-driven contract. + +## Testing Strategy + +Add focused tests for the new architecture in these categories: + +### Profile Resolution Tests + +- family base profiles resolve correctly +- explicit `model_info` overrides take precedence +- legacy capability tags still map to the correct profile during migration +- common alias patterns such as `kontext` resolve to the intended edit profile + +### Loader Tagging Tests + +- loader writes the new `_lls_profile_*` metadata +- loader keeps writing the existing compatibility fields + +### Backend Routing Tests + +- `backend_type = sdxl_native` routes to the SDXL backend +- `backend_type = flux_edit` routes to the FLUX backend +- `backend_type = none` fails clearly in Pro prepare +- forced backend mismatches are rejected clearly + +### Sampler Strategy Tests + +- `sampler_strategy = standard_k` uses the shared standard bridge path +- `sampler_strategy = flux_guided` injects FLUX guidance behavior + +### Regression Tests + +- current Pro SDXL tests still pass +- current Pro FLUX tests still pass +- current `LLSSimple*` repair tests still pass + +## Recommended Migration Outcome + +After this refactor: + +- adding a new model should usually mean adding or adjusting a profile rule +- adding a truly new model type should mean adding a new profile plus backend or sampler strategy +- neither path should require scattering new filename heuristics across loader, registry, and bridge code + +This is the core success criterion for the refactor. diff --git a/model_loader/nodes.py b/model_loader/nodes.py index cefbd6b..36b2f22 100644 --- a/model_loader/nodes.py +++ b/model_loader/nodes.py @@ -34,9 +34,9 @@ except Exception as exc: else: _CORE_NODES_ERR = None +from ..model_profiles.registry import resolve_model_profile from ..utils.model_info import ( MODEL_FAMILY_CHOICES, - build_model_info, canonicalize_family, get_family_defaults, infer_family_from_name, @@ -62,16 +62,25 @@ _FLUX_VAE_PATTERNS = ("ae.safetensors", "ae", "vae") def _build_capability_tags(model_name: str, family: str) -> dict[str, object]: - info = build_model_info( + profile = resolve_model_profile( + model=None, + model_info={ + "checkpoint_name": model_name, + "model_name": model_name, + "family": family, + }, checkpoint_name=model_name, - model_name=model_name, family=family, ) return { - "model_role": info["model_role"], - "supports_inpaint_native": info["supports_inpaint_native"], - "supports_image_edit_native": info["supports_image_edit_native"], - "preferred_edit_backend": info["preferred_edit_backend"], + "profile_id": profile["profile_id"], + "backend_type": profile["backend_type"], + "sampler_strategy": profile["sampler_strategy"], + "loader_strategy": profile["loader_strategy"], + "model_role": profile["role"], + "supports_inpaint_native": profile["supports_inpaint_native"], + "supports_image_edit_native": profile["supports_image_edit_native"], + "preferred_edit_backend": profile["preferred_edit_backend"], } diff --git a/model_profiles/__init__.py b/model_profiles/__init__.py new file mode 100644 index 0000000..f9ce97b --- /dev/null +++ b/model_profiles/__init__.py @@ -0,0 +1,9 @@ +from .base import ModelProfile, build_model_profile, normalize_profile_dict +from .registry import resolve_model_profile + +__all__ = [ + "ModelProfile", + "build_model_profile", + "normalize_profile_dict", + "resolve_model_profile", +] diff --git a/model_profiles/base.py b/model_profiles/base.py new file mode 100644 index 0000000..13a2b24 --- /dev/null +++ b/model_profiles/base.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +from dataclasses import asdict, dataclass +from typing import Any + + +def _normalize_text(value: Any, default: str = "", *, lower: bool = False) -> str: + if value is None: + return default + normalized = str(value).strip() + if not normalized: + return default + return normalized.lower() if lower else normalized + + +def _coerce_bool(value: Any) -> bool: + if isinstance(value, bool): + return value + if value is None: + return False + if isinstance(value, str): + lowered = value.strip().lower() + if lowered in {"", "0", "false", "no", "off"}: + return False + if lowered in {"1", "true", "yes", "on"}: + return True + return bool(value) + + +@dataclass(frozen=True) +class ModelProfile: + profile_id: str + family: str + role: str + backend_type: str + sampler_strategy: str + loader_strategy: str + supports_inpaint_native: bool + supports_image_edit_native: bool + preferred_edit_backend: str | None + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +def normalize_profile_dict(value: dict[str, Any]) -> dict[str, Any]: + preferred_backend = value.get("preferred_edit_backend") + return { + "profile_id": _normalize_text(value.get("profile_id"), lower=True), + "family": _normalize_text(value.get("family") or value.get("model_family"), "SD1.5"), + "role": _normalize_text(value.get("role") or value.get("model_role"), "base", lower=True), + "backend_type": _normalize_text(value.get("backend_type"), "none", lower=True), + "sampler_strategy": _normalize_text(value.get("sampler_strategy"), "standard_k", lower=True), + "loader_strategy": _normalize_text(value.get("loader_strategy"), lower=True), + "supports_inpaint_native": _coerce_bool(value.get("supports_inpaint_native", False)), + "supports_image_edit_native": _coerce_bool(value.get("supports_image_edit_native", False)), + "preferred_edit_backend": ( + None if preferred_backend in (None, "") else _normalize_text(preferred_backend, lower=True) + ), + } + + +def build_model_profile(**kwargs: Any) -> ModelProfile: + normalized = normalize_profile_dict(kwargs) + return ModelProfile(**normalized) diff --git a/model_profiles/registry.py b/model_profiles/registry.py new file mode 100644 index 0000000..bfbefe2 --- /dev/null +++ b/model_profiles/registry.py @@ -0,0 +1,145 @@ +from __future__ import annotations + +from typing import Any + +from .base import build_model_profile +from .rules import match_builtin_profile + +try: + from ..utils.model_info import canonicalize_family, get_lls_attr, parse_jsonish_info +except ImportError: + from utils.model_info import canonicalize_family, get_lls_attr, parse_jsonish_info + + +def _pick(*values: Any) -> Any: + for value in values: + if value not in (None, ""): + return value + return None + + +def _derive_role_from_legacy( + family: str | None, + role: Any, + *, + supports_inpaint_native: Any, + supports_image_edit_native: Any, + preferred_edit_backend: Any, +) -> str | None: + normalized_role = str(role or "").strip().lower() + if normalized_role and normalized_role != "base": + return normalized_role + + resolved_family = canonicalize_family(family) + preferred = str(preferred_edit_backend or "").strip().lower() + inpaint_native = supports_inpaint_native not in (None, "", False, 0, "0", "false", "False") + image_edit_native = supports_image_edit_native not in (None, "", False, 0, "0", "false", "False") + + if resolved_family.startswith("SDXL"): + if image_edit_native: + return "edit" + if inpaint_native or preferred == "sdxl": + return "inpaint" + if resolved_family.startswith("FLUX"): + if inpaint_native: + return "inpaint" + if image_edit_native or preferred == "flux": + return "edit" + return normalized_role or None + + +def _extract_profile_overrides(model=None, model_info=None, extra_info=None) -> dict[str, object]: + raw = parse_jsonish_info(model_info) + extra = parse_jsonish_info(extra_info) + + return { + "profile_id": _pick(extra.get("profile_id"), raw.get("profile_id"), get_lls_attr(model, "profile_id", None)), + "family": _pick( + extra.get("family"), + extra.get("model_family"), + raw.get("family"), + raw.get("model_family"), + get_lls_attr(model, "family", None), + ), + "role": _pick( + extra.get("role"), + extra.get("model_role"), + raw.get("role"), + raw.get("model_role"), + get_lls_attr(model, "model_role", None), + ), + "backend_type": _pick(extra.get("backend_type"), raw.get("backend_type"), get_lls_attr(model, "backend_type", None)), + "sampler_strategy": _pick( + extra.get("sampler_strategy"), + raw.get("sampler_strategy"), + get_lls_attr(model, "sampler_strategy", None), + ), + "loader_strategy": _pick(extra.get("loader_strategy"), raw.get("loader_strategy"), get_lls_attr(model, "loader_strategy", None)), + "supports_inpaint_native": ( + extra["supports_inpaint_native"] if "supports_inpaint_native" in extra + else raw["supports_inpaint_native"] if "supports_inpaint_native" in raw + else get_lls_attr(model, "supports_inpaint_native", None) + ), + "supports_image_edit_native": ( + extra["supports_image_edit_native"] if "supports_image_edit_native" in extra + else raw["supports_image_edit_native"] if "supports_image_edit_native" in raw + else get_lls_attr(model, "supports_image_edit_native", None) + ), + "preferred_edit_backend": _pick( + extra.get("preferred_edit_backend"), + raw.get("preferred_edit_backend"), + get_lls_attr(model, "preferred_edit_backend", None), + ), + "model_name": _pick( + extra.get("checkpoint_name"), + extra.get("model_name"), + raw.get("checkpoint_name"), + raw.get("model_name"), + raw.get("ckpt_name"), + raw.get("ckpt"), + get_lls_attr(model, "checkpoint_name", None), + get_lls_attr(model, "model_name", None), + ), + } + + +def resolve_model_profile( + model=None, + model_info=None, + extra_info=None, + checkpoint_name: str | None = None, + family: str | None = None, +) -> dict[str, object]: + overrides = _extract_profile_overrides(model=model, model_info=model_info, extra_info=extra_info) + resolved_family = canonicalize_family(family or overrides.get("family") or "SD1.5") + resolved_role = _derive_role_from_legacy( + resolved_family, + overrides.get("role"), + supports_inpaint_native=overrides.get("supports_inpaint_native"), + supports_image_edit_native=overrides.get("supports_image_edit_native"), + preferred_edit_backend=overrides.get("preferred_edit_backend"), + ) + model_name = str(checkpoint_name or overrides.get("model_name") or "").strip() + + merged = dict(match_builtin_profile(model_name, resolved_family, resolved_role)) + if overrides.get("profile_id") not in (None, ""): + merged["profile_id"] = overrides["profile_id"] + if overrides.get("family") not in (None, ""): + merged["family"] = canonicalize_family(overrides["family"]) + if overrides.get("role") not in (None, ""): + merged["role"] = str(overrides["role"]).strip().lower() + elif resolved_role: + merged["role"] = str(resolved_role).strip().lower() + if overrides.get("backend_type") not in (None, ""): + merged["backend_type"] = overrides["backend_type"] + if overrides.get("sampler_strategy") not in (None, ""): + merged["sampler_strategy"] = overrides["sampler_strategy"] + if overrides.get("loader_strategy") not in (None, ""): + merged["loader_strategy"] = overrides["loader_strategy"] + if overrides.get("preferred_edit_backend") not in (None, ""): + merged["preferred_edit_backend"] = overrides["preferred_edit_backend"] + for key in ("supports_inpaint_native", "supports_image_edit_native"): + if overrides.get(key) is not None: + merged[key] = overrides[key] + + return build_model_profile(**merged).to_dict() diff --git a/model_profiles/rules.py b/model_profiles/rules.py new file mode 100644 index 0000000..c454169 --- /dev/null +++ b/model_profiles/rules.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +try: + from ..utils.model_info import ( + canonicalize_family, + infer_family_from_name, + infer_model_role_from_name, + is_flux_family, + is_sdxl_family, + ) +except ImportError: + from utils.model_info import ( + canonicalize_family, + infer_family_from_name, + infer_model_role_from_name, + is_flux_family, + is_sdxl_family, + ) + + +_ROLE_TO_PROFILE = { + ("SDXL", "inpaint"): ("sdxl_inpaint", "sdxl_native", "standard_k", "sdxl_checkpoint", True, False, "sdxl"), + ("SDXL", "edit"): ("sdxl_edit", "sdxl_native", "standard_k", "sdxl_checkpoint", True, True, "sdxl"), + ("SDXL", "fill"): ("sdxl_edit", "sdxl_native", "standard_k", "sdxl_checkpoint", True, True, "sdxl"), + ("FLUX_DEV", "inpaint"): ("flux_edit", "flux_edit", "flux_guided", "flux_split_or_bundle", True, True, "flux"), + ("FLUX_DEV", "edit"): ("flux_edit", "flux_edit", "flux_guided", "flux_split_or_bundle", False, True, "flux"), + ("FLUX_DEV", "fill"): ("flux_edit", "flux_edit", "flux_guided", "flux_split_or_bundle", False, True, "flux"), + ("FLUX_SCHNELL", "inpaint"): ("flux_edit", "flux_edit", "flux_guided", "flux_split_or_bundle", True, True, "flux"), + ("FLUX_SCHNELL", "edit"): ("flux_edit", "flux_edit", "flux_guided", "flux_split_or_bundle", False, True, "flux"), + ("FLUX_SCHNELL", "fill"): ("flux_edit", "flux_edit", "flux_guided", "flux_split_or_bundle", False, True, "flux"), +} + + +def _base_profile_for_family(family: str) -> dict[str, object]: + resolved_family = canonicalize_family(family) + if is_sdxl_family(resolved_family): + return { + "profile_id": "sdxl_base", + "family": resolved_family, + "role": "base", + "backend_type": "none", + "sampler_strategy": "standard_k", + "loader_strategy": "sdxl_checkpoint", + "supports_inpaint_native": False, + "supports_image_edit_native": False, + "preferred_edit_backend": None, + } + if is_flux_family(resolved_family): + return { + "profile_id": "flux_base", + "family": resolved_family, + "role": "base", + "backend_type": "none", + "sampler_strategy": "flux_guided", + "loader_strategy": "flux_split_or_bundle", + "supports_inpaint_native": False, + "supports_image_edit_native": False, + "preferred_edit_backend": None, + } + return { + "profile_id": "sd15_base", + "family": "SD1.5", + "role": "base", + "backend_type": "none", + "sampler_strategy": "standard_k", + "loader_strategy": "sd15_checkpoint", + "supports_inpaint_native": False, + "supports_image_edit_native": False, + "preferred_edit_backend": None, + } + + +def match_builtin_profile(model_name: str | None, family: str | None, role: str | None = None) -> dict[str, object]: + resolved_family = canonicalize_family(family or infer_family_from_name(model_name, "SD1.5")) + resolved_role = str(role or infer_model_role_from_name(model_name, resolved_family) or "base").strip().lower() + matched = _ROLE_TO_PROFILE.get((resolved_family, resolved_role)) + if matched is None: + return _base_profile_for_family(resolved_family) + profile_id, backend_type, sampler_strategy, loader_strategy, supports_inpaint_native, supports_image_edit_native, preferred_backend = matched + return { + "profile_id": profile_id, + "family": resolved_family, + "role": resolved_role, + "backend_type": backend_type, + "sampler_strategy": sampler_strategy, + "loader_strategy": loader_strategy, + "supports_inpaint_native": supports_inpaint_native, + "supports_image_edit_native": supports_image_edit_native, + "preferred_edit_backend": preferred_backend, + } diff --git a/pro_edit/backends/base.py b/pro_edit/backends/base.py index fb815da..be4dfbc 100644 --- a/pro_edit/backends/base.py +++ b/pro_edit/backends/base.py @@ -9,12 +9,13 @@ class RoutingResult: backend_name: str routing_reason: str capabilities: dict[str, Any] + profile: dict[str, Any] class ProEditBackend(Protocol): backend_name: str - def supports(self, capabilities: dict[str, Any]) -> bool: + def supports(self, profile: dict[str, Any]) -> bool: raise NotImplementedError def prepare(self, **kwargs) -> dict[str, Any]: @@ -24,10 +25,10 @@ class ProEditBackend(Protocol): raise NotImplementedError -def validate_manual_backend(backend_name: str, backend: ProEditBackend, capabilities: dict[str, Any]) -> None: - if backend.supports(capabilities): +def validate_manual_backend(backend_name: str, backend: ProEditBackend, profile: dict[str, Any]) -> None: + if backend.supports(profile): return - family = capabilities.get("model_family") or "UNKNOWN" raise RuntimeError( - f"[LLS] Pro edit backend '{backend_name}' is incompatible with model family '{family}'." + f"[LLS] Pro edit backend '{backend_name}' is incompatible with profile " + f"'{profile.get('profile_id')}' (family '{profile.get('family')}', role '{profile.get('role')}')." ) diff --git a/pro_edit/backends/flux.py b/pro_edit/backends/flux.py index e103f97..9aeceb8 100644 --- a/pro_edit/backends/flux.py +++ b/pro_edit/backends/flux.py @@ -5,15 +5,8 @@ from ..pro_edit_utils import build_native_conditioning_payload, set_conditioning class FluxProEditBackend: backend_name = "flux" - def supports(self, capabilities): - family = str(capabilities.get("model_family") or "") - role = str(capabilities.get("model_role") or "") - preferred = str(capabilities.get("preferred_edit_backend") or "").strip().lower() - return family.startswith("FLUX") and ( - role in {"inpaint", "edit", "fill"} - or bool(capabilities.get("supports_image_edit_native")) - or preferred == self.backend_name - ) + def supports(self, profile): + return str(profile.get("backend_type") or "").strip().lower() == "flux_edit" def prepare(self, *, vae, work_image, work_mask, positive, negative, workspace, routing, **_kwargs): latent, concat_latent_image, concat_mask = build_native_conditioning_payload( diff --git a/pro_edit/backends/registry.py b/pro_edit/backends/registry.py index 4c87663..3d9a333 100644 --- a/pro_edit/backends/registry.py +++ b/pro_edit/backends/registry.py @@ -3,7 +3,8 @@ from __future__ import annotations import importlib from .base import RoutingResult, validate_manual_backend -from ...utils.model_info import parse_jsonish_info, resolve_edit_capabilities +from ...model_profiles.registry import resolve_model_profile +from ...utils.model_info import get_lls_attr, parse_jsonish_info _BACKENDS = {} @@ -32,58 +33,79 @@ def get_backend(name: str): return backend -def _normalize_capabilities(model=None, model_info=None, edit_info=None): - resolved = resolve_edit_capabilities(model=model, model_info=model_info) +def _resolve_profile_and_capabilities(model=None, model_info=None, edit_info=None): + raw_model_info = parse_jsonish_info(model_info) raw_edit_info = parse_jsonish_info(edit_info) - - if raw_edit_info.get("model_family"): - resolved["model_family"] = str(raw_edit_info["model_family"]) - if raw_edit_info.get("model_name"): - resolved["model_name"] = str(raw_edit_info["model_name"]) - if raw_edit_info.get("model_role"): - resolved["model_role"] = str(raw_edit_info["model_role"]) - if "supports_inpaint_native" in raw_edit_info: - resolved["supports_inpaint_native"] = bool(raw_edit_info["supports_inpaint_native"]) - if "supports_image_edit_native" in raw_edit_info: - resolved["supports_image_edit_native"] = bool(raw_edit_info["supports_image_edit_native"]) - if raw_edit_info.get("preferred_edit_backend"): - resolved["preferred_edit_backend"] = str(raw_edit_info["preferred_edit_backend"]).strip().lower() - - return resolved, raw_edit_info + merged_extra_info = dict(raw_model_info) + merged_extra_info.update(raw_edit_info) + profile = resolve_model_profile( + model=model, + model_info=model_info, + extra_info=merged_extra_info, + checkpoint_name=str( + merged_extra_info.get("checkpoint_name") + or merged_extra_info.get("model_name") + or get_lls_attr(model, "checkpoint_name", None) + or get_lls_attr(model, "model_name", "") + or "" + ), + family=( + merged_extra_info.get("family") + or merged_extra_info.get("model_family") + or get_lls_attr(model, "family", None) + or None + ), + ) + capabilities = { + "model_family": profile["family"], + "model_name": str( + merged_extra_info.get("checkpoint_name") + or merged_extra_info.get("model_name") + or get_lls_attr(model, "checkpoint_name", None) + or get_lls_attr(model, "model_name", "") + or "" + ), + "model_role": profile["role"], + "profile_id": profile["profile_id"], + "backend_type": profile["backend_type"], + "sampler_strategy": profile["sampler_strategy"], + "loader_strategy": profile["loader_strategy"], + "supports_inpaint_native": profile["supports_inpaint_native"], + "supports_image_edit_native": profile["supports_image_edit_native"], + "preferred_edit_backend": profile["preferred_edit_backend"], + } + return profile, capabilities def resolve_backend(backend_mode: str, *, model=None, model_info=None, edit_info=None): _ensure_builtin_backends() mode = str(backend_mode or "auto").strip().lower() or "auto" - capabilities, raw_edit_info = _normalize_capabilities( + profile, capabilities = _resolve_profile_and_capabilities( model=model, model_info=model_info, edit_info=edit_info, ) + backend_type = str(profile.get("backend_type") or "none").strip().lower() + + if backend_type == "none": + raise RuntimeError( + f"[LLS] Pro image edit is not available for profile '{profile['profile_id']}' " + f"(family '{profile['family']}', role '{profile['role']}'). " + "Use a supported native edit/inpaint profile or provide an explicit override." + ) + + backend_name = { + "sdxl_native": "sdxl", + "flux_edit": "flux", + }.get(backend_type) + if backend_name is None: + raise RuntimeError(f"[LLS] Unsupported pro edit backend_type '{backend_type}'.") if mode != "auto": backend = get_backend(mode) - validate_manual_backend(mode, backend, capabilities) - return backend, RoutingResult(backend.backend_name, "manual.override", capabilities) + validate_manual_backend(mode, backend, profile) + return backend, RoutingResult(backend.backend_name, "manual.override", capabilities, profile) - backend_name = str(raw_edit_info.get("backend_name") or "").strip().lower() - if backend_name: - backend = get_backend(backend_name) - validate_manual_backend(backend_name, backend, capabilities) - return backend, RoutingResult(backend.backend_name, "edit_info.backend_name", capabilities) - - preferred = str(capabilities.get("preferred_edit_backend") or "").strip().lower() - if preferred: - backend = get_backend(preferred) - validate_manual_backend(preferred, backend, capabilities) - return backend, RoutingResult(backend.backend_name, "model.preferred_edit_backend", capabilities) - - for candidate_name in sorted(_BACKENDS): - backend = get_backend(candidate_name) - if backend.supports(capabilities): - return backend, RoutingResult(backend.backend_name, "model.capability_match", capabilities) - - raise RuntimeError( - "[LLS] No professional edit backend matched the current model capability. " - "Use an edit-capable SDXL or FLUX model, or provide explicit capability metadata." - ) + backend = get_backend(backend_name) + validate_manual_backend(backend_name, backend, profile) + return backend, RoutingResult(backend.backend_name, "profile.backend_type", capabilities, profile) diff --git a/pro_edit/backends/sdxl.py b/pro_edit/backends/sdxl.py index 8be3255..f9deebf 100644 --- a/pro_edit/backends/sdxl.py +++ b/pro_edit/backends/sdxl.py @@ -5,15 +5,8 @@ from ..pro_edit_utils import build_native_conditioning_payload, set_conditioning class SDXLProEditBackend: backend_name = "sdxl" - def supports(self, capabilities): - family = str(capabilities.get("model_family") or "") - role = str(capabilities.get("model_role") or "") - preferred = str(capabilities.get("preferred_edit_backend") or "").strip().lower() - return family.startswith("SDXL") and ( - role in {"inpaint", "edit", "fill"} - or bool(capabilities.get("supports_inpaint_native")) - or preferred == self.backend_name - ) + def supports(self, profile): + return str(profile.get("backend_type") or "").strip().lower() == "sdxl_native" def prepare(self, *, vae, work_image, work_mask, positive, negative, workspace, routing, **_kwargs): latent, concat_latent_image, concat_mask = build_native_conditioning_payload( diff --git a/pro_edit/pro_edit_bridge.py b/pro_edit/pro_edit_bridge.py index 8b60250..1646861 100644 --- a/pro_edit/pro_edit_bridge.py +++ b/pro_edit/pro_edit_bridge.py @@ -20,6 +20,18 @@ from ..utils.model_info import ( ) +def _apply_sampler_strategy(strategy: str, *, positive, negative, flux_guidance, defaults): + normalized = str(strategy or "").strip().lower() or "standard_k" + if normalized == "standard_k": + return positive, negative, None + if normalized == "flux_guided": + guidance_value = _normalize_flux_guidance(flux_guidance, defaults.get("default_guidance")) + positive = set_conditioning_values(positive, {"guidance": guidance_value}) + negative = set_conditioning_values(negative, {"guidance": guidance_value}) + return positive, negative, guidance_value + raise RuntimeError(f"[LLS] Unsupported sampler_strategy '{strategy}'.") + + class LLSProKSamplerBridge: CATEGORY = "LLS/Image Edit" FUNCTION = "sample" @@ -82,13 +94,14 @@ class LLSProKSamplerBridge: elif model_info is None and normalized_model_family not in {"", "Auto", "auto"}: effective_model_info = {"family": normalized_model_family} - backend, routing = resolve_backend( + _, routing = resolve_backend( backend_mode, model=model, model_info=effective_model_info, edit_info=edit_info, ) - defaults = get_family_defaults(routing.capabilities["model_family"]) + profile = routing.profile + defaults = get_family_defaults(profile["family"]) if quality_preset == FAMILY_DEFAULT_PRESET: steps = int(defaults["default_steps"]) @@ -97,7 +110,7 @@ class LLSProKSamplerBridge: scheduler = str(defaults["default_scheduler"]) denoise = float(defaults["default_denoise"]) else: - preset = get_sampling_preset({"family": routing.capabilities["model_family"]}, quality_preset) + preset = get_sampling_preset({"family": profile["family"]}, quality_preset) if preset is not None: steps = int(preset["steps"]) cfg = float(preset["cfg"]) @@ -110,11 +123,13 @@ class LLSProKSamplerBridge: if denoise_mode == "auto_from_edit": actual_denoise = float(normalized_edit_info.get("recommended_denoise", actual_denoise)) - guidance_value = None - if routing.backend_name == "flux": - guidance_value = _normalize_flux_guidance(flux_guidance, defaults.get("default_guidance")) - positive = set_conditioning_values(positive, {"guidance": guidance_value}) - negative = set_conditioning_values(negative, {"guidance": guidance_value}) + positive, negative, guidance_value = _apply_sampler_strategy( + profile["sampler_strategy"], + positive=positive, + negative=negative, + flux_guidance=flux_guidance, + defaults=defaults, + ) actual_seed = random.randint(0, 0xFFFFFFFFFFFFFFFF) if int(seed) == -1 else int(seed) result_latent = _common_ksampler( @@ -133,8 +148,11 @@ class LLSProKSamplerBridge: { "backend_name": routing.backend_name, "routing_reason": routing.routing_reason, - "family": routing.capabilities["model_family"], - "model_role": routing.capabilities["model_role"], + "family": profile["family"], + "model_role": profile["role"], + "profile_id": profile["profile_id"], + "backend_type": profile["backend_type"], + "sampler_strategy": profile["sampler_strategy"], "seed": actual_seed, "steps": int(steps), "cfg": float(cfg), diff --git a/pro_edit/pro_edit_utils.py b/pro_edit/pro_edit_utils.py index 8939ddc..5362c73 100644 --- a/pro_edit/pro_edit_utils.py +++ b/pro_edit/pro_edit_utils.py @@ -49,6 +49,12 @@ def normalize_edit_info(edit_info): info = parse_jsonish_info(edit_info) info.setdefault("backend_name", "") info.setdefault("routing_reason", "") + info.setdefault("model_family", "SD1.5") + info.setdefault("model_role", "base") + info.setdefault("profile_id", "") + info.setdefault("backend_type", "none") + info.setdefault("sampler_strategy", "standard_k") + info.setdefault("loader_strategy", "") info.setdefault("edit_scope", "region") info.setdefault("edit_payload_version", "1.0") return info @@ -224,14 +230,19 @@ def build_workspace( def build_edit_info(workspace: dict[str, Any], routing, *, backend_hints: dict[str, Any] | None = None) -> dict[str, Any]: + profile = routing.profile info = { "backend_name": routing.backend_name, "routing_reason": routing.routing_reason, - "model_family": routing.capabilities["model_family"], - "model_role": routing.capabilities["model_role"], - "supports_inpaint_native": routing.capabilities["supports_inpaint_native"], - "supports_image_edit_native": routing.capabilities["supports_image_edit_native"], - "preferred_edit_backend": routing.capabilities["preferred_edit_backend"], + "model_family": profile["family"], + "model_role": profile["role"], + "profile_id": profile["profile_id"], + "backend_type": profile["backend_type"], + "sampler_strategy": profile["sampler_strategy"], + "loader_strategy": profile["loader_strategy"], + "supports_inpaint_native": profile["supports_inpaint_native"], + "supports_image_edit_native": profile["supports_image_edit_native"], + "preferred_edit_backend": profile["preferred_edit_backend"], "edit_scope": workspace["edit_scope"], "original_size": list(workspace["original_size"]), "work_size": list(workspace["work_size"]), diff --git a/tests/test_model_profile_loader.py b/tests/test_model_profile_loader.py new file mode 100644 index 0000000..9619a42 --- /dev/null +++ b/tests/test_model_profile_loader.py @@ -0,0 +1,46 @@ +import unittest +from types import SimpleNamespace +from unittest import mock + +try: + from .test_pro_edit_helpers import import_plugin_submodule, load_plugin_package +except ImportError: + from test_pro_edit_helpers import import_plugin_submodule, load_plugin_package + + +class TestModelProfileLoader(unittest.TestCase): + def test_loader_writes_profile_fields_to_outputs(self): + plugin = load_plugin_package() + loader_module = import_plugin_submodule(plugin, "model_loader.nodes") + loader = loader_module.LLSSimpleCheckpointLoader() + model = SimpleNamespace() + clip = SimpleNamespace() + vae = SimpleNamespace() + + with mock.patch.object(loader_module, "folder_paths", object()), \ + mock.patch.object(loader_module, "comfy_sd", object()), \ + mock.patch.object(loader_module, "_resolve_model_path", return_value=("diffusion_models", "/fake/demo-flux-kontext-dev.safetensors")), \ + mock.patch.object(loader_module, "_load_model", return_value=(model, clip, vae)), \ + mock.patch.object(loader_module, "_resolve_text_encoder", return_value=(clip, "embedded", None, None)), \ + mock.patch.object(loader_module, "_resolve_vae", return_value=(vae, "embedded", None)): + loaded_model, loaded_clip, loaded_vae, text_encoder = loader.load_checkpoint( + ckpt_name="demo-flux-kontext-dev.safetensors", + model_family="FLUX_DEV", + load_mode="simple", + vae_source="auto", + text_encoder_source="auto", + external_vae_name=loader_module.AUTO_PLACEHOLDER, + external_text_encoder_1=loader_module.AUTO_PLACEHOLDER, + external_text_encoder_2=loader_module.AUTO_PLACEHOLDER, + ) + + for obj in (loaded_model, loaded_clip, loaded_vae, text_encoder): + self.assertEqual(obj._lls_profile_id, "flux_edit") + self.assertEqual(obj._lls_backend_type, "flux_edit") + self.assertEqual(obj._lls_sampler_strategy, "flux_guided") + self.assertEqual(obj._lls_loader_strategy, "flux_split_or_bundle") + self.assertEqual(obj._lls_model_role, "edit") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_model_profiles.py b/tests/test_model_profiles.py new file mode 100644 index 0000000..7a68ae1 --- /dev/null +++ b/tests/test_model_profiles.py @@ -0,0 +1,80 @@ +import unittest + +try: + from .test_pro_edit_helpers import FakeModel, import_plugin_submodule, load_plugin_package +except ImportError: + from test_pro_edit_helpers import FakeModel, import_plugin_submodule, load_plugin_package + + +class TestModelProfiles(unittest.TestCase): + def setUp(self): + self.plugin = load_plugin_package() + self.registry = import_plugin_submodule(self.plugin, "model_profiles.registry") + + def test_resolves_sdxl_base_profile_from_family(self): + profile = self.registry.resolve_model_profile( + model=None, + model_info={"checkpoint_name": "plain-sdxl-base.safetensors", "family": "SDXL"}, + ) + + self.assertEqual(profile["profile_id"], "sdxl_base") + self.assertEqual(profile["backend_type"], "none") + self.assertEqual(profile["sampler_strategy"], "standard_k") + self.assertEqual(profile["loader_strategy"], "sdxl_checkpoint") + + def test_resolves_flux_kontext_profile_from_name(self): + profile = self.registry.resolve_model_profile( + model=None, + model_info={"checkpoint_name": "demo-flux-kontext-dev.safetensors", "family": "FLUX_DEV"}, + ) + + self.assertEqual(profile["profile_id"], "flux_edit") + self.assertEqual(profile["role"], "edit") + self.assertEqual(profile["backend_type"], "flux_edit") + self.assertEqual(profile["sampler_strategy"], "flux_guided") + self.assertEqual(profile["loader_strategy"], "flux_split_or_bundle") + + def test_model_info_override_has_highest_priority(self): + profile = self.registry.resolve_model_profile( + model=None, + model_info={ + "checkpoint_name": "plain-sdxl-base.safetensors", + "family": "SDXL", + "profile_id": "sdxl_edit", + "role": "edit", + "backend_type": "sdxl_native", + "sampler_strategy": "standard_k", + "supports_inpaint_native": True, + "supports_image_edit_native": True, + "preferred_edit_backend": "sdxl", + }, + ) + + self.assertEqual(profile["profile_id"], "sdxl_edit") + self.assertEqual(profile["role"], "edit") + self.assertEqual(profile["backend_type"], "sdxl_native") + self.assertTrue(profile["supports_image_edit_native"]) + + def test_legacy_tags_upgrade_to_matching_profile(self): + model = FakeModel( + family="FLUX_DEV", + model_role="edit", + supports_inpaint_native=False, + supports_image_edit_native=True, + preferred_edit_backend="flux", + profile_id="", + backend_type="", + sampler_strategy="", + loader_strategy="", + model_name="legacy-flux-edit.safetensors", + ) + + profile = self.registry.resolve_model_profile(model=model, model_info=None) + + self.assertEqual(profile["profile_id"], "flux_edit") + self.assertEqual(profile["backend_type"], "flux_edit") + self.assertEqual(profile["sampler_strategy"], "flux_guided") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_pro_edit_bridge.py b/tests/test_pro_edit_bridge.py index 802592e..6d7d007 100644 --- a/tests/test_pro_edit_bridge.py +++ b/tests/test_pro_edit_bridge.py @@ -21,6 +21,10 @@ class TestProEditBridge(unittest.TestCase): supports_inpaint_native=True, supports_image_edit_native=False, preferred_edit_backend="sdxl", + profile_id="sdxl_inpaint", + backend_type="sdxl_native", + sampler_strategy="standard_k", + loader_strategy="sdxl_checkpoint", ) latent = {"samples": FakeLatentTensor((1, 4, 64, 64)), "source": "pro_edit_prepare_region"} @@ -48,6 +52,9 @@ class TestProEditBridge(unittest.TestCase): "supports_inpaint_native": True, "supports_image_edit_native": False, "preferred_edit_backend": "sdxl", + "profile_id": "sdxl_inpaint", + "backend_type": "sdxl_native", + "sampler_strategy": "standard_k", "recommended_denoise": 0.63, }, model_info=None, @@ -56,6 +63,8 @@ class TestProEditBridge(unittest.TestCase): sample_info = json.loads(sample_info_json) self.assertEqual(result_latent["source"], "pro_edit_prepare_region") self.assertEqual(sample_info["backend_name"], "sdxl") + self.assertEqual(sample_info["profile_id"], "sdxl_inpaint") + self.assertEqual(sample_info["sampler_strategy"], "standard_k") self.assertEqual(sample_info["denoise"], 0.63) self.assertEqual(sample_info["denoise_mode"], "auto_from_edit") @@ -66,6 +75,10 @@ class TestProEditBridge(unittest.TestCase): supports_inpaint_native=False, supports_image_edit_native=True, preferred_edit_backend="flux", + profile_id="flux_edit", + backend_type="flux_edit", + sampler_strategy="flux_guided", + loader_strategy="flux_split_or_bundle", ) latent = {"samples": FakeLatentTensor((1, 16, 64, 64)), "source": "pro_edit_prepare_region"} @@ -93,6 +106,9 @@ class TestProEditBridge(unittest.TestCase): "supports_inpaint_native": False, "supports_image_edit_native": True, "preferred_edit_backend": "flux", + "profile_id": "flux_edit", + "backend_type": "flux_edit", + "sampler_strategy": "flux_guided", "recommended_denoise": 0.8, }, model_info=None, @@ -100,6 +116,8 @@ class TestProEditBridge(unittest.TestCase): sample_info = json.loads(sample_info_json) self.assertEqual(sample_info["backend_name"], "flux") + self.assertEqual(sample_info["profile_id"], "flux_edit") + self.assertEqual(sample_info["sampler_strategy"], "flux_guided") self.assertEqual(sample_info["guidance"], 4.2) def test_bridge_manual_backend_mismatch_raises(self): @@ -109,6 +127,10 @@ class TestProEditBridge(unittest.TestCase): supports_inpaint_native=True, supports_image_edit_native=False, preferred_edit_backend="sdxl", + profile_id="sdxl_inpaint", + backend_type="sdxl_native", + sampler_strategy="standard_k", + loader_strategy="sdxl_checkpoint", ) latent = {"samples": FakeLatentTensor((1, 4, 64, 64)), "source": "pro_edit_prepare_region"} @@ -133,6 +155,41 @@ class TestProEditBridge(unittest.TestCase): model_info=None, ) + def test_bridge_rejects_unknown_sampler_strategy(self): + model = FakeModel( + family="FLUX_DEV", + model_role="edit", + supports_inpaint_native=False, + supports_image_edit_native=True, + preferred_edit_backend="flux", + profile_id="flux_edit", + backend_type="flux_edit", + sampler_strategy="mystery_strategy", + loader_strategy="flux_split_or_bundle", + ) + latent = {"samples": FakeLatentTensor((1, 16, 64, 64)), "source": "pro_edit_prepare_region"} + + with self.assertRaisesRegex(RuntimeError, "Unsupported sampler_strategy 'mystery_strategy'"): + self.node.sample( + model=model, + positive=make_conditioning("positive"), + negative=make_conditioning("negative"), + latent_image=latent, + backend_mode="auto", + quality_preset="Manual", + seed=9, + steps=12, + cfg=1.0, + sampler_name="euler", + scheduler="simple", + denoise=0.8, + denoise_mode="manual", + flux_guidance=4.2, + model_family="Auto", + edit_info=None, + model_info=None, + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_pro_edit_capabilities.py b/tests/test_pro_edit_capabilities.py index 38d22a4..842eec1 100644 --- a/tests/test_pro_edit_capabilities.py +++ b/tests/test_pro_edit_capabilities.py @@ -41,6 +41,38 @@ class TestProEditCapabilities(unittest.TestCase): self.assertTrue(info["supports_image_edit_native"]) self.assertEqual(info["preferred_edit_backend"], "flux") + def test_parse_model_info_infers_flux_kontext_as_edit_capable(self): + plugin = load_plugin_package() + model_info = import_plugin_submodule(plugin, "utils.model_info") + + info = model_info.parse_model_info( + { + "checkpoint_name": "demo-flux-kontext-dev.safetensors", + "family": "FLUX_DEV", + } + ) + + self.assertEqual(info["model_role"], "edit") + self.assertFalse(info["supports_inpaint_native"]) + self.assertTrue(info["supports_image_edit_native"]) + self.assertEqual(info["preferred_edit_backend"], "flux") + + def test_parse_model_info_exposes_profile_fields(self): + plugin = load_plugin_package() + model_info = import_plugin_submodule(plugin, "utils.model_info") + + info = model_info.parse_model_info( + { + "checkpoint_name": "demo-flux-kontext-dev.safetensors", + "family": "FLUX_DEV", + } + ) + + self.assertEqual(info["profile_id"], "flux_edit") + self.assertEqual(info["backend_type"], "flux_edit") + self.assertEqual(info["sampler_strategy"], "flux_guided") + self.assertEqual(info["loader_strategy"], "flux_split_or_bundle") + def test_parse_model_info_infers_capabilities_from_ckpt_name_alias(self): plugin = load_plugin_package() model_info = import_plugin_submodule(plugin, "utils.model_info") @@ -144,6 +176,29 @@ class TestProEditCapabilities(unittest.TestCase): self.assertTrue(capabilities["supports_image_edit_native"]) self.assertEqual(capabilities["preferred_edit_backend"], "flux") + def test_resolve_edit_capabilities_reads_profile_tags(self): + plugin = load_plugin_package() + model_info = import_plugin_submodule(plugin, "utils.model_info") + model = FakeModel( + family="SDXL", + model_role="edit", + supports_inpaint_native=True, + supports_image_edit_native=True, + preferred_edit_backend="sdxl", + profile_id="sdxl_edit", + backend_type="sdxl_native", + sampler_strategy="standard_k", + loader_strategy="sdxl_checkpoint", + model_name="demo-sdxl-edit.safetensors", + ) + + capabilities = model_info.resolve_edit_capabilities(model=model, model_info=None) + + self.assertEqual(capabilities["profile_id"], "sdxl_edit") + self.assertEqual(capabilities["backend_type"], "sdxl_native") + self.assertEqual(capabilities["sampler_strategy"], "standard_k") + self.assertEqual(capabilities["loader_strategy"], "sdxl_checkpoint") + def test_resolve_edit_capabilities_coerces_string_overrides_to_false(self): plugin = load_plugin_package() model_info = import_plugin_submodule(plugin, "utils.model_info") diff --git a/tests/test_pro_edit_docs.py b/tests/test_pro_edit_docs.py index 6ed7cd3..6e99519 100644 --- a/tests/test_pro_edit_docs.py +++ b/tests/test_pro_edit_docs.py @@ -16,6 +16,12 @@ class TestProEditDocs(unittest.TestCase): "Simple = lightweight masked latent resampling", "Pro = true image edit / inpaint pipeline", "backend_mode = auto | sdxl | flux", + "profile-driven routing", + "profile_id", + "backend_type", + "sampler_strategy", + "LLS Simple Checkpoint Loader writes the resolved model profile", + "FLUX base models should remain on the Simple workflow unless they resolve to flux_edit", "Load Image -> Load Mask or LLS Simple Mask Draw -> LLS Simple Checkpoint Loader -> LLS Simple Prompt Encode -> LLS Pro Image Edit Prepare -> LLS Pro KSampler Bridge -> VAE Decode -> LLS Pro Image Edit Finish -> Preview Image", "Adding new professional edit models", ): diff --git a/tests/test_pro_edit_helpers.py b/tests/test_pro_edit_helpers.py index 120f599..0df1957 100644 --- a/tests/test_pro_edit_helpers.py +++ b/tests/test_pro_edit_helpers.py @@ -316,6 +316,10 @@ class FakeModel: supports_image_edit_native=False, preferred_edit_backend=None, model_name="demo-model.safetensors", + profile_id="", + backend_type="", + sampler_strategy="", + loader_strategy="", ): self._lls_family = family self._lls_model_role = model_role @@ -323,6 +327,11 @@ class FakeModel: self._lls_supports_image_edit_native = supports_image_edit_native self._lls_preferred_edit_backend = preferred_edit_backend self._lls_model_name = model_name + self._lls_checkpoint_name = model_name + self._lls_profile_id = profile_id + self._lls_backend_type = backend_type + self._lls_sampler_strategy = sampler_strategy + self._lls_loader_strategy = loader_strategy def make_conditioning(label): diff --git a/tests/test_pro_edit_prepare_flux.py b/tests/test_pro_edit_prepare_flux.py index d4805d1..b5dcf14 100644 --- a/tests/test_pro_edit_prepare_flux.py +++ b/tests/test_pro_edit_prepare_flux.py @@ -36,6 +36,10 @@ class TestProEditPrepareFlux(unittest.TestCase): supports_inpaint_native=False, supports_image_edit_native=True, preferred_edit_backend="flux", + profile_id="flux_edit", + backend_type="flux_edit", + sampler_strategy="flux_guided", + loader_strategy="flux_split_or_bundle", ) latent, work_image, work_mask, edit_info, recommended, positive, negative = self.node.prepare( @@ -66,7 +70,11 @@ class TestProEditPrepareFlux(unittest.TestCase): ) self.assertEqual(edit_info["backend_name"], "flux") + self.assertEqual(edit_info["routing_reason"], "profile.backend_type") self.assertEqual(edit_info["model_family"], "FLUX_DEV") + self.assertEqual(edit_info["profile_id"], "flux_edit") + self.assertEqual(edit_info["backend_type"], "flux_edit") + self.assertEqual(edit_info["sampler_strategy"], "flux_guided") self.assertTrue(edit_info["supports_image_edit_native"]) self.assertIn("concat_latent_image", positive[0][1]) self.assertIn("concat_mask", positive[0][1]) @@ -84,6 +92,10 @@ class TestProEditPrepareFlux(unittest.TestCase): supports_inpaint_native=False, supports_image_edit_native=True, preferred_edit_backend="flux", + profile_id="flux_edit", + backend_type="flux_edit", + sampler_strategy="flux_guided", + loader_strategy="flux_split_or_bundle", ) latent, work_image, work_mask, edit_info, recommended, positive, negative = self.node.prepare( @@ -116,6 +128,9 @@ class TestProEditPrepareFlux(unittest.TestCase): self.assertEqual(edit_info["edit_scope"], "canvas") self.assertEqual(edit_info["original_box_in_canvas"], [64, 32, 832, 800]) self.assertEqual(edit_info["work_size"], [896, 832]) + self.assertEqual(edit_info["profile_id"], "flux_edit") + self.assertEqual(edit_info["backend_type"], "flux_edit") + self.assertEqual(edit_info["sampler_strategy"], "flux_guided") self.assertEqual(positive[0][1]["edit_backend"], "flux") self.assertGreater(recommended, 0.0) diff --git a/tests/test_pro_edit_prepare_sdxl.py b/tests/test_pro_edit_prepare_sdxl.py index 682f879..196f920 100644 --- a/tests/test_pro_edit_prepare_sdxl.py +++ b/tests/test_pro_edit_prepare_sdxl.py @@ -36,6 +36,10 @@ class TestProEditPrepareSDXL(unittest.TestCase): supports_inpaint_native=True, supports_image_edit_native=False, preferred_edit_backend="sdxl", + profile_id="sdxl_inpaint", + backend_type="sdxl_native", + sampler_strategy="standard_k", + loader_strategy="sdxl_checkpoint", ) latent, work_image, work_mask, edit_info, recommended, positive, negative = self.node.prepare( @@ -66,8 +70,11 @@ class TestProEditPrepareSDXL(unittest.TestCase): ) self.assertEqual(edit_info["backend_name"], "sdxl") - self.assertEqual(edit_info["routing_reason"], "model.preferred_edit_backend") + self.assertEqual(edit_info["routing_reason"], "profile.backend_type") self.assertEqual(edit_info["edit_scope"], "region") + self.assertEqual(edit_info["profile_id"], "sdxl_inpaint") + self.assertEqual(edit_info["backend_type"], "sdxl_native") + self.assertEqual(edit_info["sampler_strategy"], "standard_k") self.assertEqual(latent["source"], "pro_edit_prepare_region") self.assertIn("concat_latent_image", positive[0][1]) self.assertIn("concat_mask", positive[0][1]) @@ -87,6 +94,10 @@ class TestProEditPrepareSDXL(unittest.TestCase): supports_inpaint_native=True, supports_image_edit_native=True, preferred_edit_backend="sdxl", + profile_id="sdxl_edit", + backend_type="sdxl_native", + sampler_strategy="standard_k", + loader_strategy="sdxl_checkpoint", ) latent, work_image, work_mask, edit_info, recommended, positive, negative = self.node.prepare( @@ -120,6 +131,9 @@ class TestProEditPrepareSDXL(unittest.TestCase): self.assertIsInstance(edit_info["crop_box"], list) self.assertEqual(edit_info["work_size"], [work_image.shape[2], work_image.shape[1]]) self.assertEqual(edit_info["backend_name"], "sdxl") + self.assertEqual(edit_info["profile_id"], "sdxl_edit") + self.assertEqual(edit_info["backend_type"], "sdxl_native") + self.assertEqual(edit_info["sampler_strategy"], "standard_k") self.assertEqual(positive[0][1]["edit_backend"], "sdxl") self.assertEqual(negative[0][1]["edit_backend"], "sdxl") self.assertGreater(recommended, 0.0) diff --git a/tests/test_pro_edit_profile_prepare.py b/tests/test_pro_edit_profile_prepare.py new file mode 100644 index 0000000..537751b --- /dev/null +++ b/tests/test_pro_edit_profile_prepare.py @@ -0,0 +1,58 @@ +import unittest + +try: + from .test_pro_edit_helpers import FakeMask, FakeModel, FakeTensor, FakeVAE, load_plugin_package, make_conditioning +except ImportError: + from test_pro_edit_helpers import FakeMask, FakeModel, FakeTensor, FakeVAE, load_plugin_package, make_conditioning + + +class TestProEditProfilePrepare(unittest.TestCase): + def test_prepare_rejects_base_profile(self): + plugin = load_plugin_package() + node = plugin.NODE_CLASS_MAPPINGS["LLSProImageEditPrepare"]() + image = FakeTensor((1, 64, 64, 3), label="image") + mask = FakeMask((1, 64, 64), mask_bbox=(8, 8, 32, 32), mask_area_ratio=0.2) + vae = FakeVAE(latent_channels=4, downscale_ratio=8) + model = FakeModel( + family="SDXL", + model_role="base", + supports_inpaint_native=False, + supports_image_edit_native=False, + preferred_edit_backend=None, + profile_id="sdxl_base", + backend_type="none", + sampler_strategy="standard_k", + loader_strategy="sdxl_checkpoint", + ) + + with self.assertRaisesRegex(RuntimeError, "Pro image edit is not available for profile 'sdxl_base'"): + node.prepare( + image=image, + mask=mask, + vae=vae, + positive=make_conditioning("positive"), + negative=make_conditioning("negative"), + backend_mode="auto", + edit_scope="region", + mask_grow=0, + mask_blur=0.0, + mask_threshold=0.5, + invert_mask=False, + crop_context=32, + crop_context_factor=1.0, + min_size=128, + max_size=512, + resize_mode="fit", + expand_left=0, + expand_right=0, + expand_top=0, + expand_bottom=0, + canvas_fill="edge", + auto_recommend="enabled", + model=model, + model_info=None, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_pro_edit_registry.py b/tests/test_pro_edit_registry.py index ba21475..61c3078 100644 --- a/tests/test_pro_edit_registry.py +++ b/tests/test_pro_edit_registry.py @@ -11,37 +11,48 @@ class TestProEditRegistry(unittest.TestCase): self.plugin = load_plugin_package() self.registry = import_plugin_submodule(self.plugin, "pro_edit.backends.registry") - def test_auto_routes_sdxl_inpaint_model(self): + def test_auto_routes_sdxl_profile_by_backend_type(self): model = FakeModel( family="SDXL", - model_role="inpaint", + model_role="edit", supports_inpaint_native=True, - supports_image_edit_native=False, + supports_image_edit_native=True, preferred_edit_backend="sdxl", + profile_id="sdxl_edit", + backend_type="sdxl_native", + sampler_strategy="standard_k", + loader_strategy="sdxl_checkpoint", ) backend, routing = self.registry.resolve_backend("auto", model=model) self.assertEqual(backend.backend_name, "sdxl") self.assertEqual(routing.backend_name, "sdxl") - self.assertEqual(routing.routing_reason, "model.preferred_edit_backend") + self.assertEqual(routing.routing_reason, "profile.backend_type") + self.assertEqual(routing.profile["profile_id"], "sdxl_edit") + self.assertEqual(routing.profile["backend_type"], "sdxl_native") - def test_auto_routes_flux_edit_model(self): + def test_auto_routes_flux_profile_by_backend_type(self): model = FakeModel( family="FLUX_DEV", model_role="edit", supports_inpaint_native=False, supports_image_edit_native=True, preferred_edit_backend="flux", + profile_id="flux_edit", + backend_type="flux_edit", + sampler_strategy="flux_guided", + loader_strategy="flux_split_or_bundle", ) backend, routing = self.registry.resolve_backend("auto", model=model) self.assertEqual(backend.backend_name, "flux") - self.assertEqual(routing.routing_reason, "model.preferred_edit_backend") + self.assertEqual(routing.routing_reason, "profile.backend_type") self.assertEqual(routing.capabilities["model_family"], "FLUX_DEV") + self.assertEqual(routing.profile["profile_id"], "flux_edit") - def test_auto_reuses_backend_name_from_edit_info(self): + def test_auto_upgrades_legacy_edit_info_to_profile_route(self): backend, routing = self.registry.resolve_backend( "auto", model=None, @@ -56,7 +67,8 @@ class TestProEditRegistry(unittest.TestCase): ) self.assertEqual(backend.backend_name, "sdxl") - self.assertEqual(routing.routing_reason, "edit_info.backend_name") + self.assertEqual(routing.routing_reason, "profile.backend_type") + self.assertEqual(routing.profile["profile_id"], "sdxl_inpaint") def test_manual_flux_override_rejects_sdxl_only_model(self): model = FakeModel( @@ -65,21 +77,29 @@ class TestProEditRegistry(unittest.TestCase): supports_inpaint_native=True, supports_image_edit_native=False, preferred_edit_backend="sdxl", + profile_id="sdxl_inpaint", + backend_type="sdxl_native", + sampler_strategy="standard_k", + loader_strategy="sdxl_checkpoint", ) - with self.assertRaisesRegex(RuntimeError, "backend 'flux' is incompatible"): + with self.assertRaisesRegex(RuntimeError, "backend 'flux' is incompatible with profile"): self.registry.resolve_backend("flux", model=model) - def test_auto_without_matching_backend_raises_clear_error(self): + def test_auto_rejects_base_profile_for_pro_chain(self): model = FakeModel( - family="SD1.5", + family="SDXL", model_role="base", supports_inpaint_native=False, supports_image_edit_native=False, preferred_edit_backend=None, + profile_id="sdxl_base", + backend_type="none", + sampler_strategy="standard_k", + loader_strategy="sdxl_checkpoint", ) - with self.assertRaisesRegex(RuntimeError, "No professional edit backend matched"): + with self.assertRaisesRegex(RuntimeError, "Pro image edit is not available for profile 'sdxl_base'"): self.registry.resolve_backend("auto", model=model) diff --git a/utils/model_info.py b/utils/model_info.py index ad0b602..acaf493 100644 --- a/utils/model_info.py +++ b/utils/model_info.py @@ -195,6 +195,7 @@ SAMPLING_PRESETS: dict[str, dict[str, dict[str, float | int | None]]] = { _LEGACY_MODEL_INFO_PATTERN = re.compile(r"(\w+)=([^|]+)") _ROLE_KEYWORDS = ( ("inpaint", "inpaint"), + ("kontext", "edit"), ("img2img", "edit"), ("imageedit", "edit"), ("image-edit", "edit"), @@ -336,6 +337,14 @@ def infer_edit_capabilities(model_name: str | None, family: str | None = None) - } +def _load_profile_resolver(): + try: + from ..model_profiles.registry import resolve_model_profile + except ImportError: + from model_profiles.registry import resolve_model_profile + return resolve_model_profile + + def parse_model_info(model_info: dict[str, Any] | str | None) -> dict[str, Any]: raw = parse_jsonish_info(model_info) raw_model_name = _get_model_name_alias_value(raw) @@ -348,25 +357,26 @@ def parse_model_info(model_info: dict[str, Any] | str | None) -> dict[str, Any]: ) ) defaults = get_family_defaults(family) - capability_defaults = infer_edit_capabilities( - raw_model_name, - family, + resolve_model_profile = _load_profile_resolver() + profile = resolve_model_profile( + model=None, + model_info=raw, + checkpoint_name=raw_model_name, + family=family, ) info: dict[str, Any] = {**defaults} info.update(raw) - info["family"] = family - info["model_role"] = str(raw.get("model_role", capability_defaults["model_role"])) - info["supports_inpaint_native"] = _coerce_bool( - raw.get("supports_inpaint_native", capability_defaults["supports_inpaint_native"]) - ) - info["supports_image_edit_native"] = _coerce_bool( - raw.get("supports_image_edit_native", capability_defaults["supports_image_edit_native"]) - ) - info["preferred_edit_backend"] = raw.get( - "preferred_edit_backend", - capability_defaults["preferred_edit_backend"], - ) + info["family"] = profile["family"] + info["model_family"] = profile["family"] + info["model_role"] = profile["role"] + info["profile_id"] = profile["profile_id"] + info["backend_type"] = profile["backend_type"] + info["sampler_strategy"] = profile["sampler_strategy"] + info["loader_strategy"] = profile["loader_strategy"] + info["supports_inpaint_native"] = _coerce_bool(profile["supports_inpaint_native"]) + info["supports_image_edit_native"] = _coerce_bool(profile["supports_image_edit_native"]) + info["preferred_edit_backend"] = profile["preferred_edit_backend"] info["text_encoder_type"] = raw.get("text_encoder_type", defaults["text_encoder_type"]) required_text_encoders = raw.get("required_text_encoders", defaults["required_text_encoders"]) if isinstance(required_text_encoders, str): @@ -402,7 +412,6 @@ def parse_model_info(model_info: dict[str, Any] | str | None) -> dict[str, Any]: info.setdefault("checkpoint_name", raw_model_name or "") info.setdefault("model_name", info["checkpoint_name"]) - info["model_family"] = family info.setdefault("vae_source", raw.get("vae_source", "auto")) info.setdefault("text_encoder_source", raw.get("text_encoder_source", "auto")) info.setdefault("load_mode", raw.get("load_mode", "simple")) @@ -516,12 +525,12 @@ def get_lls_attr(obj, name: str, default: Any = None) -> Any: def resolve_edit_capabilities(model=None, model_info: dict[str, Any] | str | None = None) -> dict[str, Any]: raw = parse_jsonish_info(model_info) info = parse_model_info(model_info) - raw_model_name = _get_model_name_alias_value(raw) model_name = str( - raw_model_name + _get_model_name_alias_value(raw) + or get_lls_attr(model, "checkpoint_name", None) + or get_lls_attr(model, "model_name", "") or info.get("checkpoint_name") or info.get("model_name") - or get_lls_attr(model, "model_name", "") or "" ) family = canonicalize_family( @@ -532,41 +541,24 @@ def resolve_edit_capabilities(model=None, model_info: dict[str, Any] | str | Non or info.get("model_family") or infer_family_from_model(model) ) - inferred = infer_edit_capabilities(model_name, family) + resolve_model_profile = _load_profile_resolver() + profile = resolve_model_profile( + model=model, + model_info=model_info, + checkpoint_name=model_name, + family=family, + ) return { - "model_family": family, + "model_family": profile["family"], "model_name": model_name, - "model_role": str( - raw.get("model_role") - or get_lls_attr(model, "model_role", None) - or info.get("model_role") - or inferred["model_role"] - ), - "supports_inpaint_native": _coerce_bool( - raw.get( - "supports_inpaint_native", - get_lls_attr( - model, - "supports_inpaint_native", - info.get("supports_inpaint_native", inferred["supports_inpaint_native"]), - ), - ) - ), - "supports_image_edit_native": _coerce_bool( - raw.get( - "supports_image_edit_native", - get_lls_attr( - model, - "supports_image_edit_native", - info.get("supports_image_edit_native", inferred["supports_image_edit_native"]), - ), - ) - ), - "preferred_edit_backend": ( - raw.get("preferred_edit_backend") - or get_lls_attr(model, "preferred_edit_backend", inferred["preferred_edit_backend"]) - or info.get("preferred_edit_backend") - ), + "model_role": profile["role"], + "profile_id": profile["profile_id"], + "backend_type": profile["backend_type"], + "sampler_strategy": profile["sampler_strategy"], + "loader_strategy": profile["loader_strategy"], + "supports_inpaint_native": _coerce_bool(profile["supports_inpaint_native"]), + "supports_image_edit_native": _coerce_bool(profile["supports_image_edit_native"]), + "preferred_edit_backend": profile["preferred_edit_backend"], }