{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from suno_utils.utils.text import read_jsonl, write_jsonl\n",
    "from suno_utils.utils.s3 import read_from_s3\n",
    "import numpy as np\n",
    "import json\n",
    "import polars as pl\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# here we need to iterate over the pass tracks and read the audio from s3\n",
    "# once we get the audio we need to apply lowpass filter, and then ideally extract semantic codes and save them out \n",
    "import polars as pl\n",
    "import os\n",
    "import json\n",
    "base_dir = \"/app/suno/data/diffusion_mix/dac_vae_tuned_25hz/\"\n",
    "\n",
    "metas_tr_filepath = os.path.join(base_dir, \"metas_tr.jsonl\")\n",
    "metas_tr = pl.read_ndjson(metas_tr_filepath)\n",
    "\n",
    "print(metas_tr[0][\"id\"])\n",
    "\n",
    "info_tr_filepath = os.path.join(base_dir, \"info_tr_t4.json\")\n",
    "with open(info_tr_filepath, \"r\") as f:\n",
    "    info_tr = json.load(f)\n",
    "\n",
    "print(len(info_tr[\"diffusion_mix_fix\"]))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# filter metas to only include the passed tracks\n",
    "passed_ids = set(info_tr[\"diffusion_mix_fix\"])\n",
    "# Use polars filtering instead of converting to dicts\n",
    "passed_metas = metas_tr.filter(pl.col(\"id\").is_in(passed_ids)).to_dicts()\n",
    "\n",
    "print(len(passed_metas))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# now we take a sample of 50_000 from the passed metas and download to the local dir\n",
    "import random\n",
    "from joblib import Parallel, delayed\n",
    "from tqdm import tqdm\n",
    "\n",
    "sample_metas = random.sample(passed_metas, 1_000)\n",
    "\n",
    "output_dir = \"/home/christian/audio/genius_ear_1k\"\n",
    "os.makedirs(output_dir, exist_ok=True)\n",
    "\n",
    "\n",
    "def download_file(w):\n",
    "    filepath = w[\"s3_filepath\"]\n",
    "    ext = filepath.split(\".\")[-1]\n",
    "    meta_id = w[\"id\"]\n",
    "    out_filepath = os.path.join(output_dir, f\"{meta_id}.{ext}\")\n",
    "    if os.path.exists(out_filepath):\n",
    "        return\n",
    "    os.system(f\"aws s3 cp {filepath} {out_filepath} > /dev/null 2>&1\")\n",
    "\n",
    "# Use joblib for parallel downloads\n",
    "Parallel(n_jobs=100, backend=\"threading\")(\n",
    "    delayed(download_file)(w) for w in tqdm(sample_metas)\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the sample metas to a jsonl file\n",
    "write_jsonl(sample_metas, os.path.join(output_dir, \"source_metas.jsonl\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import random\n",
    "import torchaudio\n",
    "\n",
    "def apply_tanh_distortion(\n",
    "    audio: torch.Tensor, sample_rate: float, gain_db: float = 0.0\n",
    "):\n",
    "    gain_lin = 10 ** (gain_db / 20.0)\n",
    "    return torch.tanh(audio * gain_lin)\n",
    "\n",
    "\n",
    "def apply_clipping_distortion(\n",
    "    audio: torch.Tensor, sample_rate: float, gain_db: float = 0.0\n",
    "):\n",
    "    gain_lin = 10 ** (gain_db / 20.0)\n",
    "    return (audio * gain_lin).clamp(-1, 1)\n",
    "\n",
    "\n",
    "def apply_audio_codec(\n",
    "    audio: torch.Tensor,\n",
    "    sample_rate: float,\n",
    "    bit_rate: int = 16000,\n",
    "    n_passes: int = 1,\n",
    "    format_str: str = \"mp3\",\n",
    "):\n",
    "    for _ in range(n_passes):\n",
    "        effector = torchaudio.io.AudioEffector(\n",
    "            format=format_str,\n",
    "            codec_config=torchaudio.io.CodecConfig(bit_rate=bit_rate),\n",
    "        )\n",
    "        audio = effector.apply(audio.T, sample_rate).T\n",
    "    return audio\n",
    "\n",
    "def apply_noise(\n",
    "    audio: torch.Tensor,\n",
    "    sample_rate: float,\n",
    "    gain_db: float = 0.0,\n",
    "    noise_type: str = \"white\",\n",
    "):\n",
    "    gain_lin = 10 ** (gain_db / 20.0)\n",
    "    noise = torch.randn_like(audio)\n",
    "\n",
    "    if noise_type == \"white\":\n",
    "        return audio + gain_lin * noise\n",
    "    elif noise_type == \"pink\":\n",
    "        b = torch.tensor([0.049922035, -0.095993537, 0.050612699, -0.004408786])\n",
    "        a = torch.tensor([1, -2.494956002, 2.017265875, -0.522189400])\n",
    "        noise = torchaudio.functional.filtfilt(noise, a, b)\n",
    "        noise /= noise.abs().max()\n",
    "        return audio + gain_lin * noise\n",
    "    else:\n",
    "        raise ValueError(f\"Invalid noise type: {noise_type}\")\n",
    "\n",
    "def apply_highpass(audio: torch.Tensor, sample_rate: float, cutoff_hz: float):\n",
    "    return torchaudio.functional.highpass_biquad(audio, sample_rate, cutoff_hz)\n",
    "\n",
    "def apply_lowpass(audio: torch.Tensor, sample_rate: float, cutoff_hz: float):\n",
    "    return torchaudio.functional.lowpass_biquad(audio, sample_rate, cutoff_hz)\n",
    "\n",
    "def apply_random_tanh_distortion(audio: torch.Tensor, sample_rate: float):\n",
    "    gain_db = random.uniform(0, 12)\n",
    "    return apply_tanh_distortion(audio, sample_rate, gain_db)\n",
    "\n",
    "\n",
    "def apply_random_clipping_distortion(audio: torch.Tensor, sample_rate: float):\n",
    "    gain_db = random.uniform(0, 12)\n",
    "    return apply_clipping_distortion(audio, sample_rate, gain_db)\n",
    "\n",
    "\n",
    "def apply_random_noise(audio: torch.Tensor, sample_rate: float):\n",
    "    noise_type = random.choice([\"white\", \"pink\"])\n",
    "    noise_gain = random.uniform(-48, -6)\n",
    "    return apply_noise(audio, sample_rate, noise_gain, noise_type)\n",
    "\n",
    "\n",
    "def apply_random_audio_codec(audio: torch.Tensor, sample_rate: float):  \n",
    "    format_str = \"mp3\"\n",
    "    bit_rate = random.choice([8000, 16000, 32000, 64000, 128000])\n",
    "    n_passes = random.choice([1, 2])\n",
    "    return apply_audio_codec(audio, sample_rate, bit_rate, n_passes, format_str)\n",
    "\n",
    "def apply_random_highpass(audio: torch.Tensor, sample_rate: float):\n",
    "    cutoff_hz = random.uniform(20, 6000)\n",
    "    return apply_highpass(audio, sample_rate, cutoff_hz)\n",
    "\n",
    "def apply_random_lowpass(audio: torch.Tensor, sample_rate: float):\n",
    "    cutoff_hz = random.uniform(100, 12000)\n",
    "    return apply_lowpass(audio, sample_rate, cutoff_hz)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# next step is to corrupt the audio \n",
    "import glob\n",
    "from joblib import Parallel, delayed\n",
    "from tqdm import tqdm\n",
    "\n",
    "output_dir = \"/home/christian/audio/genius_ear_1k\"\n",
    "source_filepaths = glob.glob(os.path.join(output_dir, \"*.webm\"))\n",
    "print(len(source_filepaths))\n",
    "\n",
    "def corrupt_audio_and_save(filepath: str, p=0.33):\n",
    "    meta_id = os.path.basename(filepath).split(\".\")[0]\n",
    "    out_filepath = os.path.join(output_dir, f\"{meta_id}__corrupted.mp3\")\n",
    "\n",
    "    #if os.path.exists(out_filepath):\n",
    "    #    return\n",
    "\n",
    "    audio, sample_rate = torchaudio.load(filepath)\n",
    "\n",
    "    if random.random() < p:\n",
    "        audio = apply_random_highpass(audio, sample_rate)\n",
    "    if random.random() < p:\n",
    "        audio = apply_random_lowpass(audio, sample_rate)\n",
    "    if random.random() < p:\n",
    "        audio = apply_random_noise(audio, sample_rate)\n",
    "    if random.random() < p:\n",
    "        audio = apply_random_tanh_distortion(audio, sample_rate)\n",
    "    if random.random() < p:\n",
    "        audio = apply_random_clipping_distortion(audio, sample_rate)\n",
    "    if random.random() < p:\n",
    "        audio = apply_random_audio_codec(audio, sample_rate)\n",
    "    # save the audio\n",
    "    torchaudio.save(out_filepath, audio, sample_rate)\n",
    "\n",
    "# normal for loop\n",
    "#for filepath in tqdm(source_filepaths):\n",
    "#    corrupt_audio_and_save(filepath)\n",
    "    \n",
    "\n",
    "# process the audio in parallel\n",
    "_ = Parallel(n_jobs=-1)(\n",
    "    delayed(corrupt_audio_and_save)(filepath) for filepath in tqdm(source_filepaths)\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# now we need to semantic encode the corrupted audio\n",
    "from suno_utils.tasks.mert_25 import (\n",
    "    preload_models as preload_semantic_models,\n",
    "    load_model as load_semantic_model,\n",
    "    encode as encode_semantic,\n",
    "    EMBEDDING_RATE as SEMANTIC_HZ,\n",
    ")\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "semantic_model_filepath=\"/home/georg/notebooks/gpu_nb/tmp/mert_25.pt\"\n",
    "semantic_clusters_filepath=\"/home/georg/notebooks/gpu_nb/tmp/mert_25_2x4k.npy\"\n",
    "_ = preload_semantic_models(semantic_model_filepath, semantic_clusters_filepath)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.audio import Audio\n",
    "\n",
    "for filepath in tqdm(source_filepaths):\n",
    "    meta_id = os.path.basename(filepath).split(\".\")[0]\n",
    "    out_filepath = os.path.join(output_dir, f\"{meta_id}__corrupted.mp3\")\n",
    "    audio = Audio.from_file(out_filepath)\n",
    "\n",
    "    # also encode semantic codes\n",
    "    semantic_codes = encode_semantic(\n",
    "        audio.convert(sample_rate=24_000, byte_width=2, n_channels=1)\n",
    "    ).astype(np.int64)[:, 0]\n",
    "    #semantic_codes = torch.from_numpy(semantic_codes)\n",
    "    \n",
    "    # save the semantic codes\n",
    "    semantic_codes_filepath = os.path.join(output_dir, f\"{meta_id}__corrupted.npz\")\n",
    "    np.savez(semantic_codes_filepath, semantic_codes=semantic_codes)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# here we need to iterate over the pass tracks and read the audio from s3\n",
    "# once we get the audio we need to apply lowpass filter, and then ideally extract semantic codes and save them out \n",
    "import polars as pl\n",
    "import os\n",
    "import json\n",
    "base_dir = \"/app/suno/data/diffusion_mix/dac_vae_tuned_25hz/\"\n",
    "\n",
    "metas_tr_filepath = os.path.join(base_dir, \"metas_tr.jsonl\")\n",
    "metas_tr = pl.read_ndjson(metas_tr_filepath)\n",
    "\n",
    "metas_val_filepath = os.path.join(base_dir, \"metas_val.jsonl\")\n",
    "metas_val = pl.read_ndjson(metas_val_filepath)\n",
    "\n",
    "print(metas_tr[0][\"id\"])\n",
    "\n",
    "info_tr_filepath = os.path.join(base_dir, \"info_tr_t4.json\")\n",
    "with open(info_tr_filepath, \"r\") as f:\n",
    "    info_tr = json.load(f)\n",
    "\n",
    "print(len(info_tr[\"diffusion_mix_fix\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Access DataFrame rows properly\n",
    "for meta in metas_val.iter_rows(named=True):\n",
    "    # meta is already a dict when using iter_rows(named=True)\n",
    "    print(meta[\"id\"])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# to create a memmap for sft we will select the higheset scoring upsample_id for each base_s3_id\n",
    "# we also need to grab the correct vae latents and semantic codes and text prompt\n",
    "from tqdm import tqdm\n",
    "from suno_utils.utils.text import write_jsonl\n",
    "import gc\n",
    "import sys\n",
    "import shutil\n",
    "\n",
    "SEMANTIC_RATE_HZ = 25\n",
    "CHUNK_SIZE_S = 30\n",
    "CHUNK_SIZE = int(CHUNK_SIZE_S * SEMANTIC_RATE_HZ)\n",
    "OUT_DATA_DIR = base_dir\n",
    "\n",
    "if not os.path.exists(OUT_DATA_DIR):\n",
    "    os.makedirs(OUT_DATA_DIR, exist_ok=True)\n",
    "else:\n",
    "    #shutil.rmtree(OUT_DATA_DIR)\n",
    "    os.makedirs(OUT_DATA_DIR, exist_ok=True)\n",
    "\n",
    "for dset_type in [\"tr\"]:\n",
    "\n",
    "    if dset_type == \"tr\":\n",
    "        metas = metas_tr\n",
    "    else:\n",
    "        metas = metas_val\n",
    "\n",
    "    new_metas = []\n",
    "\n",
    "    out_mm_semantic_filepath = os.path.join(OUT_DATA_DIR, f\"data_corrupted_semantic_{dset_type}.bin\")\n",
    "    out_metas_filepath = os.path.join(OUT_DATA_DIR, f\"metas_corrupted_{dset_type}.jsonl\")\n",
    "\n",
    "    n_offs_v = 0\n",
    "    n_offs_s = 0\n",
    "    to_write_len_v = 0\n",
    "    to_write_len_s = 0\n",
    "    total_hours = 0  # Counter for total hours of audio\n",
    "\n",
    "    out_mm_semantic = np.memmap(\n",
    "        out_mm_semantic_filepath, dtype=np.uint16, mode=\"w+\", shape=(1,)\n",
    "    )\n",
    "\n",
    "    # clear the metas file\n",
    "    with open(out_metas_filepath, \"w\") as f:\n",
    "        f.write(\"\")\n",
    "\n",
    "    # Create a tqdm progress bar with hours counter\n",
    "    pbar = tqdm(metas.iter_rows(named=True))\n",
    "    pbar.set_description(\"Hours: 0.00\")\n",
    "\n",
    "    for idx, (meta) in enumerate(pbar):\n",
    "\n",
    "        # load semantic codes from disk\n",
    "        semantic_codes_filepath = os.path.join(output_dir, f\"{meta['id']}__corrupted.npz\")\n",
    "\n",
    "        if not os.path.exists(semantic_codes_filepath):\n",
    "            continue\n",
    "\n",
    "        semantic_data = np.load(semantic_codes_filepath)[\"semantic_codes\"]\n",
    "\n",
    "        num_chunks = semantic_data.shape[0] // CHUNK_SIZE\n",
    "\n",
    "        to_write_len_s = semantic_data[:750].size * num_chunks\n",
    "        \n",
    "        if to_write_len_s == 0:\n",
    "            continue\n",
    "        \n",
    "        out_mm_semantic = np.memmap(\n",
    "            out_mm_semantic_filepath,\n",
    "            dtype=np.uint16,\n",
    "            mode=\"r+\",\n",
    "            shape=(n_offs_s + to_write_len_s,),\n",
    "        )\n",
    "\n",
    "        # Add to total hours counter\n",
    "        audio_duration_hours = (num_chunks * CHUNK_SIZE_S) / 3600\n",
    "        total_hours += audio_duration_hours\n",
    "        \n",
    "        # Update progress bar description with current total hours\n",
    "        pbar.set_description(f\"Hours: {total_hours:.2f}\")\n",
    "\n",
    "        for i in range(num_chunks):\n",
    "            # create a new meta\n",
    "            new_meta = {\n",
    "                \"id\": meta[\"id\"],\n",
    "                \"start_s\": i*CHUNK_SIZE_S,\n",
    "                \"end_s\": (i+1)*CHUNK_SIZE_S,\n",
    "                \"original_duration_s\": semantic_data.shape[0] / SEMANTIC_RATE_HZ,\n",
    "                \"n_vae_tokens\": CHUNK_SIZE,\n",
    "                \"n_semantic_tokens\": CHUNK_SIZE,\n",
    "            }\n",
    "            new_metas.append(new_meta)\n",
    "\n",
    "            semantic_chunk = semantic_data[i*CHUNK_SIZE:(i+1)*CHUNK_SIZE]\n",
    "\n",
    "            out_mm_semantic[n_offs_s : n_offs_s + semantic_chunk.size] = semantic_chunk.reshape(\n",
    "                -1,\n",
    "            )\n",
    "            n_offs_s += semantic_chunk.size\n",
    "\n",
    "    print(f\"Total hours of audio added: {total_hours:.2f} for {dset_type} set\")\n",
    "  \n",
    "    write_jsonl(\n",
    "        new_metas,\n",
    "        os.path.join(out_metas_filepath),\n",
    "        do_append=True\n",
    "    )\n",
    "\n",
    "    out_mm_semantic.flush()\n",
    "    del out_mm_semantic, f\n",
    "    gc.collect()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load metas and semantic codes\n",
    "test_metas = read_jsonl(os.path.join(OUT_DATA_DIR, f\"metas_corrupted_tr.jsonl\"))\n",
    "test_semantic_codes = np.memmap(os.path.join(OUT_DATA_DIR, f\"data_corrupted_semantic_tr.bin\"), dtype=np.uint16, mode=\"r\")\n",
    "\n",
    "# reshape the semantic codes to a 2d array\n",
    "test_semantic_codes = test_semantic_codes.reshape(-1, 750)\n",
    "\n",
    "# create a dataframe from the metas\n",
    "test_df = pl.DataFrame(test_metas)\n",
    "\n",
    "# add the semantic codes to the dataframe\n",
    "test_df = test_df.with_columns(pl.Series(name=\"semantic_codes\", values=test_semantic_codes))\n",
    "\n",
    "print(len(test_metas))\n",
    "print(test_semantic_codes.shape)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_env",
   "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}