{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"4\"\n",
    "\n",
    "import sys\n",
    "import json\n",
    "import torch\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "sys.path.insert(0, \"/home/christian/code/neon/sunoDiff/\") # make sure up to date with main"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load pretrained base diffusion model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 100hz\n",
    "DIT_MODEL_FILEPATH = \"s3://suno-data/tony/tmp/dit_1b_30s_1ergguj2.pt\"\n",
    "#DIT_MODEL_FILEPATH = \"/home/christian/code/neon/stable-audio-tools/harmonai_train/5jvuj882/checkpoints/last.ckpt\"\n",
    "DIT_CONFIG_FILEPATH = \"s3://suno-data/tony/tmp/dit_1b_30s_ctx_2.json\"\n",
    "\n",
    "from suno_utils.audio import Audio\n",
    "from generation import preload_models, generate, _retrieve_models\n",
    "_ = preload_models(\n",
    "    tokenizer_filepath=\"/home/georg/notebooks/gpu_nb/tmp/tokenizer_60k.json\",\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",
    "    codec_filepath=\"/home/georg/notebooks/gpu_nb/tmp/100hz_vae_peaq_kl_0.005.pth\",\n",
    "    dit_model_filepath=DIT_MODEL_FILEPATH,\n",
    "    dit_config_filepath=DIT_CONFIG_FILEPATH,\n",
    "    weights_precision=torch.float16,\n",
    "    compile=True,\n",
    ")\n",
    "\n",
    "models = _retrieve_models()\n",
    "model_duration_s = 30\n",
    "if models[\"dit_model\"].ctx_len is not None:\n",
    "    model_duration_s = 6 * 60\n",
    "else:\n",
    "    model_duration_s = models[\"dit_model\"].block_size // models[\"dit_model\"].io_hz\n",
    "duration_s = 2*60 if model_duration_s >= 2*60 else 30"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load preference tagging models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ear.utils import load_audio, apply_normalization\n",
    "from ear.system import EarSystem\n",
    "# load pretrained ear model\n",
    "ckpt_path = \"/home/christian/code/christian/checkpoints/w5p4nhzn-epoch=27.cpkt\"\n",
    "ckpt_path = \"/app/suno/christian/ear-logs/ear/88ufjk6s/checkpoints/last.ckpt\"\n",
    "\n",
    "if not os.path.isfile(ckpt_path):\n",
    "    os.system(\n",
    "        f\"aws s3 cp s3://suno-data/christian/ear/w5p4nhzn-epoch=27.cpkt /home/christian/code/christian/checkpoints\"\n",
    "    )\n",
    "system = EarSystem.load_from_checkpoint(ckpt_path)\n",
    "system.cuda()\n",
    "system.eval()\n",
    "\n",
    "NUM_FRAMES = 131072\n",
    "\n",
    "# load reference audio used for quality comparision\n",
    "# ref_dir = \"/app/suno/christian/data/codec_audio/reference-audio-wav-mono-24khz/\"\n",
    "# ref_filepaths = glob.glob(os.path.join(ref_dir, \"*.input.wav\"))\n",
    "# ref_filepaths = np.random.choice(ref_filepaths, num_compare)\n",
    "ref_filepaths = [\n",
    "    \"/home/christian/audio/reference-audio-wav-mono-24khz/02 Dreams.wav\",\n",
    "    \"/home/christian/audio/reference-audio-wav-mono-24khz/01 Mario Takes A Walk.wav\",\n",
    "    \"/home/christian/audio/reference-audio-wav-mono-24khz/02 Freddie Freeloader.wav\",\n",
    "    \"/home/christian/audio/reference-audio-wav-mono-24khz/09 Sounds Like Hallelujah.wav\",\n",
    "    \"/home/christian/audio/reference-audio-wav-mono-24khz/03 Your New Aesthetic.wav\",\n",
    "]\n",
    "\n",
    "ref_audios = [\n",
    "    load_audio(\n",
    "        filepath,\n",
    "        num_frames=NUM_FRAMES,\n",
    "        target_sample_rate=system.hparams.sample_rate,\n",
    "    )\n",
    "    for filepath in ref_filepaths\n",
    "]\n",
    "ref_audios = torch.stack(ref_audios)\n",
    "print(\"ref_audios\", ref_audios.shape)\n",
    "ref_audio = ref_audios.cuda()\n",
    "\n",
    "# first precompute the reference embeddings\n",
    "ref_embeds = system.embed(ref_audios)\n",
    "print(\"ref_embeds\", ref_embeds.shape)\n",
    "\n",
    "\n",
    "def evaluate_ear(audio: torch.Tensor, ref_embeds: torch.Tensor):\n",
    "    bs = audio.shape[0]\n",
    "    num_refs = ref_embeds.shape[0]\n",
    "\n",
    "    # first, embed the audio that will be evaluated\n",
    "    with torch.no_grad():\n",
    "        eval_embeds = system.embed(audio)\n",
    "\n",
    "    # aggregate the eval_embed\n",
    "    eval_embeds = eval_embeds.mean(dim=1, keepdim=True)\n",
    "    ref_embeds = ref_embeds.mean(dim=1, keepdim=True)\n",
    "\n",
    "    print(eval_embeds.shape, ref_embeds.shape)\n",
    "\n",
    "    # eval_embeds has shape (bs, embed_dim)\n",
    "    # ref_embeds has shape (num_refs, embed_dim)\n",
    "    # now copy the eval and reference embeds to evaluate against all\n",
    "    ref_embeds = ref_embeds.repeat(bs, 1, 1)\n",
    "    eval_embeds = eval_embeds.repeat(num_refs, 1, 1)\n",
    "\n",
    "    print(eval_embeds.shape, ref_embeds.shape)\n",
    "\n",
    "    # concat embeds into singular tensors\n",
    "    embeds = torch.cat((eval_embeds, ref_embeds), dim=-1)\n",
    "    # print(\"embeds\", embeds.shape)\n",
    "\n",
    "    # no run through the projection to make predictions\n",
    "    with torch.no_grad():\n",
    "        pref_preds = system.pref_classifier(embeds)\n",
    "        quant_preds = system.quant_classifier(embeds)\n",
    "\n",
    "    print(pref_preds.shape, quant_preds.shape)\n",
    "\n",
    "    # get a final score by taking mean across seq of preds\n",
    "    pref_preds = pref_preds.mean(dim=1).squeeze(1)\n",
    "    #quant_preds = quant_preds.mean(dim=1).squeeze(1)\n",
    "    pref = torch.sigmoid(pref_preds)\n",
    "    #quant = torch.argmax(quant_preds, dim=1).float()\n",
    "\n",
    "    # aggregate predictions across the reference recordings\n",
    "    prefs = pref.view(bs, -1).mean()\n",
    "    #quants = quant.view(bs, -1).mean()\n",
    "    # scores = -((prefs * 2) - 1) * (quants + 1)\n",
    "\n",
    "    return prefs\n",
    "\n",
    "def compare_quality(system, audio_a: torch.Tensor, audio_b: torch.Tensor):\n",
    "    \"\"\" Compare the quality of two audio files using the trained model.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    system : EarSystem\n",
    "        the trained model\n",
    "    audio_a : torch.Tensor\n",
    "        audio tensor of shape (2, num_frames)\n",
    "    audio_b : torch.Tensor  \n",
    "        audio tensor of shape (2, num_frames)\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    pref : torch.Tensor\n",
    "        preference prediction\n",
    "\n",
    "    quant : torch.Tensor\n",
    "        quantification prediction\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    # move audio_a and audio_b to same device as system parameters\n",
    "    audio_a = audio_a.to(system.device)\n",
    "    audio_b = audio_b.to(system.device)\n",
    "\n",
    "    # first, embed the audio that will be evaluated\n",
    "    with torch.no_grad():\n",
    "        embeds_a = system.embed(audio_a)\n",
    "        embeds_b = system.embed(audio_b)\n",
    "\n",
    "    # aggregate embeddings over time with a moving mean of frame size\n",
    "    embeds_a = torch.nn.functional.adaptive_avg_pool1d(embeds_a.permute(0, 2, 1), 137).permute(0, 2, 1)\n",
    "    embeds_b = torch.nn.functional.adaptive_avg_pool1d(embeds_b.permute(0, 2, 1), 137).permute(0, 2, 1)\n",
    "\n",
    "    # concat embeds into singular tensors\n",
    "    embeds = torch.cat((embeds_a, embeds_b), dim=-1)\n",
    "\n",
    "    # no run through the projection to make predictions\n",
    "    with torch.no_grad():\n",
    "        pref_preds = system.pref_classifier(embeds)\n",
    "        quant_preds = system.quant_classifier(embeds)\n",
    "\n",
    "    print(pref_preds.shape, quant_preds.shape)\n",
    "\n",
    "    # get a final score by taking mean across seq of preds and chunks\n",
    "    pref_preds = pref_preds.mean(dim=1).mean(dim=0)\n",
    "    quant_preds = quant_preds.mean(dim=1).mean(dim=0)\n",
    "    pref = torch.sigmoid(pref_preds)\n",
    "    quant = torch.argmax(quant_preds, dim=0).float()\n",
    "\n",
    "    return pref, quant\n",
    "\n",
    "\n",
    "def prepare_audio(\n",
    "    audio: Audio,\n",
    "    num_frames: int,\n",
    "    start_s: float = None,\n",
    "    end_s: float = None,\n",
    "):\n",
    "    sample_rate = audio.sample_rate\n",
    "    audio = torch.from_numpy(audio.array_float)\n",
    "\n",
    "    if audio.shape[0] != 2:\n",
    "        audio = audio.repeat(2, 1)\n",
    "\n",
    "    # crop audio based on metadata example\n",
    "    if start_s is not None and end_s is not None:\n",
    "        start_frame = int(start_s * sample_rate)\n",
    "        end_frame = int(end_s * sample_rate)\n",
    "        audio = audio[:, start_frame:end_frame]\n",
    "\n",
    "    # if the file is long, only take part of it\n",
    "    if audio.shape[-1] > (sample_rate * 120):\n",
    "        audio = audio[:, : sample_rate * 120]\n",
    "\n",
    "    # downmix and resample decoded audio to 24khz\n",
    "    audio = torchaudio.functional.resample(audio, sample_rate, 24_000)\n",
    "\n",
    "    # pad by repeating the signal if shorter than window\n",
    "    if audio.shape[-1] < num_frames:\n",
    "        pad_size = num_frames - audio.shape[-1]\n",
    "        audio = torch.nn.functional.pad(audio, (1, pad_size), mode=\"replicate\")\n",
    "\n",
    "    # chunk into non-overlapping blocks of num_frames\n",
    "    audio_chunks = []\n",
    "    num_chunks = audio.shape[-1] // num_frames\n",
    "    for n in range(num_chunks):\n",
    "        start_idx = n * num_frames\n",
    "        end_idx = start_idx + num_frames\n",
    "        audio_chunks.append(audio[:, start_idx:end_idx])\n",
    "\n",
    "    # loudness norm\n",
    "    meter = pyln.Meter(24_000)\n",
    "\n",
    "    for audio_chunk_idx in range(len(audio_chunks)):\n",
    "        x_lufs_db = meter.integrated_loudness(audio.T.numpy())\n",
    "        if x_lufs_db == -float(\"inf\"):\n",
    "            gain_lin = 1.0\n",
    "        else:\n",
    "            delta_lufs_db = -20.0 - x_lufs_db\n",
    "            gain_lin = 10.0 ** (np.clip(delta_lufs_db, a_min=-120, a_max=48.0) / 20.0)\n",
    "        audio_chunks[audio_chunk_idx] *= gain_lin\n",
    "\n",
    "    # take the last chunk keeping the list\n",
    "    #audio_chunks = audio_chunks[-1:]\n",
    "\n",
    "    return torch.stack(audio_chunks)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load discriminator\n",
    "import pyloudnorm as pyln\n",
    "from dac.model.discriminator2 import Discriminator as Discriminator_import\n",
    "device=\"cuda:0\"\n",
    "model_name = \"100hz_128_vae_peaq_kl_0.005\"\n",
    "ckpt_path = f\"/app/suno/christian/checkpoints/dac/100hz_vae_peaq_kl_0.005/best/discriminator/weights.pth\"\n",
    "\n",
    "if not os.path.isfile(ckpt_path):\n",
    "    raise ValueError(f\"Checkpoint not found: {ckpt_path}\") \n",
    "\n",
    "print(f\"Loading model {model_name} from {ckpt_path}\")\n",
    "sd = torch.load(ckpt_path)\n",
    "sd[\"metadata\"][\"kwargs\"] = {\n",
    "    k: v for k, v in sd[\"metadata\"][\"kwargs\"].items() if k in Discriminator_import.__init__.__code__.co_varnames\n",
    "}\n",
    "model_disc = Discriminator_import(**sd[\"metadata\"][\"kwargs\"])\n",
    "model_disc.load_state_dict(sd[\"state_dict\"])\n",
    "model_disc.eval()\n",
    "model_disc.to(device)\n",
    "\n",
    "def preprocess(y):\n",
    "    # Remove DC offset\n",
    "    y = y - y.mean(dim=-1, keepdims=True)\n",
    "    # Peak normalize the volume of input audio\n",
    "    y = 0.8 * y / (y.abs().max(dim=-1, keepdim=True)[0] + 1e-9)\n",
    "    return y\n",
    "\n",
    "def evaluate_discriminator(model, audio: torch.Tensor):\n",
    "    with torch.no_grad():\n",
    "        y = model(preprocess(audio.unsqueeze(0)))\n",
    "    loss_g = 0\n",
    "    for y_elem in y:\n",
    "        loss_g += torch.mean((1 - y_elem[-1]) ** 2)\n",
    "    return loss_g\n",
    "\n",
    "def loudness_normalize(audio: torch.Tensor, taget_loudness: float = -16.0, sr: int = 48000):\n",
    "    meter = pyln.Meter(sr) # create loudness meter\n",
    "    loudness = meter.integrated_loudness(audio.permute(1, 0).numpy())\n",
    "    loudness_delta = taget_loudness - loudness\n",
    "    loudness_delta_ln = 10 ** (loudness_delta / 20)\n",
    "    audio = audio * loudness_delta_ln\n",
    "    return audio\n",
    "\n",
    "def compare_discriminator(model, audio_a: torch.Tensor, audio_b: torch.Tensor):\n",
    "    with torch.no_grad():\n",
    "        y_a = model(preprocess(audio_a.unsqueeze(0)))\n",
    "        y_b = model(preprocess(audio_b.unsqueeze(0)))\n",
    "    loss_g_a = 0\n",
    "    loss_g_b = 0\n",
    "    for y_elem in y_a:\n",
    "        loss_g_a += torch.mean((1 - y_elem[-1]) ** 2)\n",
    "    for y_elem in y_b:\n",
    "        loss_g_b += torch.mean((1 - y_elem[-1]) ** 2)\n",
    "    return loss_g_a, loss_g_b\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Read npz and json files from 30b gens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = \"/home/tony/Data/Preference/30b_v5/interesting_clips_v4_t_5_20241015_full_with_cer_pos_gen.pkl\"\n",
    "npz_root_dir = \"/app/suno/data/dpo/30b_npz/\"\n",
    "json_root_dir = \"/app/suno/data/dpo/30b_json/\"\n",
    "df = pd.read_pickle(data_path)\n",
    "print(f\"df shape: {df.shape}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchaudio\n",
    "import uuid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Upsample generations (pairs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# upsample configurations\n",
    "steps = 32\n",
    "text_cfg = 5.0\n",
    "ctx_cfg = 1.5\n",
    "\n",
    "# create a dictionary that we will use to create a new dataframe\n",
    "# the dictionary will have the following columns:\n",
    "# id_x, tags, a_id, b_id, a_pref, b_pref\n",
    "\n",
    "name = \"interesting_clips_v4_t_5_20241015_full_with_cer_pos_gen_upsample\"\n",
    "\n",
    "data = {\n",
    "    \"id_x\": [],\n",
    "    \"tags\": [],\n",
    "    \"id_upsample\": [],\n",
    "    \"preference\": [],\n",
    "    \"pref_prob\": [],\n",
    "}\n",
    "\n",
    "out_dir = \"/app/suno/christian/data/dpo_diff_syn_data\"\n",
    "os.makedirs(out_dir, exist_ok=True)\n",
    "\n",
    "\n",
    "# Upsample generations\n",
    "for example_idx, (row_idx, row) in enumerate(tqdm(df.iterrows(), total=len(df))):\n",
    "    print(f\"example_idx: {example_idx} {row.id_x}\")\n",
    "    npz_path = os.path.join(npz_root_dir, f\"{row.id_x}.npz\")\n",
    "    json_path = os.path.join(json_root_dir, f\"{row.id_x}_hoot.json\")\n",
    "    npz_data = np.load(npz_path)\n",
    "    semantic_codes = npz_data[\"v4.0_raw\"][:, 0]\n",
    "    tags = row.tags\n",
    "    print(tags)\n",
    "    print(semantic_codes.shape)\n",
    "\n",
    "    # for now, just do the first 30 seconds\n",
    "    semantic_codes = semantic_codes[:750]\n",
    "\n",
    "    with open(json_path, \"r\") as f:\n",
    "        aligned_lyrics = json.load(f)\n",
    "    \n",
    "    # generate audio file with two different seeds\n",
    "\n",
    "\n",
    "    pref_agree = False\n",
    "    num_attempts = 0\n",
    "\n",
    "    while not pref_agree:\n",
    "        audios = []\n",
    "        scores = []\n",
    "        uuids = []\n",
    "        seeds = np.random.randint(0, 1000000, size=2)\n",
    "        for seed_idx, seed in enumerate(seeds):  \n",
    "            uid = str(uuid.uuid4())\n",
    "            uuids.append(uid)\n",
    "            if seed_idx == 0:\n",
    "                print(f\"Generating audio_a: {seed}\")\n",
    "            else:\n",
    "                print(f\"Generating audio_b: {seed}\")\n",
    "            pred_audio = generate(\n",
    "                semantic_codes=semantic_codes,\n",
    "                aligned_lyrics=aligned_lyrics,\n",
    "                tags=tags,\n",
    "                text_cfg_coef=text_cfg,\n",
    "                ctx_cfg_coef=ctx_cfg,\n",
    "                steps=steps,\n",
    "                seed=seed,\n",
    "            ).normalize_volume()\n",
    "            pred_audio.play()\n",
    "            audios.append(pred_audio)\n",
    "\n",
    "\n",
    "        #pred_audio_a = torch.from_numpy(audios[0].array_float).cuda()\n",
    "        #pred_audio_b = torch.from_numpy(audios[1].array_float).cuda()\n",
    "\n",
    "        # evaluate audio quality\n",
    "        pref, quant = compare_quality(system, prepare_audio(audios[0], NUM_FRAMES), prepare_audio(audios[1], NUM_FRAMES))\n",
    "        #scores.append(pref)\n",
    "\n",
    "        if pref < 0.4:\n",
    "            print(f\"pref={pref} audio_a is preferred by ear\")\n",
    "            pref_agree = True\n",
    "            # add one row for positive, and one row for negative\n",
    "            data[\"id_x\"].append(row.id_x)\n",
    "            data[\"tags\"].append(tags)\n",
    "            data[\"id_upsample\"].append(uuids[0])\n",
    "            data[\"preference\"].append(True)\n",
    "            data[\"pref_prob\"].append(pref)\n",
    "\n",
    "            data[\"id_x\"].append(row.id_x)\n",
    "            data[\"tags\"].append(tags)\n",
    "            data[\"id_upsample\"].append(uuids[1])\n",
    "            data[\"preference\"].append(False)\n",
    "            data[\"pref_prob\"].append(pref)\n",
    "\n",
    "        elif pref > 0.6:\n",
    "            print(f\"pref={pref} audio_b is preferred by ear\")\n",
    "            pref_agree = True\n",
    "            data[\"id_x\"].append(row.id_x)\n",
    "            data[\"tags\"].append(tags)\n",
    "            data[\"id_upsample\"].append(uuids[1])\n",
    "            data[\"preference\"].append(True)\n",
    "            data[\"pref_prob\"].append(pref)\n",
    "\n",
    "            data[\"id_x\"].append(row.id_x)\n",
    "            data[\"tags\"].append(tags)\n",
    "            data[\"id_upsample\"].append(uuids[0])\n",
    "            data[\"preference\"].append(False)\n",
    "            data[\"pref_prob\"].append(pref)\n",
    "        else:\n",
    "            print(f\"pref={pref} discriminator and ear disagree!\")\n",
    "            pref_agree = False\n",
    "\n",
    "        # evaluate discriminator\n",
    "        if False:\n",
    "            loss_g_a, loss_g_b = compare_discriminator(model_disc, pred_audio_a, pred_audio_b)\n",
    "            if loss_g_a > loss_g_b:\n",
    "                print(f\"audio_a (loss_g_a: {loss_g_a:0.2f}) is preferred over audio_b (loss_g_b: {loss_g_b:0.2f}) by discriminator\")\n",
    "            else:\n",
    "                print(f\"audio_b (loss_g_b: {loss_g_b:0.2f}) is preferred over audio_a (loss_g_a: {loss_g_a:0.2f}) by discriminator\")\n",
    "\n",
    "            delta = torch.abs(loss_g_a - loss_g_b)\n",
    "            print(f\"discriminator prefers audio_a over audio_b by {delta}\")\n",
    "            if delta > 0.1:\n",
    "                pref_agree = True\n",
    "\n",
    "        num_attempts += 1\n",
    "        if num_attempts > 5:\n",
    "            print(\"Failed to get preference\")\n",
    "            continue\n",
    "\n",
    "        # save data\n",
    "        df = pd.DataFrame(data)\n",
    "        df.to_csv(os.path.join(out_dir, f\"{name}.csv\"), index=False)\n",
    "\n",
    "        # check if discriminator and ear agree\n",
    "        #if (pref < 0.5 and loss_g_a > loss_g_b):\n",
    "        #    a_pref = True\n",
    "        #    b_pref = False\n",
    "        #    print(\"agree that audio_a is preferred\")\n",
    "        #     pref_agree = True\n",
    "        #elif (pref > 0.5 and loss_g_a < loss_g_b):\n",
    "        #    a_pref = False\n",
    "         #   b_pref = True\n",
    "        #    print(\"agree that audio_b is preferred\")\n",
    "        #    pref_agree = True\n",
    "        #else:\n",
    "       #     print(\"discriminator and ear disagree!\")\n",
    "       #     pref_agree = False\n",
    "    # we need to save out some stuff for dpo\n",
    "    print()\n",
    "\n",
    "    if example_idx > 10:\n",
    "        break\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Upsample generations (singles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# upsample configurations\n",
    "steps = 32\n",
    "text_cfg = 5.0\n",
    "ctx_cfg = 1.5\n",
    "\n",
    "# create a dictionary that we will use to create a new dataframe\n",
    "# the dictionary will have the following columns:\n",
    "# id_x, tags, a_id, b_id, a_pref, b_pref\n",
    "\n",
    "name = \"interesting_clips_v4_t_5_20241015_full_with_cer_pos_gen_upsample\"\n",
    "\n",
    "data = {\n",
    "    \"id_x\": [],\n",
    "    \"tags\": [],\n",
    "    \"id_upsample\": [],\n",
    "    \"preference\": [],\n",
    "    \"pref_prob\": [],\n",
    "}\n",
    "\n",
    "out_dir = \"/app/suno/christian/data/dpo_diff_syn_data_singles\"\n",
    "os.makedirs(out_dir, exist_ok=True)\n",
    "\n",
    "\n",
    "quality_labels = []\n",
    "\n",
    "# Upsample generations\n",
    "for example_idx, (row_idx, row) in enumerate(tqdm(df.iterrows(), total=len(df))):\n",
    "    print(f\"example_idx: {example_idx} {row.id_x}\")\n",
    "    npz_path = os.path.join(npz_root_dir, f\"{row.id_x}.npz\")\n",
    "    json_path = os.path.join(json_root_dir, f\"{row.id_x}_hoot.json\")\n",
    "    npz_data = np.load(npz_path)\n",
    "    semantic_codes = npz_data[\"v4.0_raw\"][:, 0]\n",
    "    tags = row.tags\n",
    "    print(tags)\n",
    "    print(semantic_codes.shape)\n",
    "\n",
    "    # for now, just do the first 30 seconds\n",
    "    semantic_codes = semantic_codes[:750]\n",
    "\n",
    "    with open(json_path, \"r\") as f:\n",
    "        aligned_lyrics = json.load(f)\n",
    "\n",
    "    uid = str(uuid.uuid4())\n",
    "\n",
    "    pred_audio = generate(\n",
    "        semantic_codes=semantic_codes,\n",
    "        aligned_lyrics=aligned_lyrics,\n",
    "        tags=tags,\n",
    "        text_cfg_coef=text_cfg,\n",
    "        ctx_cfg_coef=ctx_cfg,\n",
    "        steps=steps,\n",
    "        seed=seed,\n",
    "    ).normalize_volume()\n",
    "    #pred_audio.play()\n",
    "    #audios.append(pred_audio)\n",
    "\n",
    "    prep_pred_audio = prepare_audio(pred_audio, NUM_FRAMES).cuda()\n",
    "    # evaluate audio quality\n",
    "    pref = evaluate_ear(prep_pred_audio, ref_embeds)\n",
    "\n",
    "    print(f\"pref={pref} ear\")\n",
    "\n",
    "    quality_labels.append(pref)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "quality_labels = [label.item() for label in quality_labels]\n",
    "\n",
    "plt.hist(quality_labels, bins=100)\n",
    "plt.show()\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
}
