File size: 1,583 Bytes
f59a2e3
 
219ad6d
f59a2e3
 
 
 
 
 
219ad6d
f59a2e3
 
 
 
 
 
 
 
 
 
 
 
 
 
219ad6d
f59a2e3
219ad6d
 
 
 
f59a2e3
219ad6d
f59a2e3
 
219ad6d
 
 
 
 
f59a2e3
 
219ad6d
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
"""

MCP Tool: get_model_architecture_list

Get all supported model architectures, their corresponding default resolutions, and available aspect ratios.

"""

from .common import _load_yaml, _MODEL_ARCHITECTURES_PATH, _CONSTANTS_PATH


def handle_get_model_architecture_list() -> list:
    """Dynamically load all supported model architectures from model_architectures.yaml along with available aspect ratios and resolutions."""
    arch_config = _load_yaml(_MODEL_ARCHITECTURES_PATH)
    constants = _load_yaml(_CONSTANTS_PATH)
    resolution_map = constants.get("RESOLUTION_MAP", {})
    architectures = arch_config.get("architectures", {})
    architecture_order = arch_config.get("architecture_order", list(architectures.keys()))

    result = []
    for arch_name in architecture_order:
        if arch_name not in architectures:
            continue
        arch_data = architectures[arch_name]
        model_type = arch_data.get("model_type", arch_name.lower())

        default_res = [1024, 1024]
        resolutions_dict = {}
        if model_type in resolution_map:
            resolutions_dict = resolution_map[model_type]
            if resolutions_dict:
                first_key = next(iter(resolutions_dict))
                default_res = resolutions_dict[first_key]

        entry = {
            "model_architecture": arch_name,
            "default_resolution": default_res,
        }
        if resolutions_dict:
            entry["available_resolutions"] = resolutions_dict

        result.append(entry)

    return result