diff --git a/modules/model_quant.py b/modules/model_quant.py index f5baec12a..c3d2439ec 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -479,7 +479,7 @@ def apply_layerwise(sd_model, quiet:bool=False): def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weights_dtype: str = None, modules_to_not_convert: list = None, modules_dtype_dict: dict = None): global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement from modules import devices, shared, timer - from modules.sdnq import apply_sdnq_to_module + from modules.sdnq import sdnq_post_load_quant if weights_dtype is None: if op is not None and ("text_encoder" in op or op in {"TE", "LLM"}) and shared.opts.sdnq_quantize_weights_mode_te not in {"Same as model", "default"}: @@ -497,13 +497,6 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weigh if modules_dtype_dict is None: modules_dtype_dict = {} - if getattr(model, "_keep_in_fp32_modules", None) is not None: - modules_to_not_convert.extend(model._keep_in_fp32_modules) # pylint: disable=protected-access - if getattr(model, "_skip_layerwise_casting_patterns", None) is not None: - modules_to_not_convert.extend(model._skip_layerwise_casting_patterns) # pylint: disable=protected-access - if model.__class__.__name__ == "ChromaTransformer2DModel": - modules_to_not_convert.append("distilled_guidance_layer") - sdnq_modules_to_not_convert = [m.strip() for m in re.split(';|,| ', shared.opts.sdnq_modules_to_not_convert) if len(m.strip()) > 1] if len(sdnq_modules_to_not_convert) > 0: modules_to_not_convert.extend(sdnq_modules_to_not_convert) @@ -524,14 +517,9 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weigh except Exception as e: log.warning(f'Quantization: SDNQ failed to parse sdnq_modules_dtype_dict: {e}') - model.eval() - backup_embeddings = None - if hasattr(model, "get_input_embeddings"): - backup_embeddings = copy.deepcopy(model.get_input_embeddings()) - t0 = time.time() - model = apply_sdnq_to_module( + model = sdnq_post_load_quant( model, weights_dtype=weights_dtype, torch_dtype=devices.dtype, @@ -550,29 +538,8 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weigh op=op, ) - from modules.sdnq import SDNQConfig - model.quantization_config = SDNQConfig( - weights_dtype=weights_dtype, - group_size=shared.opts.sdnq_quantize_weights_group_size, - svd_rank=shared.opts.sdnq_svd_rank, - use_svd=shared.opts.sdnq_use_svd, - 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, - non_blocking=shared.opts.diffusers_offload_nonblocking, - quantization_device=quantization_device, - return_device=return_device, - modules_to_not_convert=modules_to_not_convert, - modules_dtype_dict=modules_dtype_dict.copy(), - ) - t1 = time.time() timer.load.add('sdnq', t1 - t0) - model.quantization_method = 'SDNQ' - - if hasattr(model, "set_input_embeddings") and backup_embeddings is not None: - model.set_input_embeddings(backup_embeddings) if op is not None and shared.opts.sdnq_quantize_shuffle_weights: if quant_last_model_name is not None: @@ -628,12 +595,14 @@ def optimum_quanto_model(model, op=None, sd_model=None, weights=None, activation from modules import devices, shared quanto = load_quanto('Quantize model: type=Optimum Quanto') global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement - if sd_model is not None and ("Flux" in sd_model.__class__.__name__ or "Chroma" in sd_model.__class__.__name__): # LayerNorm is not supported + if model.__class__.__name__ in {"FluxTransformer2DModel", "ChromaTransformer2DModel"}: # LayerNorm is not supported exclude_list = ["transformer_blocks.*.norm1.norm", "transformer_blocks.*.norm2", "transformer_blocks.*.norm1_context.norm", "transformer_blocks.*.norm2_context", "single_transformer_blocks.*.norm.norm", "norm_out.norm"] - if "Chroma" in sd_model.__class__.__name__: + if model.__class__.__name__ == "ChromaTransformer2DModel": # we ignore the distilled guidance layer because it degrades quality too much # see: https://github.com/huggingface/diffusers/pull/11698#issuecomment-2969717180 for more details exclude_list.append("distilled_guidance_layer.*") + elif model.__class__.__name__ == "QwenImageTransformer2DModel": + exclude_list = ["transformer_blocks.0.img_mod.1.weight", "time_text_embed", "img_in", "txt_in", "proj_out", "norm_out", "pos_embed"] else: exclude_list = None weights = getattr(quanto, weights) if weights is not None else getattr(quanto, shared.opts.optimum_quanto_weights_type) diff --git a/modules/sdnq/__init__.py b/modules/sdnq/__init__.py index 121db2c8d..c67bf0db7 100644 --- a/modules/sdnq/__init__.py +++ b/modules/sdnq/__init__.py @@ -1,12 +1,14 @@ -from .quantizer import SDNQConfig, SDNQQuantizer, apply_sdnq_to_module, sdnq_quantize_layer +from .quantizer import QuantizationMethod, SDNQConfig, SDNQQuantizer, sdnq_post_load_quant, apply_sdnq_to_module, sdnq_quantize_layer from .loader import save_sdnq_model, load_sdnq_model __all__ = [ + "QuantizationMethod", "SDNQConfig", "SDNQQuantizer", "apply_sdnq_to_module", "load_sdnq_model", "save_sdnq_model", + "sdnq_post_load_quant", "sdnq_quantize_layer", ] diff --git a/modules/sdnq/loader.py b/modules/sdnq/loader.py index 4ba68a69f..582c6b4b3 100644 --- a/modules/sdnq/loader.py +++ b/modules/sdnq/loader.py @@ -5,7 +5,7 @@ from safetensors import safe_open from diffusers.models.modeling_utils import ModelMixin from .common import use_tensorwise_fp8_matmul, use_contiguous_mm -from .quantizer import SDNQConfig, apply_sdnq_to_module +from .quantizer import SDNQConfig, sdnq_post_load_quant from .dequantizer import dequantize_symmetric, re_quantize_int8, re_quantize_fp8 @@ -66,7 +66,7 @@ def load_sdnq_model(model_path: str, model_cls: ModelMixin = None, file_name: st model = model_cls._from_config(config) # pylint: disable=protected-access else: raise ValueError(f"Dont know how to load model for {model_cls}") - model = apply_sdnq_to_module(model, **quantization_config) + model = sdnq_post_load_quant(model, **quantization_config) state_dict = {} if file_name: @@ -81,7 +81,6 @@ def load_sdnq_model(model_path: str, model_cls: ModelMixin = None, file_name: st for k in f.keys(): state_dict[k] = f.get_tensor(k) model.load_state_dict(state_dict, assign=True) - model.quantization_method = "sdnq" del state_dict if dtype is not None or dequantize_fp32 is not None or use_quantized_matmul is not None: model = apply_options_to_model(model, dtype=dtype, dequantize_fp32=dequantize_fp32, use_quantized_matmul=use_quantized_matmul) diff --git a/modules/sdnq/quantizer.py b/modules/sdnq/quantizer.py index b3b853e97..3fcf33ea7 100644 --- a/modules/sdnq/quantizer.py +++ b/modules/sdnq/quantizer.py @@ -75,7 +75,7 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz result_shape = None original_shape = layer.weight.shape if torch_dtype is None: - torch_dtype = devices.dtype + torch_dtype = layer.weight.dtype if layer_class_name in conv_types: if not quant_conv: @@ -321,6 +321,66 @@ def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_si return model +def sdnq_post_load_quant(model, weights_dtype="int8", torch_dtype=None, group_size=0, svd_rank=32, use_svd=False, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, dequantize_fp32=False, non_blocking=False, quantization_device=None, return_device=None, modules_to_not_convert: List[str] = None, modules_dtype_dict: Dict[str, List[str]] = None, op=None): # pylint: disable=unused-argument + model.eval() + if modules_to_not_convert is None: + modules_to_not_convert = [] + if modules_dtype_dict is None: + modules_dtype_dict = {} + + if getattr(model, "_keep_in_fp32_modules", None) is not None: + modules_to_not_convert.extend(model._keep_in_fp32_modules) # pylint: disable=protected-access + if getattr(model, "_skip_layerwise_casting_patterns", None) is not None: + modules_to_not_convert.extend(model._skip_layerwise_casting_patterns) # pylint: disable=protected-access + if model.__class__.__name__ == "ChromaTransformer2DModel": + modules_to_not_convert.append("distilled_guidance_layer") + elif model.__class__.__name__ == "QwenImageTransformer2DModel": + modules_to_not_convert.extend(["transformer_blocks.0.img_mod.1.weight", "time_text_embed", "img_in", "txt_in", "proj_out", "norm_out", "pos_embed"]) + if "minimum_6bit" not in modules_dtype_dict.keys(): + modules_dtype_dict["minimum_6bit"] = ["img_mod"] + else: + modules_dtype_dict["minimum_6bit"].append("img_mod") + + model = apply_sdnq_to_module( + model, + weights_dtype=weights_dtype, + torch_dtype=torch_dtype, + group_size=group_size, + svd_rank=svd_rank, + use_svd=use_svd, + quant_conv=quant_conv, + use_quantized_matmul=use_quantized_matmul, + use_quantized_matmul_conv=use_quantized_matmul_conv, + dequantize_fp32=dequantize_fp32, + non_blocking=non_blocking, + quantization_device=quantization_device, + return_device=return_device, + modules_to_not_convert=modules_to_not_convert, + modules_dtype_dict=modules_dtype_dict.copy(), + op=op, + ) + + model.quantization_config = SDNQConfig( + weights_dtype=weights_dtype, + group_size=group_size, + svd_rank=svd_rank, + use_svd=use_svd, + quant_conv=quant_conv, + use_quantized_matmul=use_quantized_matmul, + use_quantized_matmul_conv=use_quantized_matmul_conv, + dequantize_fp32=dequantize_fp32, + non_blocking=non_blocking, + quantization_device=quantization_device, + return_device=return_device, + modules_to_not_convert=modules_to_not_convert, + modules_dtype_dict=modules_dtype_dict.copy(), + ) + + model.quantization_method = QuantizationMethod.SDNQ + + return model + + class SDNQQuantizer(DiffusersQuantizer, HfQuantizer): r""" Diffusers Quantizer for SDNQ @@ -384,6 +444,8 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer): *args, **kwargs, # pylint: disable=unused-argument ): weights_dtype = self.quantization_config.weights_dtype + torch_dtype = param_value.dtype if self.torch_dtype is None else self.torch_dtype + if len(self.quantization_config.modules_dtype_dict.keys()) > 0: split_param_name = param_name.split(".") for key, value in self.quantization_config.modules_dtype_dict.items(): @@ -424,7 +486,7 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer): layer = sdnq_quantize_layer( layer, weights_dtype=weights_dtype, - torch_dtype=self.torch_dtype, + torch_dtype=torch_dtype, group_size=self.quantization_config.group_size, svd_rank=self.quantization_config.svd_rank, use_svd=self.quantization_config.use_svd, @@ -446,8 +508,6 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer): return dtype_dict[self.quantization_config.weights_dtype]["target_dtype"] def update_torch_dtype(self, torch_dtype: torch.dtype = None) -> torch.dtype: - if torch_dtype is None: - torch_dtype = devices.dtype self.torch_dtype = torch_dtype return torch_dtype diff --git a/pipelines/model_qwen.py b/pipelines/model_qwen.py index 6ed90bb15..d9a98d9be 100644 --- a/pipelines/model_qwen.py +++ b/pipelines/model_qwen.py @@ -49,7 +49,8 @@ def load_qwen(checkpoint_info, diffusers_load_config={}): subfolder=transformer_subfolder, cls_name=diffusers.QwenImageTransformer2DModel, load_config=diffusers_load_config, - modules_dtype_dict={"minimum_8bit": ["pos_embed", "time_text_embed", "img_in", "txt_in", "norm_out", "img_mod", "transformer_blocks.0.img_mod.1.weight"]}, + modules_dtype_dict={"minimum_6bit": ["img_mod"]}, + modules_to_not_convert=["transformer_blocks.0.img_mod.1.weight", "time_text_embed", "img_in", "txt_in", "proj_out", "norm_out", "pos_embed"], ) repo_te = 'Qwen/Qwen-Image'