File size: 5,742 Bytes
8a28a8d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
import os
from copy import deepcopy
from .common import _load_yaml, _CHAIN_FEATURES_PATH, _YAML_DIR
from .error_schema import make_not_found_error


def _build_feature_entry(chain_name: str, chain_data: dict, include_schema: bool = False) -> dict:
    entry = {
        "feature_name": chain_name,
        "chains": chain_data.get("chains", chain_name),
        "display_name": chain_data.get("display_name", chain_name),
        "description": chain_data.get("description", ""),
        "supported_tasks": chain_data.get("supported_tasks", []),
        "max_count": chain_data.get("max_count", 1),
        "usage_guideline": chain_data.get("usage_guideline", ""),
    }
    if include_schema:
        schema = deepcopy(chain_data.get("parameters_schema", {}))
        if chain_name in ("krea2_controlnet", "diffsynth_controlnet", "controlnet", "anima_controlnet_lllite"):
            config_key = (
                "Krea2_ControlNet" if chain_name == "krea2_controlnet"
                else "DiffSynth_ControlNet" if chain_name == "diffsynth_controlnet"
                else "Anima_ControlNet_Lllite" if chain_name == "anima_controlnet_lllite"
                else "ControlNet"
            )
            yaml_filename = f"{chain_name}_models.yaml"
            model_path = os.path.join(_YAML_DIR, yaml_filename)
            raw_models = _load_yaml(model_path).get(config_key, [])
            models_list = []
            if isinstance(raw_models, dict):
                for val in raw_models.values():
                    if isinstance(val, list):
                        models_list.extend(val)
                    elif isinstance(val, dict):
                        models_list.append(val)
            elif isinstance(raw_models, list):
                models_list = raw_models

            types_set = set()
            for m in models_list:
                t_val = m.get("Type", [])
                if isinstance(t_val, list):
                    types_set.update(t_val)
                elif isinstance(t_val, str):
                    types_set.add(t_val)
            types = sorted(list(types_set))
            series = sorted(list(set(m.get("Series") for m in models_list if m.get("Series"))))
            if "properties" in schema:
                if "type" in schema["properties"] and types:
                    schema["properties"]["type"]["enum"] = types
                if "series" in schema["properties"] and series:
                    schema["properties"]["series"]["enum"] = series
            if chain_name == "controlnet" and isinstance(raw_models, dict):
                schema["architectures"] = raw_models
        elif chain_name == "ipadapter":
            from .common import _get_ipadapter_presets_by_arch
            presets_by_arch = _get_ipadapter_presets_by_arch()
            schema["presets_by_architecture"] = presets_by_arch
            all_presets = sorted(list(set(presets_by_arch.get("SD1.5", []) + presets_by_arch.get("SDXL", []))))
            if "properties" in schema and "preset" in schema["properties"]:
                schema["properties"]["preset"]["enum"] = all_presets
        entry["parameters_schema"] = schema
    return entry


def handle_get_feature_list(
    feature_name: str | list[str] = "",
    include_schema_on_empty: bool = False,
) -> list | dict:
    """
    Dynamically load supported advanced features from chain_features.yaml.

    - If feature_name is empty: returns summaries by default; callers can set
      include_schema_on_empty for the legacy Fluxus full-schema contract.
    - If feature_name is specified (single feature name, comma-separated string, or list of strings):
      returns complete feature details INCLUDING parameters_schema for the requested feature(s).
    """
    chain_features = _load_yaml(_CHAIN_FEATURES_PATH)

    targets = []
    is_single_string_query = False

    if isinstance(feature_name, list):
        targets = [str(x).strip() for x in feature_name if str(x).strip()]
    elif isinstance(feature_name, str) and feature_name.strip():
        raw_str = feature_name.strip()
        parts = [x.strip() for x in raw_str.split(",") if x.strip()]
        targets = parts
        if len(parts) == 1 and "," not in raw_str:
            is_single_string_query = True

    # Empty discovery can serve both contracts: compact ImageGen summaries or
    # Fluxus-compatible full schemas.
    if not targets:
        return [
            _build_feature_entry(
                name, data, include_schema=include_schema_on_empty
            )
            for name, data in chain_features.items()
        ]

    # Helper function to resolve feature target by key or chains alias
    def _resolve_target(target_name: str) -> str | None:
        if target_name in chain_features:
            return target_name
        for feat_key, feat_data in chain_features.items():
            feat_chains = feat_data.get("chains")
            if isinstance(feat_chains, str) and feat_chains == target_name:
                return feat_key
            elif isinstance(feat_chains, list) and target_name in feat_chains:
                return feat_key
        return None

    resolved_targets = []
    # Case 2: Specific feature(s) requested -> validate existence
    for target in targets:
        resolved = _resolve_target(target)
        if not resolved:
            return make_not_found_error("feature_name", target)
        resolved_targets.append(resolved)

    # Case 3: Return full info including parameters_schema
    results = [
        _build_feature_entry(target, chain_features[target], include_schema=True)
        for target in resolved_targets
    ]

    if is_single_string_query and len(results) == 1:
        return results[0]

    return results