From e19fb2d8337b8e7fb4fce409c66ea9f6b7e18d60 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Thu, 9 Oct 2025 19:34:48 +0300 Subject: [PATCH] SDNQ keep the quant configs inside the module subfolder, add dtype cast and don't send to GPU --- modules/sd_detect.py | 14 +++++++++++++- modules/sd_models.py | 15 ++++++++------- modules/sdnq/loader.py | 19 ++++++++++++++----- 3 files changed, 35 insertions(+), 13 deletions(-) diff --git a/modules/sd_detect.py b/modules/sd_detect.py index 059b081a4..7d3e20c2e 100644 --- a/modules/sd_detect.py +++ b/modules/sd_detect.py @@ -138,7 +138,19 @@ def guess_by_diffusers(fn, current_guess): if pipeline is None: pipeline = cls if callable(pipeline): - is_quant = any(f for f in os.listdir(fn) if f.endswith('quantization_config.json')) + is_quant = False + for folder in os.listdir(fn): + folder = os.path.join(fn, folder) + if is_quant: + break + if folder.endswith('quantization_config.json'): + is_quant = True + break + elif os.path.isdir(folder): + for f in os.listdir(folder): + if f.endswith('quantization_config.json'): + is_quant = True + break pipelines = shared_items.get_pipelines() for k, v in pipelines.items(): if v is not None and v.__name__ == pipeline.__name__: diff --git a/modules/sd_models.py b/modules/sd_models.py index b0606dd76..4400a0aa6 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -544,20 +544,20 @@ def load_diffuser_file(model_type, pipeline, checkpoint_info, diffusers_load_con def load_sdnq_model(checkpoint_info, pipeline, diffusers_load_config, op): from modules import sdnq modules = {} - for f in os.listdir(checkpoint_info.path): - if not f.endswith('quantization_config.json'): + for module_name in os.listdir(checkpoint_info.path): + quantization_config_path = os.path.join(checkpoint_info.path, module_name, 'quantization_config.json') + if not os.path.exists(quantization_config_path): continue - module_name = f.replace('_quantization_config.json', '') - quantization_config = shared.readfile(os.path.join(checkpoint_info.path, f), silent=True) + model_path = os.path.join(checkpoint_info.path, module_name) + quantization_config = shared.readfile(quantization_config_path, silent=True) shared.log.debug(f'Load {op}: model="{checkpoint_info.name}" module="{module_name}" direct={shared.opts.diffusers_to_gpu} prequant=sdnq') - module_path = os.path.join(checkpoint_info.path, module_name) try: modules[module_name] = sdnq.load_sdnq_model( - model_path=module_path, + model_path=model_path, quantization_config=quantization_config, device=devices.device if shared.opts.diffusers_to_gpu else devices.cpu, + dtype=devices.dtype, ) - modules[module_name] = modules[module_name].to(device=devices.device) except Exception as e: shared.log.error(f'Load {op}: model="{checkpoint_info.name}" module="{module_name}" {e}') errors.display(e, 'Load') @@ -1291,6 +1291,7 @@ def save_model(name: str, path: str = None, shard: str = None, overwrite: bool = model=shared.sd_model, model_path=model_name, max_shard_size=shard, + is_pipeline=True, ) t1 = time.time() shared.log.info(f'Save model: path="{model_name}" cls={shared.sd_model.__class__.__name__} time={t1 - t0:.2f}') diff --git a/modules/sdnq/loader.py b/modules/sdnq/loader.py index 8ed2761ef..242f1e83b 100644 --- a/modules/sdnq/loader.py +++ b/modules/sdnq/loader.py @@ -15,16 +15,25 @@ def get_module_names(model: ModelMixin) -> list: return modules_names -def save_sdnq_model(model: ModelMixin, model_path: str, max_shard_size: str = "10GB", sdnq_config: SDNQConfig = None) -> None: +def save_sdnq_model(model: ModelMixin, model_path: str, max_shard_size: str = "10GB", is_pipeline: bool = False, sdnq_config: SDNQConfig = None) -> None: model.save_pretrained(model_path, max_shard_size=max_shard_size) # actual save if sdnq_config is not None: # if provided, save global config sdnq_config.to_json_file(os.path.join(model_path, "quantization_config.json")) - for module_name in get_module_names(model): # save per-module config if available - module = getattr(model, module_name, None) - if (module is not None) and hasattr(module, "quantization_config") and isinstance(module.quantization_config, SDNQConfig): - module.quantization_config.to_json_file(os.path.join(model_path, f"{module_name}_quantization_config.json")) + if is_pipeline: + for module_name in get_module_names(model): # save per-module config if available + module = getattr(model, module_name, None) + if (module is not None) and hasattr(module, "quantization_config") and isinstance(module.quantization_config, SDNQConfig): + module.quantization_config.to_json_file(os.path.join(model_path, module_name, "quantization_config.json")) + elif sdnq_config is None: + quantization_config = None + if hasattr(model, "quantization_config"): + quantization_config = model.quantization_config + elif hasattr(model, "config") and hasattr(model.config, "quantization_config"): + quantization_config = model.config.quantization_config + if quantization_config is not None: + quantization_config.to_json_file(os.path.join(model_path, "quantization_config.json")) def load_sdnq_model(model_path: str, model_cls: ModelMixin = None, file_name: str = None, dtype: torch.dtype = None, device: torch.device = 'cpu', dequantize_fp32: bool = None, use_quantized_matmul: bool = None, model_config: dict = None, quantization_config: dict = None) -> ModelMixin: