From 962cb7115d33ab863caf3b12ad4caab6d4abfdcf Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Tue, 7 Oct 2025 14:30:45 -0400 Subject: [PATCH] infra for full-model load/save with quant Signed-off-by: Vladimir Mandic --- modules/prompt_parser_diffusers.py | 3 ++ modules/sd_models.py | 28 ++++++++++++++--- modules/sd_offload.py | 1 + modules/sdnq/loader.py | 49 +++++++++++++++++++++++++----- 4 files changed, 70 insertions(+), 11 deletions(-) diff --git a/modules/prompt_parser_diffusers.py b/modules/prompt_parser_diffusers.py index 31947ee7b..571bf70d1 100644 --- a/modules/prompt_parser_diffusers.py +++ b/modules/prompt_parser_diffusers.py @@ -589,6 +589,9 @@ def get_weighted_text_embeddings(pipe, prompt: str = "", neg_prompt: str = "", c negative_weights.pop(0) embedding_providers = prepare_embedding_providers(pipe, clip_skip) + if len(embedding_providers) == 0: + shared.log.error("Prompt encode: cannot find text encoder in model") + return None, None, None, None, None, None empty_embedding_providers = None if 'StableCascade' in pipe.__class__.__name__: empty_embedding_providers = [embedding_providers[1]] diff --git a/modules/sd_models.py b/modules/sd_models.py index 449e26772..c8f78c1ae 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -542,8 +542,27 @@ def load_diffuser_file(model_type, pipeline, checkpoint_info, diffusers_load_con def load_sdnq_model(checkpoint_info, pipeline, diffusers_load_config, op): - shared.log.error(f'Load {op}: model="{checkpoint_info.name}" cls={pipeline.__name__} args={diffusers_load_config} SDNQ pre-quant not supported') - return None + from modules import sdnq + modules = {} + for f in os.listdir(checkpoint_info.path): + if not f.endswith('quantization_config.json'): + continue + module_name = f.replace('_quantization_config.json', '') + quantization_config = shared.readfile(os.path.join(checkpoint_info.path, f), silent=True) + shared.log.debug(f'Load {op}: model="{checkpoint_info.name}" module="{module_name}" prequant=sdnq') + module_path = os.path.join(checkpoint_info.path, module_name) + modules[module_name] = sdnq.load_sdnq_model( + model_path=module_path, + quantization_config=quantization_config, + ) + modules[module_name] = modules[module_name].to(device=devices.device) + sd_model = pipeline.from_pretrained( + checkpoint_info.path, + cache_dir=shared.opts.diffusers_dir, + **modules, + **diffusers_load_config, + ) + return sd_model def set_overrides(sd_model, checkpoint_info): @@ -665,8 +684,9 @@ def load_diffuser(checkpoint_info=None, op='model', revision=None): # pylint: di # load sdnq-prequantized model if sd_model is None: if model_type.endswith('SDNQ'): - allow_post_quant = False sd_model = load_sdnq_model(checkpoint_info, pipeline, diffusers_load_config, op) + allow_post_quant = False + model_type = model_type.replace(' SDNQ', '') # load from hf folder-style if sd_model is None: @@ -1056,7 +1076,7 @@ def set_diffusers_attention(pipe, quiet:bool=False): if attention is None: return # other models uses their own attention processor - if pipe.__class__.__name__.startswith("StableDiffusion") and hasattr(pipe, "unet"): + if pipe.__class__.__name__.startswith("StableDiffusion") and getattr(pipe, "unet", None) is not None and hasattr(pipe.unet, "set_attn_processor"): try: pipe.unet.set_attn_processor(attention) except Exception as e: diff --git a/modules/sd_offload.py b/modules/sd_offload.py index 54446abdf..c52220ccb 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -434,6 +434,7 @@ def apply_balanced_offload(sd_model=None, exclude=[]): module.offload_dir = os.path.join(shared.opts.accelerate_offload_path, checkpoint_name, module_name) apply_balanced_offload_to_module(module, op='apply') report_model_stats(module_name, module) + set_accelerate(sd_model) t = time.time() - t0 process_timer.add('offload', t) diff --git a/modules/sdnq/loader.py b/modules/sdnq/loader.py index 86f67d372..4ba68a69f 100644 --- a/modules/sdnq/loader.py +++ b/modules/sdnq/loader.py @@ -29,24 +29,59 @@ def save_sdnq_model(model: ModelMixin, model_path: str, max_shard_size: str = "1 module.quantization_config.to_json_file(os.path.join(model_path, f"{module_name}_quantization_config.json")) -def load_sdnq_model(model_cls: ModelMixin, model_path: str, file_name: str = "diffusion_pytorch_model.safetensors", dtype: torch.dtype = None, dequantize_fp32: bool = None, use_quantized_matmul: bool = None) -> ModelMixin: - with torch.device("meta"): - with open(os.path.join(model_path, "quantization_config.json"), "r", encoding="utf-8") as f: - quantization_config = json.load(f) +def load_sdnq_model(model_path: str, model_cls: ModelMixin = None, file_name: str = None, dtype: torch.dtype = None, dequantize_fp32: bool = None, use_quantized_matmul: bool = None, model_config: dict = None, quantization_config: dict = None) -> ModelMixin: + with torch.device("cpu"): + if quantization_config is None: + try: + with open(os.path.join(model_path, "quantization_config.json"), "r", encoding="utf-8") as f: + quantization_config = json.load(f) + except Exception: + quantization_config = {} + if model_config is None: + try: + with open(os.path.join(model_path, 'config.json'), "r", encoding="utf-8") as f: + model_config = json.load(f) + except Exception: + model_config = {} + if model_cls is None: + import transformers + import diffusers + class_name = model_config.get("_class_name", None) or model_config.get("architectures", None) + if isinstance(class_name, list): + class_name = class_name[0] + if class_name is not None: + model_cls = getattr(diffusers, class_name, None) or getattr(transformers, class_name, None) + if model_cls is None: + raise ValueError(f"Cannot determine model class for {model_path}, please provide model_cls argument") quantization_config.pop("is_integer", None) quantization_config.pop("quant_method", None) quantization_config.pop("quantization_device", None) quantization_config.pop("return_device", None) quantization_config.pop("non_blocking", None) - config = model_cls.load_config(model_path) - model = model_cls.from_config(config) + if hasattr(model_cls, "load_config"): + config = model_cls.load_config(model_path) + model = model_cls.from_config(config) + elif hasattr(model_cls, "_from_config"): + config = transformers.PretrainedConfig.from_dict(model_config) + model = model_cls._from_config(config) # pylint: disable=protected-access + else: + raise ValueError(f"Dont know how to load model for {model_cls}") model = apply_sdnq_to_module(model, **quantization_config) state_dict = {} - with safe_open(os.path.join(model_path, file_name), framework="pt") as f: + if file_name: + fn = os.path.join(model_path, file_name) + elif os.path.exists(os.path.join(model_path, "diffusion_pytorch_model.safetensors")): + fn = os.path.join(model_path, "diffusion_pytorch_model.safetensors") + elif os.path.exists(os.path.join(model_path, "model.safetensors")): + fn = os.path.join(model_path, "model.safetensors") + else: + raise ValueError(f"Cannot find safetensors file in {model_path}, please provide file_name argument") + with safe_open(fn, framework="pt") as f: for k in f.keys(): state_dict[k] = f.get_tensor(k) model.load_state_dict(state_dict, assign=True) + model.quantization_method = "sdnq" del state_dict if dtype is not None or dequantize_fp32 is not None or use_quantized_matmul is not None: model = apply_options_to_model(model, dtype=dtype, dequantize_fp32=dequantize_fp32, use_quantized_matmul=use_quantized_matmul)