set-soft-AudioSeparation/source/db/models_db.py

# Copyright (c) 2025 Salvador E. Tropea
# Copyright (c) 2025 Instituto Nacional de Tecnología Industrial
# License: GPLv3
# Project: ComfyUI-AudioSeparation
#
# Handles a JSON file with information about the supported models.
# The JSON file is similar to what UVR uses, but with more information.
import json
import logging
import os
from pathlib import Path
from ..utils.misc import NODES_NAME
from ..utils.downloader import download_model as download_model_basic
from ..utils.comfy_notification import send_toast_notification
from .hash_dir import hash_dir
from .hash import is_hash, get_hash

logger = logging.getLogger(f"{NODES_NAME}.models_db")
known_models = None
known_models_mtime = None
# ICON_REMOTE = "\u2601"         # ☁️ Cloud
ICON_REMOTE = "⬇️  "         # "\u2B07"      # ⬇️
ICON_DOWNLOADED = "\U0001F4BE "  # 💾 Floppy Disk
KNOWN_SOURCES = {'Politrees/MDXNet': 'https://huggingface.co/Politrees/UVR_resources/resolve/main/models/MDXNet',
                 'Main/MDX': 'https://huggingface.co/set-soft/audio_separation/resolve/main/MDX'}


def get_db_filename(provided=None):
    if provided is not None:
        return provided
    # Get the directory where this script is located
    script_dir = Path(__file__).resolve().parent

    # Build the path to the JSON file: go up one level, then into models/
    json_path = script_dir / ".." / ".." / "models" / "uvr_model_data.json"
    try:
        return json_path.resolve().relative_to(Path.cwd())
    except ValueError:
        pass
    return json_path.resolve()


def load_known_models(json_path=None):
    """
    Loads the uvr_model_data.json configuration file using a path relative
    to this script's location. This is the most reliable method.
    """
    global known_models
    global known_models_mtime
    json_path = get_db_filename(json_path)

    # Check if we have a fresh db
    if known_models is not None and known_models_mtime == os.path.getmtime(json_path):
        return known_models

    try:
        logger.debug(f"Attempting to load JSON from: {json_path}")

        # Open and load the JSON file
        with open(json_path, 'r', encoding='utf-8') as f:
            data = json.load(f)

        known_models = data
        known_models_mtime = os.path.getmtime(json_path)
        return data

    except FileNotFoundError:
        logger.error("Error: The models database was not found at the expected location.")
        logger.error("Please check the directory structure.")
        return None
    except json.JSONDecodeError:
        logger.error(f"Error: The file at '{json_path}' is not a valid JSON file.")
        return None
    except Exception as e:
        logger.error(f"An unexpected error occurred: {e}")
        return None


def save_known_models(data, json_path=None):
    json_path = get_db_filename(json_path)
    backup_path = os.path.join(os.path.dirname(json_path), f".{os.path.basename(json_path)}~")

    try:
        logger.debug(f"Creating DB backup at '{backup_path}'...")
        if os.path.exists(json_path):
            os.rename(json_path, backup_path)

        logger.debug(f"Saving new database to '{json_path}'...")
        with open(json_path, 'w') as f:
            json.dump(data, f, indent=4, sort_keys=True)

        logger.debug("✅ Database updated and saved successfully.")

        global known_models
        global known_models_mtime
        known_models = data
        known_models_mtime = os.path.getmtime(json_path)

    except Exception as e:
        logger.error(f"during database save: {e}")
        logger.info("Attempting to restore from backup...")
        if os.path.exists(backup_path):
            os.rename(backup_path, json_path)
            logger.info("Backup restored.")
        raise


def create_display_name(d, hash, desc, primary_stem, model_t, file_t, downloaded):
    if primary_stem and len(primary_stem) == 1:
        desc = desc.replace(next(iter(primary_stem)), "")
    if model_t and len(model_t) == 1:
        desc = desc.replace(next(iter(model_t)), "")
    if file_t is None or len(file_t) > 1:
        # We can get a name collision i.e. ONNX + safetensors
        file_extra = f" [{d['file_t']}]"
    else:
        file_extra = ""
    return " ".join(desc.split()) + file_extra


