150 lines
4.8 KiB
Python
150 lines
4.8 KiB
Python
|
|
|
|
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"
|