[Added] Tool to convert Demucs into safetensors and see its metadata
This commit is contained in:
Salvador E. Tropea
@@ -13,6 +13,13 @@ Might work for other ONNX files using the same operands.
|
||||
But you'll need a PyTorch class for its architecture that creates the layers
|
||||
in the same order as the ONNX file.
|
||||
|
||||
# Demucs to safetensors
|
||||
|
||||
File: demucs2safetensors.py
|
||||
|
||||
From a YAML file and the PyTorch components it creates safetensors version of the Demucs file.
|
||||
All metadata and models are contained in the new file.
|
||||
|
||||
# Batch converter
|
||||
|
||||
File: batch_convert.py
|
||||
@@ -59,6 +66,13 @@ Used to display our PyTorch class, also the state_dict keys.
|
||||
It can optionally export the class structure as an ONNX file that can be loaded by
|
||||
Netron, but contains too much extra names.
|
||||
|
||||
# Show Metadata
|
||||
|
||||
File: show_metadata.py
|
||||
|
||||
Used to display our safetensors metadata.
|
||||
It can decode structures in JSON format we use for Demucs.
|
||||
|
||||
# Style Fixer
|
||||
|
||||
File: style_fixer.py
|
||||
|
||||
@@ -0,0 +1,216 @@
|
||||
# Copyright (c) 2025 Salvador E. Tropea
|
||||
# Copyright (c) 2025 Instituto Nacional de Tecnología Industrial
|
||||
# License: GPLv3
|
||||
# Project: ComfyUI-AudioSeparation
|
||||
#
|
||||
# Tool to convert a Demucs v3/4 model into safetensors
|
||||
# python tool/tool/demucs2safetensors.py --yaml DEMUCS.YAML
|
||||
# You must manually download the .th files and copy them to the YAML dir
|
||||
# First version by Gemini 2.5 Pro
|
||||
from copy import deepcopy
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
from safetensors.torch import save_file
|
||||
import sys
|
||||
import torch
|
||||
from typing import Dict
|
||||
import yaml
|
||||
|
||||
try:
|
||||
# We need the original Demucs library to dequantize
|
||||
from demucs.states import set_state # noqa: F401
|
||||
with_demuc_lib = True
|
||||
except Exception:
|
||||
with_demuc_lib = False
|
||||
import bootstrap # noqa: F401
|
||||
from source.utils.misc import cli_add_verbose, FractionEncoder
|
||||
from source.utils.logger import main_logger, logger_set_standalone
|
||||
from source.db.models_db import cli_add_db, get_download_url, ModelsDB
|
||||
from source.db.hash import get_hash
|
||||
logger = main_logger
|
||||
|
||||
|
||||
def convert_demucs_model(yaml_path_str: str, model_paths: list[str], output_path_str: str, all_metadata, data: Dict):
|
||||
"""
|
||||
Loads an original Demucs model bag, extracts weights and all necessary metadata
|
||||
from the YAML and .th files, and saves it to a single, secure .safetensors file.
|
||||
"""
|
||||
yaml_path = Path(yaml_path_str)
|
||||
output_path = Path(output_path_str)
|
||||
|
||||
# 1. Load and parse the YAML file
|
||||
main_logger.info(f"\n--- 🚀 Starting conversion for {yaml_path.name} ---\n")
|
||||
main_logger.info(f"- Loading YAML definition from: {yaml_path}")
|
||||
with open(yaml_path, 'r') as f:
|
||||
yaml_data = yaml.safe_load(f)
|
||||
|
||||
signatures = yaml_data['models']
|
||||
main_logger.info(f"- Found {len(signatures)} model signatures in YAML: {signatures}")
|
||||
|
||||
# 2. Match the provided .th files to their signatures
|
||||
model_file_map = {}
|
||||
for path_str in model_paths:
|
||||
path = Path(path_str)
|
||||
# The signature is the part of the filename before the first '-' or '.'
|
||||
sig = path.stem.split('-')[0]
|
||||
if sig not in signatures:
|
||||
main_logger.warning(f"File {path.name} with signature {sig} is not listed in the YAML file. Skipping.")
|
||||
continue
|
||||
if sig in model_file_map:
|
||||
raise ValueError(f"Duplicate files found for signature {sig}.")
|
||||
model_file_map[sig] = path
|
||||
|
||||
# Verify that all signatures from the YAML have a corresponding file
|
||||
if len(model_file_map) != len(set(signatures)):
|
||||
missing_sigs = set(signatures) - set(model_file_map.keys())
|
||||
raise FileNotFoundError(f"Missing model files for signatures: {missing_sigs}")
|
||||
|
||||
main_logger.info("- Successfully mapped all signatures to model files.")
|
||||
|
||||
# 3. Load model packages and extract metadata and state dicts
|
||||
# all_metadata = {}
|
||||
consolidated_state_dict = {}
|
||||
|
||||
for sig in signatures:
|
||||
path = model_file_map[sig]
|
||||
main_logger.info(f"\n- Processing model '{sig}' from '{path.name}'...")
|
||||
|
||||
# This is the only "unsafe" part, loading the original pickle file
|
||||
pkg = torch.load(path, map_location='cpu', weights_only=False)
|
||||
|
||||
klass, args, kwargs, state = pkg["klass"], pkg["args"], pkg["kwargs"], pkg["state"]
|
||||
main_logger.info(f" - Model class: {klass.__module__}.{klass.__name__}")
|
||||
|
||||
# Dequantize if necessary by letting the original code handle it
|
||||
if state.get('__quantized'):
|
||||
if not with_demuc_lib:
|
||||
main_logger.error("Don't use quantized models. Look for the same model without `_q`")
|
||||
main_logger.error("Alternatively install the demucs Python module")
|
||||
sys.exit(3)
|
||||
else:
|
||||
main_logger.info(" - Model is quantized. Dequantizing weights...")
|
||||
model_instance = klass(*args, **kwargs)
|
||||
set_state(model_instance, state)
|
||||
clean_state_dict = model_instance.state_dict()
|
||||
else:
|
||||
clean_state_dict = state
|
||||
|
||||
# Store this model's metadata, keyed by its signature
|
||||
all_metadata[sig] = json.dumps({
|
||||
'class_module': klass.__module__,
|
||||
'class_name': klass.__name__,
|
||||
'args': args,
|
||||
'kwargs': kwargs,
|
||||
}, cls=FractionEncoder) # Use the custom encoder here
|
||||
|
||||
# Add the weights to the consolidated dict, prefixed by signature
|
||||
for key, value in clean_state_dict.items():
|
||||
consolidated_state_dict[f"{sig}.{key}"] = value
|
||||
|
||||
# 4. Add the top-level YAML data to the metadata
|
||||
all_metadata['is_bag_of_models'] = json.dumps(len(set(signatures)) > 1)
|
||||
all_metadata['signatures'] = json.dumps(signatures)
|
||||
if 'weights' in yaml_data:
|
||||
all_metadata['weights'] = json.dumps(yaml_data['weights'])
|
||||
if 'segment' in yaml_data:
|
||||
all_metadata['segment'] = str(yaml_data['segment'])
|
||||
|
||||
# 5. Calculate the total number of parameters from the final state dict
|
||||
total_params = sum(p.numel() for p in consolidated_state_dict.values())
|
||||
main_logger.info(f"- Total number of parameters in the model: {total_params:,}")
|
||||
# Add the count as a string to the metadata dictionary
|
||||
data['params'] = all_metadata['params'] = str(total_params)
|
||||
|
||||
# 6. Save the final .safetensors file
|
||||
main_logger.info(f"- \U0001F4BE Saving consolidated model and metadata to: {output_path}")
|
||||
save_file(consolidated_state_dict, output_path, metadata=all_metadata)
|
||||
main_logger.info("\n--- 🎉 Conversion complete! ---")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser(description="Convert Demucs .th models to a single .safetensors file.")
|
||||
parser.add_argument('--yaml', required=True, type=str, help="Path to the Demucs .yaml file.")
|
||||
parser.add_argument('--models', nargs='*', default=None, type=str,
|
||||
help="Optional. Paths to .th files. If not provided, assumes they are "
|
||||
"in the same directory as the YAML.")
|
||||
parser.add_argument('--output', default=None, type=str,
|
||||
help="Optional. Path for the output .safetensors file. If not provided, "
|
||||
"it's saved next to the YAML with the same name.")
|
||||
cli_add_db(parser)
|
||||
cli_add_verbose(parser)
|
||||
|
||||
args = parser.parse_args()
|
||||
logger_set_standalone(args)
|
||||
main_logger.info("⚙️ PyTorch Demucs to Safetensors converter\n")
|
||||
|
||||
# Get information about the YAML file in our database
|
||||
yaml_path = Path(args.yaml)
|
||||
db = ModelsDB(yaml_path.parent)
|
||||
known = db.get_filtered(model_t="Demucs")
|
||||
hash = get_hash(yaml_path)
|
||||
d = known.get_by_hash(hash)
|
||||
if d is None:
|
||||
# We don't have the hash for it
|
||||
d = known.get_by_file_name(yaml_path.name)
|
||||
if d is None:
|
||||
# We use some information from the database to populate the metadata so is better if we have the
|
||||
# information in the database
|
||||
main_logger.error(f"{yaml_path} not in data base, please add it first")
|
||||
sys.exit(3)
|
||||
else:
|
||||
main_logger.error(f"{yaml_path} in database, but with different hash")
|
||||
|
||||
logger.info(f"Model description: {d['desc']}")
|
||||
|
||||
# Logic for optional --models
|
||||
model_files = args.models
|
||||
if not model_files:
|
||||
main_logger.debug("No --models provided. Searching for .th files alongside the YAML...")
|
||||
yaml_dir = yaml_path.parent
|
||||
with open(yaml_path, 'r') as f:
|
||||
signatures = yaml.safe_load(f)['models']
|
||||
|
||||
model_files = []
|
||||
for sig in set(signatures): # Use set to avoid redundant searches
|
||||
found = list(yaml_dir.glob(f'{sig}*.th'))
|
||||
if not found:
|
||||
raise FileNotFoundError(f"Could not automatically find a model file for signature '{sig}' in {yaml_dir}")
|
||||
if len(found) > 1:
|
||||
full_match = yaml_dir / (sig + '.th') # UVR Demucs uses it
|
||||
if full_match in found:
|
||||
found = [full_match]
|
||||
else:
|
||||
main_logger.warning(f"Found multiple files for signature '{sig}', using the first one: {found[0]}")
|
||||
model_files.append(str(found[0]))
|
||||
main_logger.info(f"Automatically found model files: {model_files}")
|
||||
|
||||
# Logic for optional --output
|
||||
output_file = args.output
|
||||
if not output_file:
|
||||
output_file = yaml_path.with_suffix('.safetensors')
|
||||
main_logger.info(f"Defaulting to: `{output_file}` (No --output provided)")
|
||||
else:
|
||||
output_file = Path(output_file)
|
||||
|
||||
# Adjust the data to the converted version
|
||||
metadata = {}
|
||||
d = deepcopy(d)
|
||||
d["download"] = "Main/Demucs"
|
||||
d["name"] = Path(d["name"]).stem + ".safetensors"
|
||||
d["file_t"] = "safetensors"
|
||||
# Add it to the safetensors
|
||||
metadata["desc"] = d["desc"]
|
||||
metadata["download"] = get_download_url(d)
|
||||
metadata["file_t"] = "safetensors"
|
||||
metadata["model_t"] = "Demucs"
|
||||
metadata["name"] = d["name"]
|
||||
metadata["primary_stem"] = json.dumps(d["primary_stem"])
|
||||
metadata["project"] = "https://github.com/set-soft/AudioSeparation"
|
||||
|
||||
convert_demucs_model(args.yaml, model_files, output_file, metadata, d)
|
||||
|
||||
# Now update the DB
|
||||
db.remove(known.get_by_file_name(d["name"]))
|
||||
db.add(get_hash(output_file), d)
|
||||
db.save()
|
||||
@@ -0,0 +1,35 @@
|
||||
# Copyright (c) 2025 Salvador E. Tropea
|
||||
# Copyright (c) 2025 Instituto Nacional de Tecnología Industrial
|
||||
# License: GPLv3
|
||||
# Project: ComfyUI-AudioSeparation
|
||||
#
|
||||
# Tool to show the safetensors metadata
|
||||
# python tool/show_metadata.py model.safetensors
|
||||
import argparse
|
||||
import json
|
||||
from pprint import pprint
|
||||
# Local imports
|
||||
import bootstrap # noqa: F401
|
||||
from source.utils.logger import main_logger, logger_set_standalone
|
||||
from source.utils.misc import cli_add_verbose, json_object_hook
|
||||
from source.inference.get_model import get_metadata
|
||||
|
||||
|
||||
def show_metadata(args):
|
||||
d = get_metadata(args.input_file)
|
||||
main_logger.info(f"Metadata information for `{args.input_file}`")
|
||||
if 'model_t' not in d:
|
||||
main_logger.warning("Missing model_t key, this isn't an AudioSeparation file")
|
||||
expanded = {k: json.loads(v, object_hook=json_object_hook) if v[0] in '{[' else v for k, v in d.items()}
|
||||
pprint(expanded)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Shows the safetensors metadata",
|
||||
formatter_class=argparse.RawTextHelpFormatter)
|
||||
parser.add_argument('input_file', type=str, help="Path to the input ONNX model file.")
|
||||
cli_add_verbose(parser)
|
||||
|
||||
args = parser.parse_args()
|
||||
logger_set_standalone(args)
|
||||
show_metadata(args)
|
||||
Reference in New Issue
Block a user