sdnq with diffusers lora loader

Signed-off-by: Vladimir Mandic <mandic00@live.com>

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
 

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

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

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,

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

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}')