| | from pathlib import Path
|
| |
|
| | import argbind
|
| | from audiotools import ml
|
| |
|
| | from ..model import 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
|
| |
|
