From fc547a3ccdf330714b1d0773f1ca7137c4102bd8 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Mon, 18 Aug 2025 10:28:41 -0400 Subject: [PATCH] sdnq with diffusers lora loader Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 10 +++++++--- modules/lora/lora_load.py | 2 +- modules/lora/lora_overrides.py | 8 +++++++- modules/sdnq/__init__.py | 11 ++++++++++- pipelines/flux/flux_lora.py | 27 +++++++++++++++++++++++++++ pipelines/model_flux.py | 3 +++ 6 files changed, 55 insertions(+), 6 deletions(-) create mode 100644 pipelines/flux/flux_lora.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 460b7979b..15297cafc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,11 +1,12 @@ # Change Log for SD.Next -## Update for 2025-08-16 +## Update for 2025-08-18 -- **Features** +- **Features** - new setting -> huggingface -> download method default is `rust` as new `xet` is known to cause issues -- **Fixes** + - support for `flux.1-kontext` lora +- **Fixes** - fix OpenVINO with offloading - add explicit offload calls on prompt encode - error reporting on model load failure @@ -13,6 +14,9 @@ - remove extra cache clear - enable explicit sync calls for `rocm` on windows - note restart-needed on initial startup import error + - bypass diffusers-lora-fuse on quantized models + - monkey-patch diffusers to use original weights shape when loading lora + - guard against null prompt ## Update for 2025-08-15 diff --git a/modules/lora/lora_load.py b/modules/lora/lora_load.py index 0ad032c15..0f2ef570c 100644 --- a/modules/lora/lora_load.py +++ b/modules/lora/lora_load.py @@ -314,7 +314,7 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non shared.log.trace(f'Network load: type=LoRA list={shared.sd_model.get_list_adapters()}') shared.log.trace(f'Network load: type=LoRA active={shared.sd_model.get_active_adapters()}') shared.sd_model.set_adapters(adapter_names=diffuser_loaded, adapter_weights=diffuser_scales) - if shared.opts.lora_fuse_diffusers and not lora_overrides.check_fuse(): + if shared.opts.lora_fuse_diffusers and not lora_overrides.disable_fuse(): shared.sd_model.fuse_lora(adapter_names=diffuser_loaded, lora_scale=1.0, fuse_unet=True, fuse_text_encoder=True) # diffusers with fuse uses fixed scale since later apply does the scaling shared.sd_model.unload_lora_weights() l.timer.activate += time.time() - t1 diff --git a/modules/lora/lora_overrides.py b/modules/lora/lora_overrides.py index 9b51b62e2..d914f3793 100644 --- a/modules/lora/lora_overrides.py +++ b/modules/lora/lora_overrides.py @@ -47,6 +47,7 @@ force_models_diffusers = [ # forced always ] force_classes_diffusers = [ # forced always + 'FluxKontextPipeline', 'FluxKontextInpaintPipeline', ] fuse_ignore = [ @@ -68,5 +69,10 @@ def get_method(shorthash=''): else: return 'native' -def check_fuse(): + +def disable_fuse(): + if hasattr(shared.sd_model, 'quantization_config'): + return True + if hasattr(shared.sd_model, 'transformer') and hasattr(shared.sd_model.transformer, 'quantization_config'): + return True return shared.sd_model_type in fuse_ignore diff --git a/modules/sdnq/__init__.py b/modules/sdnq/__init__.py index bca210c1b..6eb2e3553 100644 --- a/modules/sdnq/__init__.py +++ b/modules/sdnq/__init__.py @@ -245,6 +245,15 @@ def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_si return model +class SDNQParameter(torch.nn.Parameter): + def __new__(cls, data=None, requires_grad=False): + return super().__new__(cls, data, requires_grad) + + def __init__(self, data=None, requires_grad=False): # pylint: disable=unused-argument + self.original_shape = data.shape + super().__init__() + + class SDNQQuantizer(DiffusersQuantizer): r""" Diffusers Quantizer for SDNQ @@ -333,7 +342,7 @@ class SDNQQuantizer(DiffusersQuantizer): param_value = param_value.to(target_device, non_blocking=self.quantization_config.non_blocking).to(dtype=torch.float32) layer, _ = get_module_from_name(model, param_name) - layer.weight = torch.nn.Parameter(param_value, requires_grad=False) + layer.weight = SDNQParameter(param_value, requires_grad=False) layer = sdnq_quantize_layer( layer, weights_dtype=weights_dtype, diff --git a/pipelines/flux/flux_lora.py b/pipelines/flux/flux_lora.py new file mode 100644 index 000000000..50f138df5 --- /dev/null +++ b/pipelines/flux/flux_lora.py @@ -0,0 +1,27 @@ +def calculate_module_shape(model, base_module=None, base_weight_param_name=None): + def _get_weight_shape(weight): + if weight.__class__.__name__ == "Params4bit": + return weight.quant_state.shape + elif weight.__class__.__name__ == "GGUFParameter": + return weight.quant_shape + elif weight.__class__.__name__ == "SDNQParameter": + return weight.original_shape + else: + return weight.shape + + if base_module is not None: + return _get_weight_shape(base_module.weight) + elif base_weight_param_name is not None: + from diffusers.utils import get_submodule_by_name + if not base_weight_param_name.endswith(".weight"): + raise ValueError(f"Invalid `base_weight_param_name` passed as it does not end with '.weight' {base_weight_param_name=}.") + module_path = base_weight_param_name.rsplit(".weight", 1)[0] + submodule = get_submodule_by_name(model, module_path) + return _get_weight_shape(submodule.weight) + + raise ValueError("Either `base_module` or `base_weight_param_name` must be provided.") + + +def apply_patch(): + from diffusers.loaders.lora_pipeline import FluxLoraLoaderMixin + FluxLoraLoaderMixin._calculate_module_shape = calculate_module_shape # pylint: disable=protected-access diff --git a/pipelines/model_flux.py b/pipelines/model_flux.py index 733b0235f..ccede0948 100644 --- a/pipelines/model_flux.py +++ b/pipelines/model_flux.py @@ -23,6 +23,9 @@ def load_flux(checkpoint_info, diffusers_load_config={}): else: cls_name = diffusers.FluxPipeline + from pipelines.flux import flux_lora + flux_lora.apply_patch() + load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) shared.log.debug(f'Load model: type=Flux repo="{repo_id}" cls={cls_name.__name__} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}')