utils.py

import os
import torch
import requests
from huggingface_hub import snapshot_download

MODEL_DIR = snapshot_download("ccmusic-database/music_genre", cache_dir="./__pycache__")


def toCUDA(x):
    if hasattr(x, "cuda"):
        if torch.cuda.is_available():
            return x.cuda()

    return x


def find_mp3_files(folder_path=f"{MODEL_DIR}/examples"):
    wav_files = []
    for root, _, files in os.walk(folder_path):
        for file in files:
            if file.endswith(".mp3"):
                file_path = os.path.join(root, file)
                wav_files.append(file_path)

    return wav_files


def get_modelist(model_dir=MODEL_DIR, assign_model=""):
    try:
        entries = os.listdir(model_dir)

    except OSError as e:
        print(f"Cannot access {model_dir}: {e}")
        return

    output = []
    for entry in entries:
        full_path = os.path.join(model_dir, entry)
        if entry == ".git" or entry == "examples":
            print(f"Skip .git / examples dir: {full_path}")
            continue

        if os.path.isdir(full_path):
            model = os.path.basename(full_path)
            if assign_model and assign_model.lower() in model:
                output.insert(0, model)
            else:
                output.append(model)

    return output


def download(url: str):
    filename = url.split("/")[-1]
    response = requests.get(url, stream=True)
    if response.status_code == 200:
        with open(filename, "wb") as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)

        print(f"The file has been downloaded to {os.getcwd()}/{filename}")

    else:
        print(f"Failed to download, status code: {response.status_code}")