print model components for safetensors load

Signed-off-by: Vladimir Mandic <mandic00@live.com>

modules/model_tools.py

@@ -5,13 +5,32 @@ import safetensors.torch
 from modules import shared, devices, model_quant
 
 
+def remove_entries_after_depth(d, depth, current_depth=0):
+    if current_depth >= depth:
+        return None
+    if isinstance(d, dict):
+        return {k: remove_entries_after_depth(v, depth, current_depth + 1) for k, v in d.items() if remove_entries_after_depth(v, depth, current_depth + 1) is not None}
+    return d
+
+
+def list_to_dict(flat_list):
+    result_dict = {}
+    for item in flat_list:
+        keys = item.split('.')
+        d = result_dict
+        for key in keys[:-1]:
+            d = d.setdefault(key, {})
+        d[keys[-1]] = None
+    return result_dict
+
+
 def get_safetensor_keys(filename):
     keys = []
     try:
         with safetensors.torch.safe_open(filename, framework="pt", device="cpu") as f:
             keys = f.keys()
-    except Exception as e:
-        shared.log.error(f'Load dict: path="{filename}" {e}')
+    except Exception:
+        pass
     return keys
 
 

modules/sd_detect.py

@@ -1,7 +1,8 @@
 import os
+import time
 import torch
 import diffusers
-from modules import shared, shared_items, devices, errors
+from modules import shared, shared_items, devices, errors, model_tools
 
 
 debug_load = os.environ.get('SD_LOAD_DEBUG', None)
@@ -103,6 +104,13 @@ def detect_pipeline(f: str, op: str = 'model', warning=True, quiet=False):
             pipeline = shared_items.get_pipelines().get(guess, None) if pipeline is None else pipeline
             if not quiet:
                 shared.log.info(f'Autodetect {op}: detect="{guess}" class={getattr(pipeline, "__name__", None)} file="{f}" size={size}MB')
+                t0 = time.time()
+                keys = model_tools.get_safetensor_keys(f)
+                if keys is not None:
+                    modules = model_tools.list_to_dict(keys)
+                    modules = model_tools.remove_entries_after_depth(modules, 3)
+                    t1 = time.time()
+                    shared.log.debug(f'Autodetect {op}: modules={modules} time={t1-t0:.2f}')
         except Exception as e:
             shared.log.error(f'Autodetect {op}: file="{f}" {e}')
             if debug_load: