onediff/onediff_comfy_nodes/__init__.py at 0f491f53528f25c6a0503dc94bf923d388772ef5 · siliconflow/onediff · 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
"""OneFlow ComfyUI Speedup Module"""
from ._config import _USE_UNET_INT8
from ._nodes import (
    ModelSpeedup,
    ModelGraphLoader,
    ModelGraphSaver,
    VaeSpeedup,
    VaeGraphLoader,
    VaeGraphSaver,
    ControlNetSpeedup,
    ControlNetGraphLoader,
    ControlNetGraphSaver,
    SVDSpeedup,
    ModuleDeepCacheSpeedup,
    OneDiffCheckpointLoaderSimple,
)
from ._compare_node import CompareModel, ShowImageDiff


NODE_CLASS_MAPPINGS = {
    "ModelSpeedup": ModelSpeedup,
    "CompareModel": CompareModel,
    "ShowImageDiff": ShowImageDiff,
    "ModelGraphLoader": ModelGraphLoader,
    "ModelGraphSaver": ModelGraphSaver,
    "VaeSpeedup": VaeSpeedup,
    "VaeGraphSaver": VaeGraphSaver,
    "VaeGraphLoader": VaeGraphLoader,
    "ControlNetSpeedup": ControlNetSpeedup,
    "ControlNetGraphLoader": ControlNetGraphLoader,
    "ControlNetGraphSaver": ControlNetGraphSaver,
    "SVDSpeedup": SVDSpeedup,
    "ModuleDeepCacheSpeedup": ModuleDeepCacheSpeedup,
    "OneDiffCheckpointLoaderSimple": OneDiffCheckpointLoaderSimple,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "ModelSpeedup": "Model Speedup",
    "CompareModel": "Model Weight Comparator",
    "ShowImageDiff": "Image Distinction Scanner",
    "ModelGraphLoader": "Model Graph Loader",
    "ModelGraphSaver": "Model Graph Saver",
    "VaeSpeedup": "VAE Speedup",
    "VaeGraphLoader": "VAE Graph Loader",
    "VaeGraphSaver": "VAE Graph Saver",
    "ControlNetSpeedup": "ControlNet Speedup",
    "ControlNetGraphLoader": "ControlNet Graph Loader",
    "ControlNetGraphSaver": "ControlNet Graph Saver",
    "SVDSpeedup": "SVD Speedup",
    "ModuleDeepCacheSpeedup": "Model DeepCache Speedup",
    "OneDiffCheckpointLoaderSimple": "Load Checkpoint - OneDiff",
}

if _USE_UNET_INT8:
    from ._nodes import UNETLoaderInt8, Quant8Model

    NODE_CLASS_MAPPINGS.update(
        {"UNETLoaderInt8": UNETLoaderInt8, "Quant8Model": Quant8Model}
    )
    NODE_DISPLAY_NAME_MAPPINGS.update(
        {
            "UNETLoaderInt8": "UNET Loader Int8",
            "Quant8Model": "Model Quantization(int8)",
        }
    )