From b44e22911613823d524d6201af2104de8f573ddb Mon Sep 17 00:00:00 2001 From: "Salvador E. Tropea" Date: Tue, 8 Jul 2025 10:51:16 -0300 Subject: [PATCH] [Added] Support for partial models created from bigger ones I guess they are to test individual portions --- .gitignore | 1 + models/uvr_model_data.json | 66 ++++++++++++++++++++++++++++++++ source/db/hash_dir.py | 2 +- source/db/load_model.py | 10 +++++ source/db/models_db.py | 18 +++++++++ source/inference/get_model.py | 27 +++++++++++-- source/utils/load_safetensors.py | 10 +++-- 7 files changed, 126 insertions(+), 8 deletions(-) diff --git a/.gitignore b/.gitignore index 6cb7a3a..dfec806 100644 --- a/.gitignore +++ b/.gitignore @@ -12,4 +12,5 @@ models/all_public_uvr_models models/demucs __no__ models/.catalog.csv +models/*.yaml 0LEEME diff --git a/models/uvr_model_data.json b/models/uvr_model_data.json index c98d6cc..6b2fd8f 100644 --- a/models/uvr_model_data.json +++ b/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, diff --git a/source/db/hash_dir.py b/source/db/hash_dir.py index 8100fcb..0e142ba 100644 --- a/source/db/hash_dir.py +++ b/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: diff --git a/source/db/load_model.py b/source/db/load_model.py index 273b57f..db4582b 100644 --- a/source/db/load_model.py +++ b/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": diff --git a/source/db/models_db.py b/source/db/models_db.py index 7290d49..4cacb3b 100644 --- a/source/db/models_db.py +++ b/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: diff --git a/source/inference/get_model.py b/source/inference/get_model.py index a2afc71..302e41b 100644 --- a/source/inference/get_model.py +++ b/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']) diff --git a/source/utils/load_safetensors.py b/source/utils/load_safetensors.py index e6e19ad..99e86e9 100644 --- a/source/utils/load_safetensors.py +++ b/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)}