SDNQ keep the quant configs inside the module subfolder, add dtype cast and don't send to GPU

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__:

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}')

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: