{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"4\"\n",
    "import math\n",
    "import random\n",
    "import torch\n",
    "import torchaudio\n",
    "\n",
    "from suno_utils.utils.text import read_jsonl\n",
    "\n",
    "from suno_utils.tasks.mert_25 import (\n",
    "    encode as encode_semantic,\n",
    "    preload_models as preload_semantic_models\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "preload_semantic_models(\n",
    "    checkpoint_filepath=\"s3://suno-data/georg/models/semantic/mert_25.pt\",\n",
    "    centroids_filepath=\"s3://suno-data/georg/models/semantic/mert_25_2x4k.npy\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def mdct(x):\n",
    "    N = x.shape[-1]\n",
    "    n = torch.arange(N, device=x.device)\n",
    "    k = torch.arange(N // 2, device=x.device)\n",
    "\n",
    "    arg = (\n",
    "        (math.pi / (2 * N))\n",
    "        * ((2 * n + 1 + N // 2).view(-1, 1))\n",
    "        * ((2 * k + 1).view(1, -1))\n",
    "    )\n",
    "    mdct_matrix = torch.cos(arg) * (2.0 / N) ** 0.5\n",
    "\n",
    "    return torch.matmul(x, mdct_matrix)\n",
    "\n",
    "\n",
    "def imdct(X):\n",
    "    half_N = X.shape[-1]\n",
    "    N = half_N * 2\n",
    "    n = torch.arange(N, device=X.device)\n",
    "    k = torch.arange(half_N, device=X.device)\n",
    "\n",
    "    arg = (\n",
    "        (math.pi / (2 * N))\n",
    "        * ((2 * n + 1 + N // 2).view(-1, 1))\n",
    "        * ((2 * k + 1).view(1, -1))\n",
    "    )\n",
    "    imdct_matrix = torch.cos(arg) * (2.0 / N) ** 0.5\n",
    "\n",
    "    return torch.matmul(X, imdct_matrix.T) * 2.0\n",
    "\n",
    "@torch.no_grad()\n",
    "def audio_to_mdct_frames(audio, frame_size=1920, midside=False):\n",
    "    batch, channels, samples = audio.shape\n",
    "    hop_size = frame_size // 2\n",
    "\n",
    "    if midside:\n",
    "        mid = (audio[:, 0, :] + audio[:, 1, :]) / 2.0\n",
    "        side = (audio[:, 0, :] - audio[:, 1, :]) / 2.0\n",
    "        audio = torch.stack([mid, side], dim=1)\n",
    "\n",
    "    n_frames = (samples - frame_size) // hop_size + 1\n",
    "\n",
    "    frames = []\n",
    "    for i in range(n_frames):\n",
    "        start = i * hop_size\n",
    "        frame = audio[:, :, start : start + frame_size]\n",
    "        if frame.shape[-1] == frame_size:\n",
    "            frames.append(frame)\n",
    "\n",
    "    frames = torch.stack(frames, dim=2)\n",
    "\n",
    "    window = torch.sin(\n",
    "        torch.pi / frame_size * (torch.arange(frame_size, device=audio.device) + 0.5)\n",
    "    )\n",
    "    frames = frames * window.view(1, 1, 1, -1)\n",
    "\n",
    "    shape = frames.shape\n",
    "    frames_reshaped = frames.reshape(-1, frame_size)\n",
    "    mdct_coeffs = mdct(frames_reshaped)\n",
    "\n",
    "    return mdct_coeffs.reshape(shape[0], shape[1], shape[2], -1)\n",
    "\n",
    "\n",
    "def mdct_frames_to_audio_slow(mdct_coeffs, frame_size=1920, midside=False):\n",
    "    batch, channels, n_frames, half_frame_size = mdct_coeffs.shape\n",
    "    hop_size = frame_size // 2\n",
    "\n",
    "    shape = mdct_coeffs.shape\n",
    "    coeffs_reshaped = mdct_coeffs.reshape(-1, half_frame_size)\n",
    "    frames = imdct(coeffs_reshaped)\n",
    "    frames = frames.reshape(shape[0], shape[1], shape[2], -1)\n",
    "\n",
    "    window = torch.sin(\n",
    "        torch.pi\n",
    "        / frame_size\n",
    "        * (torch.arange(frame_size, device=mdct_coeffs.device) + 0.5)\n",
    "    )\n",
    "    frames = frames * window.view(1, 1, 1, -1)\n",
    "\n",
    "    total_samples = (n_frames - 1) * hop_size + frame_size\n",
    "    output = torch.zeros(batch, channels, total_samples, device=mdct_coeffs.device)\n",
    "\n",
    "    for i in range(n_frames):\n",
    "        start = i * hop_size\n",
    "        output[:, :, start : start + frame_size] += frames[:, :, i]\n",
    "\n",
    "    if midside:\n",
    "        mid = output[:, 0, :]\n",
    "        side = output[:, 1, :]\n",
    "        left = mid + side\n",
    "        right = mid - side\n",
    "        output = torch.stack([left, right], dim=1)\n",
    "\n",
    "    return output\n",
    "\n",
    "@torch.no_grad()\n",
    "def mdct_frames_to_audio(mdct_coeffs, frame_size=1920, midside=False):\n",
    "    batch, channels, n_frames, half_frame_size = mdct_coeffs.shape\n",
    "    hop_size = frame_size // 2\n",
    "\n",
    "    # Reshape and apply IMDCT\n",
    "    shape = mdct_coeffs.shape\n",
    "    coeffs_reshaped = mdct_coeffs.reshape(-1, half_frame_size)\n",
    "    frames = imdct(coeffs_reshaped)\n",
    "    frames = frames.reshape(shape[0], shape[1], shape[2], -1)\n",
    "\n",
    "    # Apply window\n",
    "    window = torch.sin(\n",
    "        torch.pi\n",
    "        / frame_size\n",
    "        * (torch.arange(frame_size, device=mdct_coeffs.device) + 0.5)\n",
    "    )\n",
    "    frames = frames * window.view(1, 1, 1, -1)\n",
    "\n",
    "    # Calculate total samples and create output shape\n",
    "    total_samples = (n_frames - 1) * hop_size + frame_size\n",
    "\n",
    "    # Reshape frames to prepare for folding\n",
    "    frames = frames.permute(0, 1, 3, 2)  # [batch, channels, frame_size, n_frames]\n",
    "    frames = frames.reshape(batch * channels, frame_size, n_frames)\n",
    "\n",
    "    # Use fold operation to overlap-add frames\n",
    "    output = torch.nn.functional.fold(\n",
    "        frames,\n",
    "        output_size=(1, total_samples),\n",
    "        kernel_size=(1, frame_size),\n",
    "        stride=(1, hop_size),\n",
    "    )\n",
    "\n",
    "    # Reshape output to expected dimensions\n",
    "    output = output.view(batch, channels, total_samples)\n",
    "\n",
    "    if midside:\n",
    "        mid = output[:, 0, :]\n",
    "        side = output[:, 1, :]\n",
    "        left = mid + side\n",
    "        right = mid - side\n",
    "        output = torch.stack([left, right], dim=1)\n",
    "\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "codec_encode = audio_to_mdct_frames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DynamicEncodingDataset(torch.utils.data.Dataset):\n",
    "    def __init__(\n",
    "        self,\n",
    "        dataset_dir: str,\n",
    "        metas_filename: str,\n",
    "        mode: str = \"pretraining\",\n",
    "        audio_duration_s: float = 10.02,\n",
    "        audio_sample_rate: int = 48000,\n",
    "        mdct_frame_size: int = 1920,  # hop size is always 1/2 frame size\n",
    "        is_training: bool = True,\n",
    "    ):\n",
    "        self.dataset_dir = dataset_dir\n",
    "        self.mode = mode\n",
    "        self.audio_duration_s = audio_duration_s\n",
    "        self.audio_duration_samples = int(audio_duration_s * audio_sample_rate)\n",
    "        self.mdct_frame_size = mdct_frame_size\n",
    "        self.metas_filename = metas_filename\n",
    "        self.is_training = is_training\n",
    "        # load metas\n",
    "        self.metas = read_jsonl(os.path.join(dataset_dir, metas_filename))\n",
    "        print(f\"Loaded {len(self.metas)} metas\")\n",
    "\n",
    "        # load tokenizer\n",
    "        self.tokenizer = load_tokenizer()\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.metas)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        info = {}\n",
    "        info[\"idx\"] = idx\n",
    "\n",
    "        # get the metadata, load text and tags\n",
    "        meta = self.metas[idx]\n",
    "        tags = meta.get(\"tags\", [])\n",
    "        lyrics = meta.get(\"text\", \"\")\n",
    "        local_filepath = meta[\"local_filepath\"]\n",
    "\n",
    "        # Build condition tensors\n",
    "        if self.is_training:\n",
    "            if random.random() <= 0.1:\n",
    "                text = \"\"\n",
    "            else:\n",
    "                text = \"\" #augment_text_training(tags, lyrics)\n",
    "        else:\n",
    "            text = \"\" #prepare_text_inference(tags, lyrics)\n",
    "        text_codes = self.tokenizer.encode(text).ids[: self.cond_text_len]\n",
    "        text_codes = text_codes + [self.tokenizer.pad_idx] * max(0, self.cond_text_len - len(text_codes))\n",
    "        text_codes = torch.tensor(text_codes).long()\n",
    "        info[\"text_codes\"] = text_codes\n",
    "\n",
    "        # load with torchaudio\n",
    "        audio, sr = torchaudio.load(local_filepath)\n",
    "        assert sr == 48000\n",
    "\n",
    "        # select a random chunk of desired duration\n",
    "        # start the chunk at 30s intervals\n",
    "        # Calculate valid start points at 30s intervals\n",
    "        max_start_time_s = max(0, meta[\"duration_s\"] - self.audio_duration_s)\n",
    "        num_intervals = int(max_start_time_s // 30) + 1  # +1 to include the start at 0\n",
    "        \n",
    "        # Choose a random interval (0, 1, 2, ..., num_intervals-1)\n",
    "        interval_idx = random.randint(0, num_intervals - 1)\n",
    "        start_time_s = interval_idx * 30\n",
    "        print(f\"start_time_s: {start_time_s}\")\n",
    "        \n",
    "        # Ensure we don't go beyond the audio duration\n",
    "        start_time_s = min(start_time_s, max_start_time_s)\n",
    "        start_sample = start_time_s * sr\n",
    "\n",
    "        end_sample = start_sample + self.audio_duration_samples\n",
    "        start_sample = int(start_sample)\n",
    "        end_sample = int(end_sample)\n",
    "        audio = audio[:, start_sample:end_sample]\n",
    "\n",
    "        # now we have to semantic encode, 24khz mono\n",
    "        audio_mert = torchaudio.functional.resample(audio, sr, 24000).mean(dim=0).unsqueeze(0)\n",
    "        semantic_codes = torch.from_numpy(encode_semantic(audio_mert)[:,0])\n",
    "        info[\"semantic_codes\"] = semantic_codes\n",
    "\n",
    "        # now we have to codec encode, 48khz stereo\n",
    "        latents = codec_encode(audio.unsqueeze(0))\n",
    "        print(latents.shape)\n",
    "\n",
    "        # move the l/r channels to the sequence dimension\n",
    "        # latents shape: bs, stereo, seq_len, channels\n",
    "        # reshape to move stereo channels into sequence dimension\n",
    "        bs, stereo, seq_len, channels = latents.shape\n",
    "        latents = latents.reshape(bs, seq_len * stereo, channels).squeeze(0)\n",
    "        latents = latents.permute(1, 0) # shape channels, seq_len\n",
    "        print(latents.shape)\n",
    "        \n",
    "        info[\"padding_mask\"] = torch.ones_like(latents).bool()\n",
    "\n",
    "        # todo:\n",
    "        # if its available, load the previous chunk embeddings\n",
    "\n",
    "        # empty ctx amd infill latents\n",
    "        empty_ctx_vae_embeds = torch.zeros_like(latents)\n",
    "        empty_ctx_vae_mask = torch.zeros_like(latents)\n",
    "        info[\"ctx_vae\"] = empty_ctx_vae_embeds\n",
    "        info[\"ctx_mask\"] = empty_ctx_vae_mask\n",
    "        info[\"infill_ctx_vae\"] = empty_ctx_vae_embeds\n",
    "        info[\"infill_ctx_mask\"] = empty_ctx_vae_mask\n",
    "\n",
    "        return (latents, info)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = DynamicEncodingDataset(\n",
    "    dataset_dir=\"/app2/suno/data/auk_v0/\",\n",
    "    metas_filename=\"metas_v3_val.jsonl\",\n",
    "    mode=\"pretraining\",\n",
    "    audio_duration_s=10.02,\n",
    "    mdct_frame_size=1920,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metas = read_jsonl(os.path.join(dataset.dataset_dir, dataset.metas_filename))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset[10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_diff",
   "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.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}