Spaces:
Running on Zero
Running on Zero
| from . import modules | |
| from pathlib import Path | |
| from . import scheduler | |
| from .interface import Interface | |
| from .modules.transformer import VampNet | |
| __version__ = "0.0.1" | |
| ROOT = Path(__file__).parent.parent | |
| MODELS_DIR = ROOT / "models" / "vampnet" | |
| from huggingface_hub import hf_hub_download, HfFileSystem | |
| DEFAULT_HF_MODEL_REPO_DIR = ROOT / "DEFAULT_HF_MODEL_REPO" | |
| DEFAULT_HF_MODEL_REPO = DEFAULT_HF_MODEL_REPO_DIR.read_text().strip() | |
| # DEFAULT_HF_MODEL_REPO = "hugggof/vampnet" | |
| # The weights' revision, which every download below names, so that the models cannot change | |
| # under the Space | |
| REVISION = "69e9c7dba72b54a45f93122ef91aa738d2e803ba" | |
| FS = HfFileSystem() | |
| def download_codec(): | |
| # from dac.model.dac import DAC | |
| from lac.model.lac import LAC as DAC | |
| repo_id = DEFAULT_HF_MODEL_REPO | |
| filename = "codec.pth" | |
| codec_path = hf_hub_download( | |
| repo_id=repo_id, | |
| filename=filename, | |
| subfolder=None, | |
| local_dir=MODELS_DIR, | |
| revision=REVISION, | |
| ) | |
| return codec_path | |
| # hf_hub_download writes each checkpoint beside its final path and moves it there only once its | |
| # size is confirmed, so an interrupted download cannot leave a truncated file that a later start | |
| # would take for a complete one | |
| def download_default(): | |
| filenames = ["coarse.pth", "c2f.pth", "wavebeat.pth"] | |
| paths = [ | |
| hf_hub_download(repo_id=DEFAULT_HF_MODEL_REPO, filename=filename, local_dir=MODELS_DIR, | |
| revision=REVISION) | |
| for filename in filenames | |
| ] | |
| # load the models | |
| return paths[0], paths[1] | |
| def download_finetuned(name, repo_id=DEFAULT_HF_MODEL_REPO): | |
| filenames = ["coarse.pth", "c2f.pth"] | |
| paths = [ | |
| hf_hub_download(repo_id=repo_id, filename=f"loras/{name}/{filename}", local_dir=MODELS_DIR, | |
| revision=REVISION) | |
| for filename in filenames | |
| ] | |
| # load the models | |
| return paths[0], paths[1] | |
| def list_finetuned(repo_id=DEFAULT_HF_MODEL_REPO): | |
| diritems = FS.listdir(f"{repo_id}@{REVISION}/loras") | |
| # iterate through all the names | |
| valid_diritems = [] | |
| for item in diritems: | |
| model_file_items = FS.listdir(item["name"]) | |
| item_names = [item["name"].split("/")[-1] for item in model_file_items] | |
| # check that theres a "c2f.pth" and "coarse.pth" in the items | |
| c2f_exists = "c2f.pth" in item_names | |
| coarse_exists = "coarse.pth" in item_names | |
| if c2f_exists and coarse_exists: | |
| valid_diritems.append(item) | |
| # get the names of the valid items | |
| names = [item["name"].split("/")[-1] for item in valid_diritems] | |
| return names | |