From b55a341fc5d6840c347f9726c4e046dfcc770bbf Mon Sep 17 00:00:00 2001 From: wailovet Date: Thu, 15 Aug 2024 01:29:37 +0800 Subject: [PATCH] 'update' --- __init__.py | 50 ++++++++++++++++++ copytoww.bat | 1 + exclude.txt | 2 + mz_fluxext_core.py | 124 +++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 177 insertions(+) create mode 100644 __init__.py create mode 100644 copytoww.bat create mode 100644 exclude.txt create mode 100644 mz_fluxext_core.py diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..520ec92 --- /dev/null +++ b/__init__.py @@ -0,0 +1,50 @@ + + +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_Flux1VRAM_MT_Patch: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ("MODEL", ) + }} + 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.MZ_Flux1VRAM_MT_Patch_call(kwargs) + + +NODE_CLASS_MAPPINGS["MZ_Flux1VRAM_MT_Patch"] = MZ_Flux1VRAM_MT_Patch +NODE_DISPLAY_NAME_MAPPINGS["MZ_Flux1VRAM_MT_Patch"] = f"{AUTHOR_NAME} - Flux1VRAM_MT_Patch" diff --git a/copytoww.bat b/copytoww.bat new file mode 100644 index 0000000..61c81ec --- /dev/null +++ b/copytoww.bat @@ -0,0 +1 @@ +xcopy . E:\data\ComfyUI\custom_nodes\ComfyUI-FluxExt-MZ /Y/E/H/C/I /EXCLUDE:exclude.txt \ No newline at end of file diff --git a/exclude.txt b/exclude.txt new file mode 100644 index 0000000..d436eab --- /dev/null +++ b/exclude.txt @@ -0,0 +1,2 @@ +.git + diff --git a/mz_fluxext_core.py b/mz_fluxext_core.py new file mode 100644 index 0000000..d6f1eec --- /dev/null +++ b/mz_fluxext_core.py @@ -0,0 +1,124 @@ + +import gc +import json +from types import MethodType +import safetensors.torch +import torch +import torch.nn as nn +import safetensors + + +from torch import Tensor, nn + + +def MZ_Flux1VRAM_MT_Patch_call(args={}): + model = args.get("model") + + def other_to_cpu(): + model.model.diffusion_model.img_in.to("cpu") + model.model.diffusion_model.time_in.to("cpu") + model.model.diffusion_model.guidance_in.to("cpu") + model.model.diffusion_model.vector_in.to("cpu") + model.model.diffusion_model.txt_in.to("cpu") + model.model.diffusion_model.pe_embedder.to("cpu") + + def other_to_cuda(): + model.model.diffusion_model.img_in.to("cuda") + model.model.diffusion_model.time_in.to("cuda") + model.model.diffusion_model.guidance_in.to("cuda") + model.model.diffusion_model.vector_in.to("cuda") + model.model.diffusion_model.txt_in.to("cuda") + model.model.diffusion_model.pe_embedder.to("cuda") + + def double_blocks_to_cpu(layer_start=0, layer_size=-1): + if layer_size == -1: + model.model.diffusion_model.double_blocks.to("cpu") + else: + model.model.diffusion_model.double_blocks[layer_start:layer_start + + layer_size].to("cpu") + torch.cuda.empty_cache() + gc.collect() + + def double_blocks_to_cuda(layer_start=0, layer_size=-1): + if layer_size == -1: + model.model.diffusion_model.double_blocks.to("cuda") + else: + model.model.diffusion_model.double_blocks[layer_start:layer_start + + layer_size].to("cuda") + + def single_blocks_to_cpu(layer_start=0, layer_size=-1): + if layer_size == -1: + model.model.diffusion_model.single_blocks.to("cpu") + else: + model.model.diffusion_model.single_blocks[layer_start:layer_start + + layer_size].to("cpu") + torch.cuda.empty_cache() + gc.collect() + + def single_blocks_to_cuda(layer_start=0, layer_size=-1): + if layer_size == -1: + model.model.diffusion_model.single_blocks.to("cuda") + else: + model.model.diffusion_model.single_blocks[layer_start:layer_start + + layer_size].to("cuda") + + def generate_double_blocks_forward_hook(layer_start, layer_size): + def pre_only_double_blocks_forward_hook(module, inp): + + other_to_cpu() + + if layer_start > 0: + double_blocks_to_cpu(layer_start=0, layer_size=layer_start) + + double_blocks_to_cuda(layer_start=layer_start, + layer_size=layer_size) + # print("pre_only_double_blocks_forward_hook: ", + # layer_start, layer_size) + # input("Press Enter to continue...") + return inp + return pre_only_double_blocks_forward_hook + + def generate_single_blocks_forward_hook(layer_start, layer_size): + def pre_only_single_blocks_forward_hook(module, inp): + double_blocks_to_cpu() + if layer_start > 0: + single_blocks_to_cpu(layer_start=0, layer_size=layer_start) + + single_blocks_to_cuda(layer_start=layer_start, + layer_size=layer_size) + # print("pre_only_single_blocks_forward_hook: ", + # layer_start, layer_size) + # input("Press Enter to continue...") + return inp + return pre_only_single_blocks_forward_hook + + def pre_only_model_forward_hook(module, inp): + print("double_blocks to cpu") + double_blocks_to_cpu() + print("single_blocks to cpu") + single_blocks_to_cpu() + print("other to cuda") + other_to_cuda() + return inp + model.model.diffusion_model.register_forward_pre_hook( + pre_only_model_forward_hook) + + double_blocks_depth = len(model.model.diffusion_model.double_blocks) + steps = 7 + for i in range(0, double_blocks_depth, steps): + s = steps + if i + s > double_blocks_depth: + s = double_blocks_depth - i + model.model.diffusion_model.double_blocks[i].register_forward_pre_hook( + generate_double_blocks_forward_hook(i, s)) + + single_blocks_depth = len(model.model.diffusion_model.single_blocks) + steps = 7 + for i in range(0, single_blocks_depth, steps): + s = steps + if i + s > single_blocks_depth: + s = single_blocks_depth - i + model.model.diffusion_model.single_blocks[i].register_forward_pre_hook( + generate_single_blocks_forward_hook(i, s)) + + return (model,)