From ebb26ac123c548981c5a9b7d38d739a3e217da5a Mon Sep 17 00:00:00 2001 From: Disty0 Date: Mon, 6 Oct 2025 01:04:00 +0300 Subject: [PATCH] SDNQ make load file name configurable --- modules/sdnq/loader.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/modules/sdnq/loader.py b/modules/sdnq/loader.py index 3473e70ad..fc2ed831c 100644 --- a/modules/sdnq/loader.py +++ b/modules/sdnq/loader.py @@ -14,7 +14,7 @@ def save_sdnq_model(model: ModelMixin, sdnq_config: SDNQConfig, model_path: str) sdnq_config.to_json_file(os.path.join(model_path, "quantization_config.json")) -def load_sdnq_model(model_cls: ModelMixin, model_path: str, use_quantized_matmul: bool = False) -> ModelMixin: +def load_sdnq_model(model_cls: ModelMixin, model_path: str, file_name: str = "diffusion_pytorch_model.safetensors", use_quantized_matmul: bool = False) -> ModelMixin: with torch.device("meta"): with open(os.path.join(model_path, "quantization_config.json"), "r") as f: quantization_config = json.load(f) @@ -28,7 +28,7 @@ def load_sdnq_model(model_cls: ModelMixin, model_path: str, use_quantized_matmul model = apply_sdnq_to_module(model, **quantization_config) state_dict = {} - with safe_open(os.path.join(model_path, "diffusion_pytorch_model.safetensors"), framework="pt") as f: + with safe_open(os.path.join(model_path, file_name), framework="pt") as f: for k in f.keys(): state_dict[k] = f.get_tensor(k) model.load_state_dict(state_dict, assign=True)