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

Unify quant options

Browse Source
This commit is contained in:
Disty0
2025-06-27 21:05:14 +03:00
parent 3b8ced444c
commit 71f7474de2
17 changed files with 155 additions and 154 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/control/units/controlnet.py
+11 -3
View File
@@ -303,21 +303,29 @@ class ControlNet():
return
if self.dtype is not None:
self.model.to(self.dtype)
if "ControlNet" in opts.sdnq_quantize_weights:
if "Control" in opts.sdnq_quantize_weights:
try:
log.debug(f'Control {what} model SDNQ Compress: id="{model_id}"')
from modules.model_quant import sdnq_quantize_model
self.model = sdnq_quantize_model(self.model)
except Exception as e:
log.error(f'Control {what} model SDNQ Compression failed: id="{model_id}" {e}')
elif "ControlNet" in opts.optimum_quanto_weights:
elif "Control" in opts.optimum_quanto_weights:
try:
log.debug(f'Control {what} model Optimum Quanto: id="{model_id}"')
model_quant.load_quanto('Load model: type=ControlNet')
model_quant.load_quanto('Load model: type=Control')
from modules.model_quant import optimum_quanto_model
self.model = optimum_quanto_model(self.model)
except Exception as e:
log.error(f'Control {what} model Optimum Quanto: id="{model_id}" {e}')
elif "Control" in opts.torchao_quantization:
try:
log.debug(f'Control {what} model Torch AO: id="{model_id}"')
model_quant.load_torchao('Load model: type=Control')
from modules.model_quant import torchao_quantization
self.model = torchao_quantization(self.model)
except Exception as e:
log.error(f'Control {what} model Torch AO: id="{model_id}" {e}')
if self.device is not None:
self.model.to(self.device)
t1 = time.time()
modules/infiniteyou/pipeline_infu_flux.py
+1 -1
View File
@@ -151,7 +151,7 @@ class InfUFluxPipeline:
local_path = snapshot_download(repo_id='ByteDance/InfiniteYou', cache_dir=shared.opts.hfcache_dir)
infiniteyou_path = os.path.join(local_path, f'infu_flux_{infu_flux_version}', model_version)
infusenet_path = os.path.join(infiniteyou_path, 'InfuseNetModel')
quant_args = model_quant.create_config(module='ControlNet')
quant_args = model_quant.create_config(module='Control')
shared.log.debug(f'InfiniteYou: fn="{infusenet_path}" load infusenet')
self.infusenet = FluxControlNetModel.from_pretrained(
infusenet_path,
modules/model_chroma.py
+5 -5
View File
@@ -112,10 +112,10 @@ def load_chroma_bnb(checkpoint_info, diffusers_load_config): # pylint: disable=u
def load_quants(kwargs, pretrained_model_name_or_path, cache_dir, allow_quant):
try:
if 'transformer' not in kwargs and model_quant.check_nunchaku('Transformer'):
if 'transformer' not in kwargs and model_quant.check_nunchaku('Model'):
raise NotImplementedError('Nunchaku does not support Chroma Model yet. See https://github.com/mit-han-lab/nunchaku/issues/167')
elif 'transformer' not in kwargs and model_quant.check_quant('Transformer'):
quant_args = model_quant.create_config(allow=allow_quant, module='Transformer', modules_to_not_convert=["distilled_guidance_layer"])
elif 'transformer' not in kwargs and model_quant.check_quant('Model'):
quant_args = model_quant.create_config(allow=allow_quant, module='Model', modules_to_not_convert=["distilled_guidance_layer"])
if quant_args:
if os.path.isfile(pretrained_model_name_or_path):
kwargs['transformer'] = diffusers.ChromaTransformer2DModel.from_single_file(pretrained_model_name_or_path, cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_args)
@@ -166,7 +166,7 @@ def load_transformer(file_path): # triggered by opts.sd_unet change
if _transformer is not None:
transformer = _transformer
else:
quant_args = model_quant.create_config(module='Transformer', modules_to_not_convert=["distilled_guidance_layer"])
quant_args = model_quant.create_config(module='Model', modules_to_not_convert=["distilled_guidance_layer"])
if quant_args:
shared.log.info(f'Load module: type=UNet/Transformer file="{file_path}" offload={shared.opts.diffusers_offload_mode} quant=torchao dtype={devices.dtype}')
transformer = diffusers.ChromaTransformer2DModel.from_single_file(file_path, **diffusers_load_config, **quant_args)
@@ -284,7 +284,7 @@ def load_chroma(checkpoint_info, diffusers_load_config): # triggered by opts.sd_
else:
pipe = cls.from_pretrained(repo_id or fn, cache_dir=shared.opts.diffusers_dir, **kwargs, **diffusers_load_config)
if shared.opts.teacache_enabled and model_quant.check_nunchaku('Transformer'):
if shared.opts.teacache_enabled and model_quant.check_nunchaku('Model'):
from nunchaku.caching.diffusers_adapters import apply_cache_on_pipe
apply_cache_on_pipe(pipe, residual_diff_threshold=0.12)
modules/model_cogview.py
+2 -2
View File
@@ -7,7 +7,7 @@ def load_cogview3(checkpoint_info, diffusers_load_config={}):
modelloader.hf_login()
repo_id = sd_models.path_to_repo(checkpoint_info.name)
load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Transformer')
load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model')
shared.log.debug(f'Load model: type=CogView3 transformer="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
transformer = diffusers.CogView3PlusTransformer2DModel.from_pretrained(
repo_id,
@@ -44,7 +44,7 @@ def load_cogview4(checkpoint_info, diffusers_load_config={}):
modelloader.hf_login()
repo_id = sd_models.path_to_repo(checkpoint_info.name)
load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Transformer')
load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model')
shared.log.debug(f'Load model: type=CogView4 transformer="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
transformer = diffusers.CogView4Transformer2DModel.from_pretrained(
repo_id,
modules/model_cosmos.py
+2 -2
View File
@@ -6,7 +6,7 @@ from modules import shared, devices, sd_models, model_quant, modelloader, sd_hij
def load_transformer(repo_id, diffusers_load_config={}):
load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Transformer', device_map=True)
load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model', device_map=True)
fn = None
if shared.opts.sd_unet is not None and shared.opts.sd_unet != 'Default':
@@ -49,7 +49,7 @@ def load_text_encoder(repo_id, diffusers_load_config={}):
if shared.opts.diffusers_offload_mode != 'none' and text_encoder is not None:
sd_models.move_model(text_encoder, devices.cpu)
load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='LLM', device_map=True)
load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True)
llama_repo = shared.opts.model_h1_llama_repo if shared.opts.model_h1_llama_repo != 'Default' else 'meta-llama/Meta-Llama-3.1-8B-Instruct'
shared.log.debug(f'Load model: type=HiDream te4="{llama_repo}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
modules/model_flex.py
+2 -2
View File
@@ -6,7 +6,7 @@ from modules import shared, devices, sd_models, model_quant, modelloader, sd_hij
def load_transformer(repo_id, diffusers_load_config={}):
load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Transformer', device_map=True)
load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model', device_map=True)
fn = None
if shared.opts.sd_unet is not None and shared.opts.sd_unet != 'Default':
@@ -24,7 +24,7 @@ def load_transformer(repo_id, diffusers_load_config={}):
elif fn is not None and 'safetensors' in fn.lower():
shared.log.debug(f'Load model: type=FLEX transformer="{repo_id}" quant="{model_quant.get_quant(repo_id)}" args={load_args}')
transformer = diffusers.FluxTransformer2DModel.from_single_file(fn, cache_dir=shared.opts.hfcache_dir, **load_args)
# elif model_quant.check_nunchaku('Transformer'):
# elif model_quant.check_nunchaku('Model'):
# shared.log.error(f'Load model: type=HiDream transformer="{repo_id}" quant="Nunchaku" unsupported')
# transformer = None
else:
modules/model_flux.py
+5 -5
View File
@@ -110,7 +110,7 @@ def load_flux_bnb(checkpoint_info, diffusers_load_config): # pylint: disable=unu
def load_quants(kwargs, repo_id, cache_dir, allow_quant):
try:
if 'transformer' not in kwargs and model_quant.check_nunchaku('Transformer'):
if 'transformer' not in kwargs and model_quant.check_nunchaku('Model'):
import nunchaku
nunchaku_precision = nunchaku.utils.get_precision()
nunchaku_repo = None
@@ -128,8 +128,8 @@ def load_quants(kwargs, repo_id, cache_dir, allow_quant):
kwargs['transformer'].quantization_method = 'SVDQuant'
if shared.opts.nunchaku_attention:
kwargs['transformer'].set_attention_impl("nunchaku-fp16")
elif 'transformer' not in kwargs and model_quant.check_quant('Transformer'):
quant_args = model_quant.create_config(allow=allow_quant, module='Transformer')
elif 'transformer' not in kwargs and model_quant.check_quant('Model'):
quant_args = model_quant.create_config(allow=allow_quant, module='Model')
if quant_args:
kwargs['transformer'] = diffusers.FluxTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_args)
if 'text_encoder_2' not in kwargs and model_quant.check_nunchaku('TE'):
@@ -186,7 +186,7 @@ def load_transformer(file_path): # triggered by opts.sd_unet change
transformer, _text_encoder_2 = load_flux_nf4(file_path, prequantized=False)
if transformer is not None:
return transformer
quant_args = model_quant.create_config(module='Transformer')
quant_args = model_quant.create_config(module='Model')
if quant_args:
shared.log.info(f'Load module: type=UNet/Transformer file="{file_path}" offload={shared.opts.diffusers_offload_mode} quant=torchao dtype={devices.dtype}')
transformer = diffusers.FluxTransformer2DModel.from_single_file(file_path, **diffusers_load_config, **quant_args)
@@ -347,7 +347,7 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
else:
pipe = cls.from_pretrained(repo_id, cache_dir=shared.opts.diffusers_dir, **kwargs, **diffusers_load_config)
if shared.opts.teacache_enabled and model_quant.check_nunchaku('Transformer'):
if shared.opts.teacache_enabled and model_quant.check_nunchaku('Model'):
from nunchaku.caching.diffusers_adapters import apply_cache_on_pipe
apply_cache_on_pipe(pipe, residual_diff_threshold=0.12)
modules/model_hidream.py
+3 -3
View File
@@ -6,7 +6,7 @@ from modules import shared, devices, sd_models, model_quant, modelloader, sd_hij
def load_transformer(repo_id, diffusers_load_config={}):
load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Transformer', device_map=True)
load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model', device_map=True)
fn = None
if shared.opts.sd_unet is not None and shared.opts.sd_unet != 'Default':
@@ -24,7 +24,7 @@ def load_transformer(repo_id, diffusers_load_config={}):
elif fn is not None and 'safetensors' in fn.lower():
shared.log.debug(f'Load model: type=HiDream transformer="{repo_id}" quant="{model_quant.get_quant(repo_id)}" args={load_args}')
transformer = diffusers.HiDreamImageTransformer2DModel.from_single_file(fn, cache_dir=shared.opts.hfcache_dir, **load_args)
# elif model_quant.check_nunchaku('Transformer'):
# elif model_quant.check_nunchaku('Model'):
# shared.log.error(f'Load model: type=HiDream transformer="{repo_id}" quant="Nunchaku" unsupported')
# transformer = None
else:
@@ -56,7 +56,7 @@ def load_text_encoders(repo_id, diffusers_load_config={}):
if shared.opts.diffusers_offload_mode != 'none' and text_encoder_3 is not None:
sd_models.move_model(text_encoder_3, devices.cpu)
load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='LLM', device_map=True)
load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True)
llama_repo = shared.opts.model_h1_llama_repo if shared.opts.model_h1_llama_repo != 'Default' else 'meta-llama/Meta-Llama-3.1-8B-Instruct'
shared.log.debug(f'Load model: type=HiDream te4="{llama_repo}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
modules/model_lumina.py
+1 -1
View File
@@ -30,7 +30,7 @@ def load_lumina2(checkpoint_info, diffusers_load_config={}):
shared.log.debug(f'Transformers cache: type=teacache patch=forward cls={diffusers.Lumina2Transformer2DModel.__name__}')
diffusers.Lumina2Transformer2DModel.forward = teacache.teacache_lumina2_forward # patch must be done before transformer is loaded
load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, module='Transformer')
load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, module='Model')
if shared.opts.sd_unet != 'Default':
try:
debug(f'Load model: type=Lumina2 unet="{shared.opts.sd_unet}"')
modules/model_omnigen.py
+1 -1
View File
@@ -24,7 +24,7 @@ def load_omnigen(checkpoint_info, diffusers_load_config={}): # pylint: disable=u
if debug:
errors.display(e, 'OmniGen VAE:')
load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, module='Transformer')
load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, module='Model')
transformer = diffusers.OmniGenTransformer2DModel.from_pretrained(
repo_id,
subfolder="transformer",
modules/model_pixart.py
+1 -1
View File
@@ -15,7 +15,7 @@ def load_pixart(checkpoint_info, diffusers_load_config={}):
if not file_exists(repo_id_pipe, "model_index.json"):
repo_id_pipe = "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS"
load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Transformer')
load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model')
transformer = diffusers.PixArtTransformer2DModel.from_pretrained(
repo_id,
subfolder='transformer',
modules/model_quant.py
+103 -109
View File
@@ -37,128 +37,122 @@ def get_quant(name):
return 'none'
def create_bnb_config(kwargs = None, allow_bnb: bool = True, module: str = 'Model', modules_to_not_convert: list = []):
def create_bnb_config(kwargs = None, allow: bool = True, module: str = 'Model', modules_to_not_convert: list = []):
from modules import shared, devices
if len(shared.opts.bnb_quantization) > 0 and allow_bnb:
if 'Model' in shared.opts.bnb_quantization or (module is not None and module in shared.opts.bnb_quantization) or module == 'any':
load_bnb()
if bnb is None:
return kwargs
bnb_config = diffusers.BitsAndBytesConfig(
load_in_8bit=shared.opts.bnb_quantization_type in ['fp8'],
load_in_4bit=shared.opts.bnb_quantization_type in ['nf4', 'fp4'],
bnb_4bit_quant_storage=shared.opts.bnb_quantization_storage,
bnb_4bit_quant_type=shared.opts.bnb_quantization_type,
bnb_4bit_compute_dtype=devices.dtype,
llm_int8_skip_modules=modules_to_not_convert,
)
log.debug(f'Quantization: module={module} type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
if kwargs is None:
return bnb_config
else:
kwargs['quantization_config'] = bnb_config
return kwargs
if allow and (module == 'any' or module in shared.opts.bnb_quantization):
load_bnb()
if bnb is None:
return kwargs
bnb_config = diffusers.BitsAndBytesConfig(
load_in_8bit=shared.opts.bnb_quantization_type in ['fp8'],
load_in_4bit=shared.opts.bnb_quantization_type in ['nf4', 'fp4'],
bnb_4bit_quant_storage=shared.opts.bnb_quantization_storage,
bnb_4bit_quant_type=shared.opts.bnb_quantization_type,
bnb_4bit_compute_dtype=devices.dtype,
llm_int8_skip_modules=modules_to_not_convert,
)
log.debug(f'Quantization: module={module} type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
if kwargs is None:
return bnb_config
else:
kwargs['quantization_config'] = bnb_config
return kwargs
return kwargs
def create_ao_config(kwargs = None, allow_ao: bool = True, module: str = 'Model', modules_to_not_convert: list = []):
def create_ao_config(kwargs = None, allow: bool = True, module: str = 'Model', modules_to_not_convert: list = []):
from modules import shared
if len(shared.opts.torchao_quantization) > 0 and (shared.opts.torchao_quantization_mode in {'pre', 'auto'}) and allow_ao:
if 'Model' in shared.opts.torchao_quantization or (module is not None and module in shared.opts.torchao_quantization) or module == 'any':
torchao = load_torchao()
if torchao is None:
return kwargs
if module in {'TE', 'LLM'}:
ao_config = transformers.TorchAoConfig(quant_type=shared.opts.torchao_quantization_type, modules_to_not_convert=modules_to_not_convert)
else:
ao_config = diffusers.TorchAoConfig(shared.opts.torchao_quantization_type, modules_to_not_convert=modules_to_not_convert)
log.debug(f'Quantization: module={module} type=torchao dtype={shared.opts.torchao_quantization_type}')
if kwargs is None:
return ao_config
else:
kwargs['quantization_config'] = ao_config
return kwargs
if allow and (shared.opts.torchao_quantization_mode in {'pre', 'auto'}) and (module == 'any' or module in shared.opts.torchao_quantization):
torchao = load_torchao()
if torchao is None:
return kwargs
if module in {'TE', 'LLM'}:
ao_config = transformers.TorchAoConfig(quant_type=shared.opts.torchao_quantization_type, modules_to_not_convert=modules_to_not_convert)
else:
ao_config = diffusers.TorchAoConfig(shared.opts.torchao_quantization_type, modules_to_not_convert=modules_to_not_convert)
log.debug(f'Quantization: module={module} type=torchao dtype={shared.opts.torchao_quantization_type}')
if kwargs is None:
return ao_config
else:
kwargs['quantization_config'] = ao_config
return kwargs
return kwargs
def create_quanto_config(kwargs = None, allow_quanto: bool = True, module: str = 'Model', modules_to_not_convert: list = []):
def create_quanto_config(kwargs = None, allow: bool = True, module: str = 'Model', modules_to_not_convert: list = []):
from modules import shared
if len(shared.opts.quanto_quantization) > 0 and allow_quanto:
if 'Model' in shared.opts.quanto_quantization or (module is not None and module in shared.opts.quanto_quantization) or module == 'any':
load_quanto(silent=True)
if optimum_quanto is None:
return kwargs
if module in {'TE', 'LLM'}:
quanto_config = transformers.QuantoConfig(weights=shared.opts.quanto_quantization_type, modules_to_not_convert=modules_to_not_convert)
quanto_config.weights_dtype = quanto_config.weights
else:
quanto_config = diffusers.QuantoConfig(weights_dtype=shared.opts.quanto_quantization_type, modules_to_not_convert=modules_to_not_convert)
quanto_config.activations = None # patch so it works with transformers
quanto_config.weights = quanto_config.weights_dtype
log.debug(f'Quantization: module={module} type=quanto dtype={shared.opts.quanto_quantization_type}')
if kwargs is None:
return quanto_config
else:
kwargs['quantization_config'] = quanto_config
return kwargs
if allow and (module == 'any' or module in shared.opts.quanto_quantization):
load_quanto(silent=True)
if optimum_quanto is None:
return kwargs
if module in {'TE', 'LLM'}:
quanto_config = transformers.QuantoConfig(weights=shared.opts.quanto_quantization_type, modules_to_not_convert=modules_to_not_convert)
quanto_config.weights_dtype = quanto_config.weights
else:
quanto_config = diffusers.QuantoConfig(weights_dtype=shared.opts.quanto_quantization_type, modules_to_not_convert=modules_to_not_convert)
quanto_config.activations = None # patch so it works with transformers
quanto_config.weights = quanto_config.weights_dtype
log.debug(f'Quantization: module={module} type=quanto dtype={shared.opts.quanto_quantization_type}')
if kwargs is None:
return quanto_config
else:
kwargs['quantization_config'] = quanto_config
return kwargs
return kwargs
def create_sdnq_config(kwargs = None, allow_sdnq: bool = True, module: str = 'Model', weights_dtype: str = None, modules_to_not_convert: list = []):
def create_sdnq_config(kwargs = None, allow: bool = True, module: str = 'Model', weights_dtype: str = None, modules_to_not_convert: list = []):
from modules import devices, shared
if len(shared.opts.sdnq_quantize_weights) > 0 and (shared.opts.sdnq_quantize_mode in {'pre', 'auto'}) and allow_sdnq:
if 'Model' in shared.opts.sdnq_quantize_weights or (module is not None and module in shared.opts.sdnq_quantize_weights) or module == 'any':
from modules.sdnq import SDNQQuantizer, SDNQConfig
diffusers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer
transformers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer
diffusers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig
transformers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig
if allow and (shared.opts.sdnq_quantize_mode in {'pre', 'auto'}) and (module == 'any' or module in shared.opts.sdnq_quantize_weights):
from modules.sdnq import SDNQQuantizer, SDNQConfig
diffusers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer
transformers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer
diffusers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig
transformers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig
if weights_dtype is None:
if module in {"TE", "LLM"} and shared.opts.sdnq_quantize_weights_mode_te not in {"same as model", "default"}:
weights_dtype = shared.opts.sdnq_quantize_weights_mode_te
else:
weights_dtype = shared.opts.sdnq_quantize_weights_mode
if weights_dtype is None or weights_dtype == 'none':
return kwargs
if shared.opts.device_map == "gpu":
quantization_device = devices.device
return_device = devices.device
elif shared.opts.diffusers_offload_mode in {"none", "model"}:
quantization_device = devices.device if shared.opts.sdnq_quantize_with_gpu else devices.cpu
return_device = devices.device
elif shared.opts.sdnq_quantize_with_gpu:
quantization_device = devices.device
return_device = devices.cpu
if weights_dtype is None:
if module in {"TE", "LLM"} and shared.opts.sdnq_quantize_weights_mode_te not in {"same as model", "default"}:
weights_dtype = shared.opts.sdnq_quantize_weights_mode_te
else:
quantization_device = None
return_device = None
weights_dtype = shared.opts.sdnq_quantize_weights_mode
if weights_dtype is None or weights_dtype == 'none':
return kwargs
sdnq_config = SDNQConfig(
weights_dtype=weights_dtype,
group_size=shared.opts.sdnq_quantize_weights_group_size,
quant_conv=shared.opts.sdnq_quantize_conv_layers,
use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul,
use_quantized_matmul_conv=shared.opts.sdnq_use_quantized_matmul_conv,
dequantize_fp32=shared.opts.sdnq_dequantize_fp32,
quantization_device=quantization_device,
return_device=return_device,
modules_to_not_convert=modules_to_not_convert,
)
log.debug(f'Quantization: module="{module}" type=sdnq dtype={weights_dtype} matmul={shared.opts.sdnq_use_quantized_matmul} group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} dequantize_fp32={shared.opts.sdnq_dequantize_fp32} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} quantization_device={quantization_device} return_device={return_device}')
if kwargs is None:
return sdnq_config
else:
kwargs['quantization_config'] = sdnq_config
return kwargs
if shared.opts.device_map == "gpu":
quantization_device = devices.device
return_device = devices.device
elif shared.opts.diffusers_offload_mode in {"none", "model"}:
quantization_device = devices.device if shared.opts.sdnq_quantize_with_gpu else devices.cpu
return_device = devices.device
elif shared.opts.sdnq_quantize_with_gpu:
quantization_device = devices.device
return_device = devices.cpu
else:
quantization_device = None
return_device = None
sdnq_config = SDNQConfig(
weights_dtype=weights_dtype,
group_size=shared.opts.sdnq_quantize_weights_group_size,
quant_conv=shared.opts.sdnq_quantize_conv_layers,
use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul,
use_quantized_matmul_conv=shared.opts.sdnq_use_quantized_matmul_conv,
dequantize_fp32=shared.opts.sdnq_dequantize_fp32,
quantization_device=quantization_device,
return_device=return_device,
modules_to_not_convert=modules_to_not_convert,
)
log.debug(f'Quantization: module="{module}" type=sdnq dtype={weights_dtype} matmul={shared.opts.sdnq_use_quantized_matmul} group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} dequantize_fp32={shared.opts.sdnq_dequantize_fp32} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} quantization_device={quantization_device} return_device={return_device}')
if kwargs is None:
return sdnq_config
else:
kwargs['quantization_config'] = sdnq_config
return kwargs
return kwargs
def check_quant(module: str = ''):
from modules import shared
if 'Model' in shared.opts.bnb_quantization or 'Model' in shared.opts.torchao_quantization or 'Model' in shared.opts.quanto_quantization or 'Model' in shared.opts.sdnq_quantize_weights:
return True
if module in shared.opts.bnb_quantization or module in shared.opts.torchao_quantization or module in shared.opts.quanto_quantization or module in shared.opts.sdnq_quantize_weights:
return True
return False
@@ -166,7 +160,7 @@ def check_quant(module: str = ''):
def check_nunchaku(module: str = ''):
from modules import shared
if 'Model' not in shared.opts.nunchaku_quantization and module not in shared.opts.nunchaku_quantization:
if module not in shared.opts.nunchaku_quantization:
return False
from modules import mit_nunchaku
mit_nunchaku.install_nunchaku()
@@ -178,22 +172,22 @@ def check_nunchaku(module: str = ''):
def create_config(kwargs = None, allow: bool = True, module: str = 'Model', modules_to_not_convert = []):
if kwargs is None:
kwargs = {}
kwargs = create_sdnq_config(kwargs, allow_sdnq=allow, module=module, modules_to_not_convert=modules_to_not_convert)
kwargs = create_sdnq_config(kwargs, allow=allow, module=module, modules_to_not_convert=modules_to_not_convert)
if kwargs is not None and 'quantization_config' in kwargs:
if debug:
log.trace(f'Quantization: type=sdnq config={kwargs.get("quantization_config", None)}')
return kwargs
kwargs = create_bnb_config(kwargs, allow_bnb=allow, module=module, modules_to_not_convert=modules_to_not_convert)
kwargs = create_bnb_config(kwargs, allow=allow, module=module, modules_to_not_convert=modules_to_not_convert)
if kwargs is not None and 'quantization_config' in kwargs:
if debug:
log.trace(f'Quantization: type=bnb config={kwargs.get("quantization_config", None)}')
return kwargs
kwargs = create_quanto_config(kwargs, allow_quanto=allow, module=module, modules_to_not_convert=modules_to_not_convert)
kwargs = create_quanto_config(kwargs, allow=allow, module=module, modules_to_not_convert=modules_to_not_convert)
if kwargs is not None and 'quantization_config' in kwargs:
if debug:
log.trace(f'Quantization: type=quanto config={kwargs.get("quantization_config", None)}')
return kwargs
kwargs = create_ao_config(kwargs, allow_ao=allow, module=module, modules_to_not_convert=modules_to_not_convert)
kwargs = create_ao_config(kwargs, allow=allow, module=module, modules_to_not_convert=modules_to_not_convert)
if kwargs is not None and 'quantization_config' in kwargs:
if debug:
log.trace(f'Quantization: type=torchao config={kwargs.get("quantization_config", None)}')
@@ -309,13 +303,13 @@ def apply_layerwise(sd_model, quiet:bool=False):
m.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=non_blocking)
m.quantization_method = 'LayerWise'
log.quiet(quiet, f'Quantization: type=layerwise module={module} cls={cls} storage={storage_dtype} compute={devices.dtype} blocking={not non_blocking}')
if module.startswith('transformer') and ('Model' in shared.opts.layerwise_quantization or 'Transformer' in shared.opts.layerwise_quantization):
if module.startswith('transformer') and ('Model' in shared.opts.layerwise_quantization):
m = getattr(sd_model, module)
if hasattr(m, 'enable_layerwise_casting'):
m.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=non_blocking)
m.quantization_method = 'LayerWise'
log.quiet(quiet, f'Quantization: type=layerwise module={module} cls={cls} storage={storage_dtype} compute={devices.dtype} blocking={not non_blocking}')
if module.startswith('text_encoder') and ('Model' in shared.opts.layerwise_quantization or 'TE' in shared.opts.layerwise_quantization) and ('clip' not in cls.lower()):
if module.startswith('text_encoder') and ('TE' in shared.opts.layerwise_quantization) and ('clip' not in cls.lower()):
m = getattr(sd_model, module)
if hasattr(m, 'enable_layerwise_casting'):
m.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=non_blocking)
@@ -575,7 +569,7 @@ def get_dit_args(load_config:dict={}, module:str=None, device_map:bool=False, al
# if 'variant' in config:
# del config['variant']
if device_map:
if devices.backend == "ipex" and os.environ.get('UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS', '0') != '1' and module in {'TE', 'LLM'}:
if devices.backend == 'ipex' and os.environ.get('UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS', '0') != '1' and module in {'TE', 'LLM'}:
config['device_map'] = 'cpu' # alchemist gpus hits the 4GB allocation limit with transformers, UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS emulates above 4GB allocations
elif shared.opts.device_map == 'cpu':
config['device_map'] = 'cpu'
modules/model_sana.py
+3 -3
View File
@@ -7,14 +7,14 @@ from modules import shared, sd_models, devices, modelloader, model_quant
def load_quants(kwargs, repo_id, cache_dir):
kwargs_copy = kwargs.copy()
if model_quant.check_nunchaku('Transformer') and 'Sana_1600M' in repo_id: # only sana-1600m
if model_quant.check_nunchaku('Model') and 'Sana_1600M' in repo_id: # only sana-1600m
import nunchaku
nunchaku_precision = nunchaku.utils.get_precision()
nunchaku_repo = f"mit-han-lab/svdq-{nunchaku_precision}-sana-1600m"
shared.log.debug(f'Load module: quant=Nunchaku module=transformer repo="{nunchaku_repo}" precision={nunchaku_precision} attention={shared.opts.nunchaku_attention}')
kwargs['transformer'] = nunchaku.NunchakuSanaTransformer2DModel.from_pretrained(nunchaku_repo, torch_dtype=devices.dtype)
elif model_quant.check_quant('Transformer'):
load_args, quant_args = model_quant.get_dit_args(kwargs_copy, module='Transformer')
elif model_quant.check_quant('Model'):
load_args, quant_args = model_quant.get_dit_args(kwargs_copy, module='Model')
if quant_args:
kwargs['transformer'] = diffusers.SanaTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=cache_dir, **load_args, **quant_args)
if model_quant.check_quant('TE'):
modules/model_sd3.py
+1 -1
View File
@@ -57,7 +57,7 @@ def load_overrides(kwargs, cache_dir):
def load_quants(kwargs, repo_id, cache_dir):
quant_args = model_quant.create_config(module='Transformer')
quant_args = model_quant.create_config(module='Model')
if quant_args and 'quantization_config' in quant_args:
kwargs['transformer'] = diffusers.SD3Transformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_args)
quant_args = model_quant.create_config(module='TE')
modules/sd_models_utils.py
+2 -3
View File
@@ -151,14 +151,13 @@ def patch_diffuser_config(sd_model, model_file):
def apply_function_to_model(sd_model, function, options, op=None):
if "Model" in options or "Transformer" in options:
if hasattr(sd_model, 'transformer') and hasattr(sd_model.transformer, 'config'):
sd_model.transformer = function(sd_model.transformer, op="transformer", sd_model=sd_model)
if "Model" in options:
if hasattr(sd_model, 'model') and (hasattr(sd_model.model, 'config') or isinstance(sd_model.model, torch.nn.Module)):
sd_model.model = function(sd_model.model, op="model", sd_model=sd_model)
if hasattr(sd_model, 'unet') and hasattr(sd_model.unet, 'config'):
sd_model.unet = function(sd_model.unet, op="unet", sd_model=sd_model)
if hasattr(sd_model, 'transformer') and hasattr(sd_model.transformer, 'config'):
sd_model.transformer = function(sd_model.transformer, op="transformer", sd_model=sd_model)
if hasattr(sd_model, 'decoder_pipe') and hasattr(sd_model, 'decoder'):
sd_model.decoder = None
sd_model.decoder = sd_model.decoder_pipe.decoder = function(sd_model.decoder_pipe.decoder, op="decoder_pipe.decoder", sd_model=sd_model)
modules/shared.py
+11 -11
View File
@@ -502,7 +502,7 @@ options_templates.update(options_section(('backends', "Backend Settings"), {
"olive_cache_optimized": OptionInfo(True, 'Olive cache optimized models'),
"ipex_sep": OptionInfo("<h2>IPEX</h2>", "", gr.HTML, {"visible": devices.backend == "ipex"}),
"ipex_optimize": OptionInfo([], "IPEX Optimize", gr.CheckboxGroup, {"choices": ["Model", "VAE", "TE", "Upscaler"], "visible": devices.backend == "ipex"}),
"ipex_optimize": OptionInfo([], "IPEX Optimize", gr.CheckboxGroup, {"choices": ["Model", "TE", "VAE", "Upscaler"], "visible": devices.backend == "ipex"}),
"openvino_sep": OptionInfo("<h2>OpenVINO</h2>", "", gr.HTML, {"visible": cmd_opts.use_openvino}),
"openvino_devices": OptionInfo([], "OpenVINO devices to use", gr.CheckboxGroup, {"choices": get_openvino_device_list() if cmd_opts.use_openvino else [], "visible": cmd_opts.use_openvino}), # pylint: disable=E0606
@@ -517,7 +517,7 @@ options_templates.update(options_section(('backends', "Backend Settings"), {
options_templates.update(options_section(("quantization", "Quantization Settings"), {
"sdnq_quantize_sep": OptionInfo("<h2>SDNQ: SD.Next Quantization</h2>", "", gr.HTML),
"sdnq_quantize_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}),
"sdnq_quantize_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "LLM", "Control", "VAE"], "visible": native}),
"sdnq_quantize_mode": OptionInfo("auto", "Quantization mode", gr.Dropdown, {"choices": ["auto", "pre", "post"], "visible": native}),
"sdnq_quantize_weights_mode": OptionInfo("int8", "Quantization type", gr.Dropdown, {"choices": sdnq_quant_modes, "visible": native}),
"sdnq_quantize_weights_mode_te": OptionInfo("default", "Quantization type for Text Encoders", gr.Dropdown, {"choices": ['default'] + sdnq_quant_modes, "visible": native}),
@@ -531,41 +531,41 @@ options_templates.update(options_section(("quantization", "Quantization Settings
"sdnq_quantize_shuffle_weights": OptionInfo(False, "Shuffle weights in post mode", gr.Checkbox, {"visible": native}),
"bnb_quantization_sep": OptionInfo("<h2>BitsAndBytes</h2>", "", gr.HTML),
"bnb_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}),
"bnb_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "LLM", "VAE"], "visible": native}),
"bnb_quantization_type": OptionInfo("nf4", "Quantization type", gr.Dropdown, {"choices": ["nf4", "fp8", "fp4"], "visible": native}),
"bnb_quantization_storage": OptionInfo("uint8", "Backend storage", gr.Dropdown, {"choices": ["float16", "float32", "int8", "uint8", "float64", "bfloat16"], "visible": native}),
"quanto_quantization_sep": OptionInfo("<h2>Optimum Quanto</h2>", "", gr.HTML),
"quanto_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}),
"quanto_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "LLM"], "visible": native}),
"quanto_quantization_type": OptionInfo("int8", "Quantization weights type", gr.Dropdown, {"choices": ["float8", "int8", "int4", "int2"], "visible": native}),
"optimum_quanto_sep": OptionInfo("<h2>Optimum Quanto: post-load</h2>", "", gr.HTML),
"optimum_quanto_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}),
"optimum_quanto_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "Control", "VAE"], "visible": native}),
"optimum_quanto_weights_type": OptionInfo("qint8", "Quantization weights type", gr.Dropdown, {"choices": ["qint8", "qfloat8_e4m3fn", "qfloat8_e5m2", "qint4", "qint2"], "visible": native}),
"optimum_quanto_activations_type": OptionInfo("none", "Quantization activations type ", gr.Dropdown, {"choices": ["none", "qint8", "qfloat8_e4m3fn", "qfloat8_e5m2"], "visible": native}),
"optimum_quanto_shuffle_weights": OptionInfo(False, "Shuffle weights in post mode", gr.Checkbox, {"visible": native}),
"torchao_sep": OptionInfo("<h2>TorchAO</h2>", "", gr.HTML),
"torchao_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM"], "visible": native}),
"torchao_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "LLM", "Control", "VAE"], "visible": native}),
"torchao_quantization_mode": OptionInfo("auto", "Quantization mode", gr.Dropdown, {"choices": ["auto", "pre", "post"], "visible": native}),
"torchao_quantization_type": OptionInfo("int8_weight_only", "Quantization type", gr.Dropdown, {"choices": ["int4_weight_only", "int8_dynamic_activation_int4_weight", "int8_weight_only", "int8_dynamic_activation_int8_weight", "float8_weight_only", "float8_dynamic_activation_float8_weight", "float8_static_activation_float8_weight"], "visible": native}),
"layerwise_quantization_sep": OptionInfo("<h2>Layerwise Casting</h2>", "", gr.HTML),
"layerwise_quantization": OptionInfo([], "Layerwise casting enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "TE"], "visible": native}),
"layerwise_quantization": OptionInfo([], "Layerwise casting enabled", gr.CheckboxGroup, {"choices": ["Model", "TE"], "visible": native}),
"layerwise_quantization_storage": OptionInfo("float8_e4m3fn", "Layerwise casting storage", gr.Dropdown, {"choices": ["float8_e4m3fn", "float8_e5m2"], "visible": native}),
"layerwise_quantization_nonblocking": OptionInfo(False, "Layerwise non-blocking operations", gr.Checkbox, {"visible": native}),
"nunchaku_sep": OptionInfo("<h2>Nunchaku Engine</h2>", "", gr.HTML),
"nunchaku_quantization": OptionInfo([], "SVDQuant enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}),
"nunchaku_quantization": OptionInfo([], "SVDQuant enabled", gr.CheckboxGroup, {"choices": ["Model", "TE"], "visible": native}),
"nunchaku_attention": OptionInfo(False, "Nunchaku attention", gr.Checkbox, {"visible": native}),
"nunchaku_offload": OptionInfo(False, "Nunchaku offloading", gr.Checkbox, {"visible": native}),
"nncf_compress_sep": OptionInfo("<h2>NNCF: Neural Network Compression Framework</h2>", "", gr.HTML, {"visible": cmd_opts.use_openvino}),
"nncf_compress_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": cmd_opts.use_openvino}),
"nncf_compress_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "VAE"], "visible": cmd_opts.use_openvino}),
"nncf_compress_weights_mode": OptionInfo("INT8_SYM", "Quantization type", gr.Dropdown, {"choices": ["INT8", "INT4_ASYM", "INT8_SYM", "INT4_SYM", "NF4"], "visible": cmd_opts.use_openvino}),
"nncf_compress_weights_raito": OptionInfo(0, "Compress ratio", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01, "visible": cmd_opts.use_openvino}),
"nncf_compress_weights_group_size": OptionInfo(0, "Group size", gr.Slider, {"minimum": -1, "maximum": 4096, "step": 1, "visible": cmd_opts.use_openvino}),
"nncf_quantize": OptionInfo([], "Static Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "TE"], "visible": cmd_opts.use_openvino}),
"nncf_quantize": OptionInfo([], "Static Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "VAE"], "visible": cmd_opts.use_openvino}),
"nncf_quantize_mode": OptionInfo("INT8", "OpenVINO activations mode", gr.Dropdown, {"choices": ["INT8", "FP8_E4M3", "FP8_E5M2"], "visible": cmd_opts.use_openvino}),
}))
@@ -650,7 +650,7 @@ options_templates.update(options_section(('advanced', "Pipeline Modifiers"), {
options_templates.update(options_section(('compile', "Model Compile"), {
"cuda_compile_sep": OptionInfo("<h2>Model Compile</h2>", "", gr.HTML),
"cuda_compile": OptionInfo([] if not cmd_opts.use_openvino else ["Model", "VAE", "Upscaler"], "Compile Model", gr.CheckboxGroup, {"choices": ["Model", "VAE", "TE", "Upscaler"]}),
"cuda_compile": OptionInfo([] if not cmd_opts.use_openvino else ["Model", "VAE", "Upscaler"], "Compile Model", gr.CheckboxGroup, {"choices": ["Model", "TE", "VAE", "Upscaler"]}),
"cuda_compile_backend": OptionInfo("none" if not cmd_opts.use_openvino else "openvino_fx", "Model compile backend", gr.Radio, {"choices": ['none', 'inductor', 'cudagraphs', 'aot_ts_nvfuser', 'hidet', 'migraphx', 'ipex', 'onediff', 'stable-fast', 'deep-cache', 'olive-ai', 'openvino_fx']}),
"cuda_compile_mode": OptionInfo("default", "Model compile mode", gr.Radio, {"choices": ['default', 'reduce-overhead', 'max-autotune', 'max-autotune-no-cudagraphs']}),
"cuda_compile_fullgraph": OptionInfo(True if not cmd_opts.use_openvino else False, "Model compile fullgraph"),
modules/video_models/video_load.py
+1 -1
View File
@@ -38,7 +38,7 @@ def load_model(selected: models_def.Model):
# transformer
try:
quant_args = model_quant.create_config(module='Video')
quant_args = model_quant.create_config(module='Model')
debug(f'Video load: module=transformer repo="{selected.dit or selected.repo}" folder="{selected.dit_folder}" cls={selected.dit_cls.__name__} quant={model_quant.get_quant_type(quant_args)}')
transformer = selected.dit_cls.from_pretrained(
pretrained_model_name_or_path=selected.dit or selected.repo,