| from pathlib import Path
|
|
|
| import argbind
|
| from audiotools import ml
|
|
|
| import dac
|
|
|
| DAC = dac.model.DAC
|
| Accelerator = ml.Accelerator
|
|
|
| __MODEL_LATEST_TAGS__ = {
|
| ("44khz", "8kbps"): "0.0.1",
|
| ("24khz", "8kbps"): "0.0.4",
|
| ("16khz", "8kbps"): "0.0.5",
|
| ("44khz", "16kbps"): "1.0.0",
|
| }
|
|
|
| __MODEL_URLS__ = {
|
| (
|
| "44khz",
|
| "0.0.1",
|
| "8kbps",
|
| ): "https://github.com/descriptinc/descript-audio-codec/releases/download/0.0.1/weights.pth",
|
| (
|
| "24khz",
|
| "0.0.4",
|
| "8kbps",
|
| ): "https://github.com/descriptinc/descript-audio-codec/releases/download/0.0.4/weights_24khz.pth",
|
| (
|
| "16khz",
|
| "0.0.5",
|
| "8kbps",
|
| ): "https://github.com/descriptinc/descript-audio-codec/releases/download/0.0.5/weights_16khz.pth",
|
| (
|
| "44khz",
|
| "1.0.0",
|
| "16kbps",
|
| ): "https://github.com/descriptinc/descript-audio-codec/releases/download/1.0.0/weights_44khz_16kbps.pth",
|
| }
|
|
|
|
|
| @argbind.bind(group="download", positional=True, without_prefix=True)
|
| def download(
|
| model_type: str = "44khz", model_bitrate: str = "8kbps", tag: str = "latest"
|
| ):
|
| """
|
| Function that downloads the weights file from URL if a local cache is not found.
|
|
|
| Parameters
|
| ----------
|
| model_type : str
|
| The type of model to download. Must be one of "44khz", "24khz", or "16khz". Defaults to "44khz".
|
| model_bitrate: str
|
| Bitrate of the model. Must be one of "8kbps", or "16kbps". Defaults to "8kbps".
|
| Only 44khz model supports 16kbps.
|
| tag : str
|
| The tag of the model to download. Defaults to "latest".
|
|
|
| Returns
|
| -------
|
| Path
|
| Directory path required to load model via audiotools.
|
| """
|
| model_type = model_type.lower()
|
| tag = tag.lower()
|
|
|
| assert model_type in [
|
| "44khz",
|
| "24khz",
|
| "16khz",
|
| ], "model_type must be one of '44khz', '24khz', or '16khz'"
|
|
|
| assert model_bitrate in [
|
| "8kbps",
|
| "16kbps",
|
| ], "model_bitrate must be one of '8kbps', or '16kbps'"
|
|
|
| if tag == "latest":
|
| tag = __MODEL_LATEST_TAGS__[(model_type, model_bitrate)]
|
|
|
| download_link = __MODEL_URLS__.get((model_type, tag, model_bitrate), None)
|
|
|
| if download_link is None:
|
| raise ValueError(
|
| f"Could not find model with tag {tag} and model type {model_type}"
|
| )
|
|
|
| local_path = (
|
| Path.home()
|
| / ".cache"
|
| / "descript"
|
| / "dac"
|
| / f"weights_{model_type}_{model_bitrate}_{tag}.pth"
|
| )
|
| if not local_path.exists():
|
| local_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
| import requests
|
|
|
| response = requests.get(download_link)
|
|
|
| if response.status_code != 200:
|
| raise ValueError(
|
| f"Could not download model. Received response code {response.status_code}"
|
| )
|
| local_path.write_bytes(response.content)
|
|
|
| return local_path
|
|
|
|
|
| def load_model(
|
| model_type: str = "44khz",
|
| model_bitrate: str = "8kbps",
|
| tag: str = "latest",
|
| load_path: str = None,
|
| ):
|
| if not load_path:
|
| load_path = download(
|
| model_type=model_type, model_bitrate=model_bitrate, tag=tag
|
| )
|
| generator = DAC.load(load_path)
|
| return generator
|
|
|
