Bugfix lorahotload

diffsynth/core/loader/model.py

@@ -119,6 +119,9 @@ def load_model(model_class, path, config=None, torch_dtype=torch.bfloat16, devic
         # Because some models override the behavior of `to()`,
         # especially those from libraries like Transformers.
         model = model.to(dtype=torch_dtype, device=device)
+    if quantize is not None:
+        # Downstream steps (e.g. LoRA hot-loading) need the config to handle the quantized layers.
+        model.quantize_config = quantize
     if hasattr(model, "eval"):
         model = model.eval()
     return model

diffsynth/core/vram/layers.py

@@ -442,7 +442,7 @@ def forward(self, x, *args, **kwargs):
         return out
 
 
-class AutoWrappedQuantizedModule(AutoTorchModule, LoRAHotLoadMixin):
+class AutoWrappedQuantizedModule(torch.nn.Linear, AutoTorchModule, LoRAHotLoadMixin):
     def __init__(
         self,
         module: torch.nn.Module,
@@ -467,7 +467,8 @@ def __init__(
                 "Disk offload for quantized layers requires both `disk_map` and `quantize`, "
                 "so each layer can rebuild its packed weight and quant state lazily."
             )
-        super().__init__(
+        AutoTorchModule.__init__(
+            self,
             offload_dtype,
             offload_device,
             onload_dtype,
@@ -478,6 +479,8 @@ def __init__(
             computation_device,
             vram_limit,
         )
+        self.in_features = module.in_features
+        self.out_features = module.out_features
         self.module = module
         self.name = name
         self.disk_offload = disk_offload

diffsynth/diffusion/base_pipeline.py

@@ -240,6 +240,13 @@ def flush_vram_management_device(self, device):
                 module.computation_device = device
 
 
+    def check_quant_hot_load(self, module: torch.nn.Module):
+        # A quantized weight cannot absorb a fused LoRA, so hot-loading is the only option for it.
+        if getattr(module, "quantize_config", None) is None:
+            return module
+        return self.enable_lora_hot_loading(module)
+
+
     def load_lora(
         self,
         module: torch.nn.Module,
@@ -249,6 +256,7 @@ def load_lora(
         state_dict=None,
         verbose=1,
     ):
+        module = self.check_quant_hot_load(module)
         if state_dict is None:
             if isinstance(lora_config, str):
                 lora = load_state_dict(lora_config, torch_dtype=self.torch_dtype, device=self.device)
@@ -370,7 +378,7 @@ def enable_lora_hot_loading(self, model: torch.nn.Module):
             "computation_dtype": self.torch_dtype,
             "computation_device": self.device,
         }
-        model = enable_vram_management(model, module_map, vram_config=vram_config)
+        model = enable_vram_management(model, module_map, vram_config=vram_config, quantize=getattr(model, "quantize_config", None))
         return model
 
     def compile_pipeline(self, mode: str = "default", dynamic: bool = True, fullgraph: bool = False, compile_models: list = None, **kwargs):