From f1841269a07ff6a544f6b1db7db54d8a9b583f56 Mon Sep 17 00:00:00 2001
From: "Salvador E. Tropea" <stropea@inti.gob.ar>
Date: Sun, 6 Jul 2025 15:05:41 -0300
Subject: [PATCH] [Load Safetensors][Added] Support for multiple models in one
 file

"Bag of Models" in Demucs tongue
---
 source/utils/load_safetensors.py | 38 +++++++++++++++++++++++++++-----
 1 file changed, 33 insertions(+), 5 deletions(-)

diff --git a/source/utils/load_safetensors.py b/source/utils/load_safetensors.py
index 4d19dc6..e6e19ad 100644
--- a/source/utils/load_safetensors.py
+++ b/source/utils/load_safetensors.py
@@ -12,11 +12,7 @@ from .misc import NODES_NAME
 logger = logging.getLogger(f"{NODES_NAME}.load_safetensors")
 
 
-def load_safetensors(model_path, model_run, device):
-    logger.info("Loading PyTorch model from .safetensors file...")
-    # 1. Load the state_dict from the file, EXPLICITLY forcing all tensors onto the CPU.
-    state_dict = load_file(model_path, device="cpu")
-    # 2. Load the CPU state_dict into the CPU model. This is now a safe operation.
+def load_state_dict(model_run, state_dict):
     try:
         missing_keys, unexpected_keys = model_run.load_state_dict(state_dict, strict=False)
         if missing_keys:
@@ -29,4 +25,36 @@ def load_safetensors(model_path, model_run, device):
         logger.error(f"RuntimeError during model_run.load_state_dict: {e}")
         logger.error("This might indicate a mismatch between saved weights and model architecture.")
         raise
+
+
+def load_safetensors(model_path, model_run, device):
+    logger.info("Loading PyTorch model from .safetensors file...")
+    # 1. Load the state_dict from the file, EXPLICITLY forcing all tensors onto the CPU.
+    state_dict = load_file(model_path, device="cpu")
+
+    # 2. Load the CPU state_dict into the CPU model. This is now a safe operation.
+    if hasattr(model_run, 'signatures'):
+        # This is the way we store Demucs models
+        signatures = model_run.signatures
+        if len(signatures) == 1:
+            # Single model
+            sig = signatures[0]
+            logger.debug(f"Single model with filtered keys, prefix: {sig}")
+            # Just remove the prefix
+            state_dict = {k[len(sig)+1:]: v for k, v in state_dict.items()}
+            # The rest is as a regular model
+        else:
+            # Bag of models, load the keys for each sub-model
+            logger.debug("Multiple models with filtered keys")
+            for sub_model, sig in zip(model_run.models, signatures):
+                prefix = f"{sig}."
+                logger.debug(f"  - Filtering {prefix}")
+                sub_state_dict = {k[len(prefix):]: v for k, v in state_dict.items() if k.startswith(prefix)}
+                load_state_dict(sub_model, sub_state_dict)
+            # Finished
+            return model_run
+
+    load_state_dict(model_run, state_dict)
+    model_run.target_device = device
+
     return model_run