def get_models_full(primary_stem=None, model_t=None, file_t=None, json_path=None, downloaded=None, default=None,
                    repeat_dl=False):
    """ Returns a dict with models that satisfy the provided criteria """

    # Allow for multiple values in the filters
    if isinstance(primary_stem, str):
        primary_stem = {primary_stem}
    if isinstance(model_t, str):
        model_t = {model_t}
    if isinstance(file_t, str):
        file_t = {file_t}

    # Filter the models db
    models = load_known_models(json_path)
    found = {}
    found_hashes = {}
    found_disk = {}
    def_sep = []
    on_disk = []
    to_down = []
    on_disk_as_down = []
    for hash, d in models.items():
        # Skip unnamed models, valid, these are models we don't have or support
        try:
            name = d['name']
        except KeyError:
            continue
        # Description is mandatory
        try:
            desc = d['desc']
        except KeyError:
            logger.error(f"Missing `desc` for {name}")
            continue
        # Check the stem
        try:
            if primary_stem is not None and d['primary_stem'] not in primary_stem:
                continue
        except KeyError:
            logger.error(f"Missing `primary_stem` for {name}")
            continue
        # Check the network type
        try:
            if model_t is not None and d['model_t'] not in model_t:
                continue
        except KeyError:
            logger.error(f"Missing `model_t` for {name}")
            continue
        # Check the container type
        try:
            if file_t is not None and d['file_t'] not in file_t:
                continue
        except KeyError:
            logger.error(f"Missing `file_t` for {name}")
            continue
        filtered_name = create_display_name(d, hash, desc, primary_stem, model_t, file_t, downloaded)
        d['hash'] = hash
        d['filtered_name'] = filtered_name
        if downloaded is not None:
            file_name = downloaded.get(hash)
            d['indicator'] = ICON_DOWNLOADED if file_name is not None else ICON_REMOTE
        else:
            file_name = None
            d['indicator'] = ""
        if default and default == name:
            if downloaded is None:
                def_sep.append(filtered_name)
            else:
                if file_name is not None:
                    def_sep.append(ICON_DOWNLOADED + filtered_name)
                    if repeat_dl:
                        # A copy for ComfyUI, so the user doesn't need to refresh the node and choose the one downloaded
                        # Also helps to allow forcing save a node with the file as "not downloaded"
                        # Is a hack, but is the best I came up
                        on_disk_as_down.append(ICON_REMOTE + filtered_name)
                else:
                    def_sep.append(ICON_REMOTE + filtered_name)
        else:
            if file_name is not None:
                on_disk.append(ICON_DOWNLOADED + filtered_name)
                if repeat_dl:
                    on_disk_as_down.append(ICON_REMOTE + filtered_name)
            else:
                to_down.append(ICON_REMOTE + filtered_name)
        found[filtered_name] = d
        found_hashes[hash] = d
        if file_name is not None:
            d['model_path'] = downloaded[hash]
            found_disk[os.path.realpath(file_name)] = d

    return found, found_hashes, found_disk, def_sep + sorted(on_disk) + sorted(to_down) + sorted(on_disk_as_down)


def get_models(primary_stem=None, model_t=None, file_t=None, json_path=None, downloaded=None, default=None, by_hash=False):
    dnames, hashes, fnames, ldnames = get_models_full(primary_stem, model_t, file_t, json_path, downloaded, default)
    return hashes if by_hash else dnames, ldnames


def get_download_url(data):
    try:
        name = data['name']
        dn_t = data['download']
    except KeyError:
        return None
    try:
        return os.path.join(KNOWN_SOURCES[dn_t], name)
    except KeyError:
        logger.error(f"Unknown download source `{dn_t}`")
        return None


def cli_add_models_and_db(parser):
    # Compute the models dir assuming the script is run from a clone of the repo
    default_json_file = get_db_filename()

    parser.add_argument('--models_dir', type=str, default=os.path.dirname(default_json_file),
                        help="Path to the directory containing model files.")
    parser.add_argument('--json_file', type=str, default=default_json_file,
                        help="Path to the models database JSON file.")


def download_model(data, models_dir):
    # Check we can download it
    url = get_download_url(data)
    if url is None:
        raise ValueError("Model is not downloadable")

    # Download the file
    name = data['name']
    send_toast_notification(f"Downloading `{name}`", "Download")
    try:
        fname = download_model_basic(url, models_dir, name)
        # Mark it as downloaded
        data['model_path'] = fname
        if data['indicator']:
            data['indicator'] = ICON_DOWNLOADED
        # Notify the user
        send_toast_notification("Finished downloading", "Download", 'success')
        return fname
    except Exception as e:
        raise ValueError(f"Failed to download {name} from {url}\n{e}")


class FilteredModels(object):
    def __init__(self, primary_stem=None, model_t=None, file_t=None, json_path=None, downloaded=None, default=None,
                 repeat_dl=False):
        self.primary_stem = primary_stem
        self.model_t = model_t
        self.file_t = file_t
        self.by_dname, self.by_hash, self.by_fname, self.dnames = get_models_full(primary_stem, model_t, file_t, json_path,
                                                                                  downloaded, default, repeat_dl)

    def get_by_display_name(self, name):
        if name.startswith(ICON_REMOTE) or name.startswith(ICON_DOWNLOADED):
            name = name[name.index(' ')+1:].strip()
        return self.by_dname.get(name)

    def get_by_hash(self, hash):
        return self.by_hash.get(hash)

    def get_by_file_name(self, file):
        return self.by_fname.get(file)

    def get(self, value):
        # If it looks like a hash try it first
        if is_hash(value):
            d = self.by_hash.get(value)
            if d:
                return d
        # Then try by a display name
        d = self.by_dname.get(value)
        if d:
            return d
        # Is this a file name?
        if os.path.isfile(value):
            # Try using its hash
            hash = get_hash(value)
            return self.by_hash.get(hash)
        return None

    def get_display_names(self, clean=False):
        return [m[m.index(' ')+1:].strip() for m in self.dnames] if clean else self.dnames


class ModelsDB(object):
    def __init__(self, models_dir: str, json_path: str = None):
        super().__init__()
        self.models_dir = models_dir
        self.json_path = json_path
        self.refresh()

    def refresh(self):
        self.downloaded = hash_dir(self.models_dir)
        self.models = load_known_models(self.json_path)

    def get_filtered(self, primary_stem=None, model_t=None, file_t=None, default=None, repeat_dl=False):
        return FilteredModels(primary_stem=primary_stem, model_t=model_t, file_t=file_t, json_path=self.json_path,
                              downloaded=self.downloaded, default=default, repeat_dl=repeat_dl)