{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"id": "186304df",
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\""
]
},
{
"cell_type": "code",
"execution_count": 47,
"id": "1c5e1be8",
"metadata": {},
"outputs": [],
"source": [
"from encodec import EncodecModel\n",
"from encodec.utils import convert_audio\n",
"\n",
"import torchaudio\n",
"import torch\n",
"import shutil\n",
"import numpy as np\n",
"\n",
"from suno_utils.audio import Audio\n",
"\n",
"# Causal model: \n",
"# monophonic, variety of audio\n",
"# supports 1.5, 3, 6, 12 or 24 kbps\n",
"model = EncodecModel.encodec_model_24khz() \n",
"\n",
"# Non-causal model:\n",
"# stereophonic audio, music-only audio\n",
"# model = EncodecModel.encodec_model_48khz()\n",
"# supports 3, 6, 12 and 24 kbps\n",
"model.set_target_bandwidth(6)"
]
},
{
"cell_type": "code",
"execution_count": 48,
"id": "97b83f7f",
"metadata": {},
"outputs": [],
"source": [
"# Load and pre-process the audio waveform\n",
"wav, sr = torchaudio.load(\"../s2t/sample_data/audio_16k/russia.wav\")\n",
"wav = wav.unsqueeze(0)\n",
"wav = convert_audio(wav, sr, model.sample_rate, model.channels)\n",
"\n",
"with torch.no_grad():\n",
" encoded_frames = model.encode(wav)\n",
"\n",
"# codes = torch.cat([encoded[0] for encoded in encoded_frames], dim=-1) # [B, n_q, T]\n",
"\n",
"decoded_frames = model.decode(encoded_frames)"
]
},
{
"cell_type": "code",
"execution_count": 50,
"id": "8803c597",
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n"
],
"text/plain": [
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"Audio.from_array_float(decoded_frames.detach().cpu().numpy().squeeze(), sample_rate=24_000).play()"
]
},
{
"cell_type": "code",
"execution_count": 52,
"id": "b3fc6db9",
"metadata": {},
"outputs": [],
"source": [
"## measure size (this is probably not the best way to save, also can still lossless compress with LM)\n",
"# shutil.rmtree(\"tmp_conv\", ignore_errors=True)\n",
"# os.makedirs(\"tmp_conv\")\n",
"# for n, (e, _) in enumerate(encoded_frames):\n",
"# en = e.detach().cpu().numpy().squeeze()\n",
"# np.save(f\"tmp_conv/{n}.npz\", en)\n",
"# print(round(np.sum([\n",
"# os.path.getsize(os.path.join(\"tmp_conv\", f)) \n",
"# for f in os.listdir(\"tmp_conv\") if os.path.isfile(os.path.join(\"tmp_conv\", f))\n",
"# ]) / 1_000, 1), \"kB\")\n",
"# shutil.rmtree(\"tmp_conv\", ignore_errors=True)"
]
},
{
"cell_type": "code",
"execution_count": 51,
"id": "a43b5b87",
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n"
],
"text/plain": [
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"Audio.play_audio(\"../s2t/sample_data/audio_16k/russia.wav\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4ea9a1df",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "c3b4530c",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "8b109c8b",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.8.10"
}
},
"nbformat": 4,
"nbformat_minor": 5
}