ComfyUI-Lightx2vWrapper/model_utils.py at main · ModelTC/ComfyUI-Lightx2vWrapper · GitHub

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
# coding: utf-8

import json
from pathlib import Path
from typing import Dict, List

import folder_paths


def get_model_base_path() -> Path:
    models_base = folder_paths.models_dir
    lightx2v_path = Path(models_base) / "lightx2v"
    if not lightx2v_path.exists():
        lightx2v_path.mkdir(parents=True, exist_ok=True)
    return lightx2v_path


def scan_models() -> List[str]:
    models = []
    base_path = get_model_base_path()

    if base_path.exists():
        for item in base_path.iterdir():
            if item.is_dir() and item.name != "loras":
                models.append(item.name)

    models.sort()

    return ["None"] + models if models else ["None"]


def support_model_cls_list() -> List[str]:
    return [
        "wan2.1",
        "wan2.1_distill",
        "wan2.1_vace",
        "cogvideox",
        "seko_talk",
        "wan2.2_moe",
        "wan2.2",
        "wan2.2_moe_audio",
        "wan2.2_audio",
        "wan2.2_moe_distill",
        "qwen_image",
    ]


def get_loras_models(model_path: Path) -> List[str]:
    loras = []
    if model_path.exists():
        for item in model_path.iterdir():
            if item.is_file():
                if item.suffix.lower() in [".safetensors", ".pt", ".pth", ".ckpt"]:
                    loras.append(item.name)
    return loras


def scan_loras() -> List[str]:
    base_path = get_model_base_path()
    loras_path = base_path / "loras"
    loras = get_loras_models(loras_path)

    models_base = folder_paths.models_dir
    loras_path = Path(models_base) / "loras"

    loras2 = get_loras_models(loras_path)
    loras.extend(loras2)

    loras.sort()

    return ["None"] + loras if loras else ["None"]


def get_model_full_path(model_name: str) -> str:
    if model_name == "None" or not model_name:
        return ""

    base_path = get_model_base_path()
    model_path = base_path / model_name

    if model_path.exists():
        return str(model_path)
    return ""


def get_lora_full_path(lora_name: str) -> str:
    if lora_name == "None" or not lora_name:
        return ""

    base_path = get_model_base_path()
    lora_path = base_path / "loras" / lora_name

    if lora_path.exists():
        return str(lora_path)

    models_base = folder_paths.models_dir
    loras_path = Path(models_base) / "loras" / lora_name

    if loras_path.exists():
        return str(loras_path)

    return ""


def get_model_info(model_name: str) -> Dict:
    if model_name == "None" or not model_name:
        return {}

    base_path = get_model_base_path()
    config_path = base_path / model_name / "config.json"

    if config_path.exists():
        with open(config_path, "r") as f:
            return json.load(f)
    return {}