[Added] Support for partial models created from bigger ones

I guess they are to test individual portions

.gitignore

@@ -12,4 +12,5 @@ models/all_public_uvr_models
 models/demucs
 __no__
 models/.catalog.csv
+models/*.yaml
 0LEEME

models/uvr_model_data.json

@@ -591,6 +591,24 @@
         "primary_stem": "Other",
         "stages": 5
     },
+    "65dd0f678ed827e4e434d97255639c38": {
+        "desc": "Hybrid Demucs MDX A Rep. TO",
+        "download": "Main/Demucs",
+        "file_t": "safetensors",
+        "is_bag_of_models": "true",
+        "model_t": "Demucs",
+        "name": "repro_mdx_a_time_only.yaml",
+        "params": "267506128",
+        "parent": "e2bccab7f9797842079390749eee8178",
+        "primary_stem": [
+            "Vocals",
+            "Drums",
+            "Bass",
+            "Other"
+        ],
+        "segment": "44",
+        "signatures": "[\"9a6b4851\", \"9a6b4851\", \"1ef250f1\", \"1ef250f1\"]"
+    },
     "6703e39f36f18aa7855ee1047765621d": {
         "channels": 32,
         "compensate": 1.035,
@@ -773,6 +791,24 @@
         "mdx_n_fft_scale_set": 5120,
         "primary_stem": "Instrumental"
     },
+    "8abca393bc679ca0336a0ba8ca0e732d": {
+        "desc": "Hybrid Demucs MDX A Rep. HO",
+        "download": "Main/Demucs",
+        "file_t": "safetensors",
+        "is_bag_of_models": "true",
+        "model_t": "Demucs",
+        "name": "repro_mdx_a_hybrid_only.yaml",
+        "params": "167275664",
+        "parent": "e2bccab7f9797842079390749eee8178",
+        "primary_stem": [
+            "Vocals",
+            "Drums",
+            "Bass",
+            "Other"
+        ],
+        "segment": "44",
+        "signatures": "[\"fa0cb7f9\", \"902315c2\", \"fa0cb7f9\", \"902315c2\"]"
+    },
     "8b15e58e9f33ee39346be7f636ad1d63": {
         "channels": 32,
         "compensate": 1.03,
@@ -1135,6 +1171,21 @@
         "primary_stem": "Bass",
         "stages": 5
     },
+    "c50d9f98739139dbc9fd2477931d686e": {
+        "desc": "UVR Hybrid Demucs 2",
+        "download": "Main/Demucs",
+        "file_t": "safetensors",
+        "model_t": "Demucs",
+        "name": "UVR_Demucs_Model_2.yaml",
+        "params": "83633212",
+        "parent": "e4e0c3604695bafc2227555268e37942",
+        "primary_stem": [
+            "Vocals",
+            "Other"
+        ],
+        "segment": "44",
+        "signatures": "[\"ebf34a2d\"]"
+    },
     "c7500d7fdb1c0fc24b14b698515462d2": {
         "config_yaml": "config_mdx23c_similarity.yaml",
         "is_karaoke": false,
@@ -1219,6 +1270,21 @@
         "primary_stem": "Instrumental",
         "stages": 5
     },
+    "ccdd9985f1f7cbe62bf1e1a3abfcbd3c": {
+        "desc": "UVR Hybrid Demucs 1",
+        "download": "Main/Demucs",
+        "file_t": "safetensors",
+        "model_t": "Demucs",
+        "name": "UVR_Demucs_Model_1.yaml",
+        "params": "83633212",
+        "parent": "e4e0c3604695bafc2227555268e37942",
+        "primary_stem": [
+            "Vocals",
+            "Other"
+        ],
+        "segment": "44",
+        "signatures": "[\"ebf34a2db\"]"
+    },
     "cd5b2989ad863f116c855db1dfe24e39": {
         "channels": 48,
         "compensate": 1.035,

source/db/hash_dir.py

@@ -80,7 +80,7 @@ def hash_dir(directory_path: str) -> dict:
         # 4. Filter by file size.
         try:
             file_size = os.path.getsize(file_path)
-            if file_size < MIN_FILE_SIZE_BYTES:
+            if file_size < MIN_FILE_SIZE_BYTES and not filename.lower().endswith('.yaml'):
                 logger.debug(f"Skipping small file: '{filename}' ({file_size / 1024**2:.2f}MB)")
                 continue
         except OSError as e:

source/db/load_model.py

@@ -32,6 +32,16 @@ def load_model(d, device, models_dir):
     if model_path is None:
         # It means it wasn't on disk
         model_path = download_model(d, models_dir)
+    # Is this a child model?
+    parent = d.get("parent")
+    if parent:
+        # Yes, we need the parent file
+        if not isinstance(parent, dict):
+            raise ValueError("Trying to load a broken child model")
+        model_path = parent.get('model_path')
+        if model_path is None:
+            # It means it wasn't on disk
+            model_path = download_model(parent, models_dir)
 
     # ONNX
     if file_t == "onnx":

source/db/models_db.py

@@ -218,6 +218,17 @@ def get_models_full(primary_stem=None, model_t=None, file_t=None, json_path=None
             d['model_path'] = downloaded[hash]
             found_disk[os.path.realpath(file_name)] = d
 
+    # Link the child models
+    for d in found.values():
+        parent_hash = d.get("parent")
+        if not parent_hash:
+            continue
+        parent_obj = found_hashes.get(parent_hash)
+        if not parent_obj:
+            logger.error(f"Inconsistency in database, model `{d['name']}` points to unknown hash `{parent_hash}`")
+            continue
+        d["parent"] = parent_obj
+
     return found, found_hashes, found_disk, def_sep + sorted(on_disk) + sorted(to_down) + sorted(on_disk_as_down)
 
 
@@ -339,6 +350,13 @@ class ModelsDB(object):
     def save(self):
         # Remove run-time information
         for k, v in self.models.items():
+            try:
+                # Make children just refer to parent's hash, not the actual parent
+                parent = v["parent"]
+                if isinstance(parent, dict):
+                    v["parent"] = parent["hash"]
+            except KeyError:
+                pass
             try:
                 del v["hash"]
             except KeyError:

source/inference/get_model.py

@@ -16,7 +16,7 @@ from .demucs_api import BagOfModels
 logger = logging.getLogger(f"{NODES_NAME}.get_model")
 
 
-def get_metadata(file_path):
+def get_metadata(file_path, d):
     """ Read the metadata from a safetensors file """
     logger.debug(f"Reading metadata from {file_path}")
     metadata = {}
@@ -24,6 +24,20 @@ def get_metadata(file_path):
         metadata = f.metadata()
     if not metadata:
         raise ValueError(f"Could not read metadata from safetensors file: {file_path}")
+
+    # Is this a child model?
+    parent = d.get('parent')
+    if parent:
+        # Ok, this is a child model changing details of a parent model
+        # Currently used by Demucs models to create simplified versions of the same model
+        metadata['is_bag_of_models'] = d.get('is_bag_of_models', 'false')
+        try:
+            metadata['signatures'] = d['signatures']
+        except KeyError:
+            logger.error("Child model without signatures")
+            raise
+        metadata['segment'] = d.get('segment', '0')
+
     return metadata
 
 
@@ -51,13 +65,20 @@ def get_mdx_model(d):
     return MDX_Net(dim_f=dim_f, ch=channels, num_stages=stages)
 
 
+def get_model_path(d):
+    parent = d.get("parent")
+    if parent is None:
+        return d.get('model_path')
+    return parent.get('model_path')
+
+
 def get_demucs_model(d):
     """ Create a Demucs, HDemucs (Hybrid) or HTDemucs (Hybrid Transformer) object.
         All metadata comes from the safetensors """
-    file_path = d['model_path']
+    file_path = get_model_path(d)
 
     # 1. First, open the file safely to read only the metadata header.
-    metadata = get_metadata(file_path)
+    metadata = get_metadata(file_path, d)
 
     is_bag = json.loads(metadata.get('is_bag_of_models', 'false'))
     signatures = json.loads(metadata['signatures'])

source/utils/load_safetensors.py

@@ -38,15 +38,17 @@ def load_safetensors(model_path, model_run, device):
         signatures = model_run.signatures
         if len(signatures) == 1:
             # Single model
-            sig = signatures[0]
-            logger.debug(f"Single model with filtered keys, prefix: {sig}")
+            prefix = f"{signatures[0]}."
+            logger.debug(f"Single model with filtered keys, prefix: {prefix}")
             # Just remove the prefix
-            state_dict = {k[len(sig)+1:]: v for k, v in state_dict.items()}
+            # Note: child models needs the if k.startswith(prefix) because they contain extra keys
+            state_dict = {k[len(prefix):]: v for k, v in state_dict.items() if k.startswith(prefix)}
             # 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):
+            # Load weights, but just once when the same model is used more than once
+            for sig, sub_model in {s: m for m, s in zip(model_run.models, signatures)}.items():
                 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)}