Spaces:
Configuration error
Configuration error
| 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 | |
| import copy | |
| def Flux1PartialLoad_Patch(args={}): | |
| model = args.get("model") | |
| double_blocks_cuda_size = args.get("double_blocks_cuda_size") | |
| single_blocks_cuda_size = args.get("single_blocks_cuda_size") | |
| 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") | |
| torch.cuda.empty_cache() | |
| 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 = double_blocks_cuda_size | |
| 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 = single_blocks_cuda_size | |
| 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,) | |