Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 17:24:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

Flux load quant model to cpu

Browse Source
This commit is contained in:
Disty0
2024-08-15 22:17:41 +03:00
parent f3f721e39a
commit 80be079130
1 changed files with 7 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/model_flux.py
+7 -8
View File
@@ -5,23 +5,23 @@ import transformers
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download
from modules import devices, shared
from modules import shared
def load_quanto_transformer(repo_path, device):
def load_quanto_transformer(repo_path):
from optimum.quanto import requantize
with open(repo_path + "/" + "transformer/quantization_map.json", "r") as f:
quantization_map = json.load(f)
with torch.device("meta"):
transformer = diffusers.FluxTransformer2DModel.from_config(repo_path + "/" + "transformer/config.json").to(torch.bfloat16)
state_dict = load_file(repo_path + "/" + "transformer/diffusion_pytorch_model.safetensors")
requantize(transformer, state_dict, quantization_map, device=torch.device(device))
requantize(transformer, state_dict, quantization_map, device=torch.device("cpu"))
transformer.eval()
return transformer
def load_quanto_text_encoder_2(repo_path, device):
def load_quanto_text_encoder_2(repo_path):
from optimum.quanto import requantize
with open(repo_path + "/" + "text_encoder_2/quantization_map.json", "r") as f:
quantization_map = json.load(f)
@@ -30,7 +30,7 @@ def load_quanto_text_encoder_2(repo_path, device):
with torch.device("meta"):
text_encoder_2 = transformers.T5EncoderModel(t5_config).to(torch.bfloat16)
state_dict = load_file(repo_path + "/" + "text_encoder_2/model.safetensors")
requantize(text_encoder_2, state_dict, quantization_map, device=torch.device(device))
requantize(text_encoder_2, state_dict, quantization_map, device=torch.device("cpu"))
text_encoder_2.eval()
return text_encoder_2
@@ -41,10 +41,9 @@ def load_flux(checkpoint_info, diffusers_load_config):
install('optimum-quanto', quiet=True)
from optimum import quanto
quanto.tensor.qbits.QBitsTensor.create = lambda *args, **kwargs: quanto.tensor.qbits.QBitsTensor(*args, **kwargs)
device = devices.device if shared.opts.diffusers_offload_mode == "none" else "cpu"
pipe = diffusers.FluxPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, transformer=None, text_encoder_2=None, **diffusers_load_config)
pipe.transformer = load_quanto_transformer(checkpoint_info.path, device)
pipe.text_encoder_2 = load_quanto_text_encoder_2(checkpoint_info.path, device)
pipe.transformer = load_quanto_transformer(checkpoint_info.path)
pipe.text_encoder_2 = load_quanto_text_encoder_2(checkpoint_info.path)
else:
pipe = diffusers.FluxPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
shared.log.debug(f'Loading FLUX: model="{checkpoint_info.name}" quant=False')