Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-20 01:31:13 +02:00
Code Issues Packages Projects Releases Wiki Activity

Quanto disable gemm kernels

Browse Source
This commit is contained in:
Disty0
2024-08-14 20:26:46 +03:00
parent 237cab2aa1
commit f3f721e39a
2 changed files with 14 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/model_flux.py
+12 -7
View File
@@ -9,18 +9,19 @@ from modules import devices, shared
def load_quanto_transformer(repo_path):
def load_quanto_transformer(repo_path, device):
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("cpu"))
requantize(transformer, state_dict, quantization_map, device=torch.device(device))
transformer.eval()
return transformer
def load_quanto_text_encoder_2(repo_path):
def load_quanto_text_encoder_2(repo_path, device):
from optimum.quanto import requantize
with open(repo_path + "/" + "text_encoder_2/quantization_map.json", "r") as f:
quantization_map = json.load(f)
@@ -29,17 +30,21 @@ def load_quanto_text_encoder_2(repo_path):
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("cpu"))
requantize(text_encoder_2, state_dict, quantization_map, device=torch.device(device))
text_encoder_2.eval()
return text_encoder_2
def load_flux(checkpoint_info, diffusers_load_config):
if "qint8" in checkpoint_info.name.lower() or "qint4" in checkpoint_info.name.lower():
shared.log.debug(f'Loading FLUX: model="{checkpoint_info.name}" quant=True')
from installer import install
install('optimum-quanto', quiet=True)
shared.log.debug(f'Loading FLUX: model="{checkpoint_info.name}" quant=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)
pipe.text_encoder_2 = load_quanto_text_encoder_2(checkpoint_info.path)
pipe.transformer = load_quanto_transformer(checkpoint_info.path, device)
pipe.text_encoder_2 = load_quanto_text_encoder_2(checkpoint_info.path, device)
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')
modules/sd_models_compile.py
+2
View File
@@ -222,6 +222,8 @@ def optimum_quanto_weights(sd_model):
global quant_last_model_name, quant_last_model_device
from installer import install
install('optimum-quanto', quiet=True)
from optimum import quanto
quanto.tensor.qbits.QBitsTensor.create = lambda *args, **kwargs: quanto.tensor.qbits.QBitsTensor(*args, **kwargs)
sd_model = apply_compile_to_model(sd_model, optimum_quanto_model, shared.opts.optimum_quanto_weights, op="optimum-quanto")
if quant_last_model_name is not None: