Hugging Face's logo

Spaces:
JasonSmithSO
/
FooocusEnhanced
Configuration error

App Files Community
FooocusEnhanced
File size: 4,941 Bytes 
8866644
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
 


import json
import os
import sys
from nodes import MAX_RESOLUTION
import comfy.utils
import shutil
import comfy.samplers
import folder_paths


WEB_DIRECTORY = "./web"

AUTHOR_NAME = u"MinusZone"
CATEGORY_NAME = f"{AUTHOR_NAME} - FluxExt"


import importlib

NODE_CLASS_MAPPINGS = {
}


NODE_DISPLAY_NAME_MAPPINGS = {
}

from . import mz_fluxext_core
import importlib


class MZ_Flux1PartialLoad_Patch:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
            "model": ("MODEL", ),
            "double_blocks_cuda_size": ("INT", {"min": 0, "max": 16, "default": 7}),
            "single_blocks_cuda_size": ("INT", {"min": 0, "max": 37, "default": 7}),
        }}
    RETURN_TYPES = ("MODEL",)
    FUNCTION = "load_unet"

    CATEGORY = f"{CATEGORY_NAME}"

    def load_unet(self, **kwargs):
        from . import mz_fluxext_core
        importlib.reload(mz_fluxext_core)
        return mz_fluxext_core.Flux1PartialLoad_Patch(kwargs)


NODE_CLASS_MAPPINGS["MZ_Flux1PartialLoad_Patch"] = MZ_Flux1PartialLoad_Patch
NODE_DISPLAY_NAME_MAPPINGS["MZ_Flux1PartialLoad_Patch"] = f"{AUTHOR_NAME} - Flux1PartialLoad_Patch"

import nodes


class MZ_Flux1CheckpointLoaderNF4_cpuDynOffload:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
            "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
            "double_blocks_cuda_size": ("INT", {"min": 0, "max": 16, "default": 7}),
            "single_blocks_cuda_size": ("INT", {"min": 0, "max": 37, "default": 7}),
        }}

    RETURN_TYPES = ("MODEL", "CLIP", "VAE")
    FUNCTION = "load_checkpoint"

    CATEGORY = f"{CATEGORY_NAME}"

    def load_checkpoint(self, ckpt_name, **kwargs):
        CheckpointLoaderNF4 = nodes.NODE_CLASS_MAPPINGS.get(
            "CheckpointLoaderNF4", None)
        if CheckpointLoaderNF4 is None:
            # 必须安装 https://github.com/comfyanonymous/ComfyUI_bitsandbytes_NF4
            raise Exception(
                "Please install comfyanonymous/ComfyUI_bitsandbytes_NF4 to use this node.")

        model, clip, vae = CheckpointLoaderNF4().load_checkpoint(ckpt_name)
        return mz_fluxext_core.Flux1PartialLoad_Patch({
            "model": model,
            "double_blocks_cuda_size": kwargs.get("double_blocks_cuda_size", 7),
            "single_blocks_cuda_size": kwargs.get("single_blocks_cuda_size", 7),
        })[0], clip, vae


NODE_CLASS_MAPPINGS["MZ_Flux1CheckpointLoaderNF4_cpuDynOffload"] = MZ_Flux1CheckpointLoaderNF4_cpuDynOffload
NODE_DISPLAY_NAME_MAPPINGS[
    "MZ_Flux1CheckpointLoaderNF4_cpuDynOffload"] = f"{AUTHOR_NAME} - Flux1CheckpointLoaderNF4_cpuDynOffload"


class MZ_Flux1CheckpointLoader_cpuDynOffload:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
            "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
            "double_blocks_cuda_size": ("INT", {"min": 0, "max": 16, "default": 7}),
            "single_blocks_cuda_size": ("INT", {"min": 0, "max": 37, "default": 7}),
        }}

    RETURN_TYPES = ("MODEL", "CLIP", "VAE")
    FUNCTION = "load_checkpoint"

    CATEGORY = f"{CATEGORY_NAME}"

    def load_checkpoint(self, ckpt_name, **kwargs):
        model, clip, vae = nodes.CheckpointLoaderSimple().load_checkpoint(
            ckpt_name=ckpt_name)
        return mz_fluxext_core.Flux1PartialLoad_Patch({
            "model": model,
            "double_blocks_cuda_size": kwargs.get("double_blocks_cuda_size", 7),
            "single_blocks_cuda_size": kwargs.get("single_blocks_cuda_size", 7),
        })[0], clip, vae


NODE_CLASS_MAPPINGS["MZ_Flux1CheckpointLoader_cpuDynOffload"] = MZ_Flux1CheckpointLoader_cpuDynOffload
NODE_DISPLAY_NAME_MAPPINGS[
    "MZ_Flux1CheckpointLoader_cpuDynOffload"] = f"{AUTHOR_NAME} - Flux1CheckpointLoader_cpuDynOffload"


class MZ_Flux1UnetLoader_cpuDynOffload:
    @classmethod
    def INPUT_TYPES(s):
        args = nodes.UNETLoader().INPUT_TYPES()
        args["required"]["double_blocks_cuda_size"] = (
            "INT", {"min": 0, "max": 16, "default": 7})
        args["required"]["single_blocks_cuda_size"] = (
            "INT", {"min": 0, "max": 37, "default": 7})
        return args

    RETURN_TYPES = ("MODEL",)
    FUNCTION = "load_unet"

    CATEGORY = f"{CATEGORY_NAME}"

    def load_unet(self, **kwargs):
        model = nodes.UNETLoader().load_unet(
            **{k: v for k, v in kwargs.items() if k != "double_blocks_cuda_size" and k != "single_blocks_cuda_size"})[0]

        return mz_fluxext_core.Flux1PartialLoad_Patch({
            "model": model,
            "double_blocks_cuda_size": kwargs.get("double_blocks_cuda_size", 7),
            "single_blocks_cuda_size": kwargs.get("single_blocks_cuda_size", 7),
        })


NODE_CLASS_MAPPINGS["MZ_Flux1UnetLoader_cpuDynOffload"] = MZ_Flux1UnetLoader_cpuDynOffload
NODE_DISPLAY_NAME_MAPPINGS[
    "MZ_Flux1UnetLoader_cpuDynOffload"] = f"{AUTHOR_NAME} - Flux1UnetLoader_cpuDynOffload"