diff --git a/diffsynth/core/loader/model.py b/diffsynth/core/loader/model.py index 64a22cca6..340e87aa1 100644 --- a/diffsynth/core/loader/model.py +++ b/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 diff --git a/diffsynth/core/vram/layers.py b/diffsynth/core/vram/layers.py index fde062426..716d1d09b 100644 --- a/diffsynth/core/vram/layers.py +++ b/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 diff --git a/diffsynth/diffusion/base_pipeline.py b/diffsynth/diffusion/base_pipeline.py index 907c731b2..5af48a110 100644 --- a/diffsynth/diffusion/base_pipeline.py +++ b/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):