{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Consolidate midi datasets\n",
    "\n",
    "This notebook consolidates the midi datasets into a single dataset. It also serves as documentation for the midi datasets.\n",
    "\n",
    "It produces a metadata file pointing to the midi and audio files.\n",
    "\n",
    "## Datasets\n",
    "\n",
    "- Infinite synthetic MIDI\n",
    "- Mew\n",
    "\t- `/app2/suno/data/mew/scoretube`\n",
    "\t- `/app2/suno/data/mew/mmd_chunks`\n",
    "\t- Match midi instruments to stems.\n",
    "- A dataset of loosely paired real music with arrangements\n",
    "\t- [Slack](https://suno-main.slack.com/archives/C03A7347CSK/p1719682717375969)\n",
    "\t- 20k songs matched with high confidence\n",
    "\t- [The AI workspace that works for you. | Notion](https://www.notion.so/suno-ai/MooTube-187b01573ccf8065915ecb04af262d6c#187b01573ccf80de941bc4412c900eca)\n",
    "- Hooktheory melodys, 50k hooks only\n",
    "\t- [The AI workspace that works for you. | Notion](https://www.notion.so/suno-ai/hooktheory-91267f75401f4087b506f086abefc65f)\n",
    "- Trombone champ melodys. 4k full songs\n",
    "\t- [The AI workspace that works for you. | Notion](https://www.notion.so/suno-ai/Trombone-Champ-206b01573ccf8009aa1fe22f3440687c)\n",
    "- Piano Maestro\n",
    "- Sara's synthetic VST stems\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# mew\n",
    "\n",
    "Victors's personal stash"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import random\n",
    "from suno_utils.audio import Audio\n",
    "from suno_utils.audio.midi import Midi\n",
    "\n",
    "\n",
    "class MidiPair:\n",
    "    def __init__(self, midi_file: str, audio_file: str):\n",
    "        self.midi_file = midi_file\n",
    "        self.audio_file = audio_file\n",
    "        self.midi = None\n",
    "        self.audio = None\n",
    "\n",
    "    def load_midi(self):\n",
    "        if self.midi is not None:\n",
    "            return\n",
    "        self.midi = Midi.from_path(self.midi_file)\n",
    "\n",
    "    def load_audio(self):\n",
    "        if self.audio is not None:\n",
    "            return\n",
    "        self.audio = Audio.from_file(self.audio_file)\n",
    "\n",
    "    def load_all(self):\n",
    "        self.load_midi()\n",
    "        self.load_audio()\n",
    "\n",
    "    def __str__(self):\n",
    "        return f\"MidiPair(midi_file={self.midi_file}, audio_file={self.audio_file})\"\n",
    "\n",
    "    def __repr__(self):\n",
    "        return self.__str__()\n",
    "\n",
    "    def play(self):\n",
    "        self.load_all()\n",
    "        stereo_audio = self.midi.make_stereo_comparison(self.audio)\n",
    "        stereo_audio.play()\n",
    "\n",
    "\n",
    "scoretube_dir = \"/app2/suno/data/victor/mew/scoretube\"\n",
    "\n",
    "\n",
    "def load_pairs(\n",
    "    dir: str,\n",
    "    audio_ext: str = \"mp3\",\n",
    "    midi_ext: str = \"mid\",\n",
    "    include_parent_dir: bool = False,\n",
    "):\n",
    "    midi_files = glob.glob(os.path.join(dir, \"**\", f\"*.{midi_ext}\"), recursive=True)\n",
    "    audio_files = glob.glob(os.path.join(dir, \"**\", f\"*.{audio_ext}\"), recursive=True)\n",
    "\n",
    "    # Create a mapping of base filenames to audio files\n",
    "    audio_map = {}\n",
    "    for audio_file in audio_files:\n",
    "        base_name = os.path.splitext(os.path.basename(audio_file))[0]\n",
    "        if include_parent_dir:\n",
    "            parent_dir = os.path.basename(os.path.dirname(audio_file))\n",
    "            base_name = os.path.join(parent_dir, base_name)\n",
    "        audio_map[base_name] = audio_file\n",
    "\n",
    "    pairs = []\n",
    "    for midi_file in midi_files:\n",
    "        base_name = os.path.splitext(os.path.basename(midi_file))[0]\n",
    "        if include_parent_dir:\n",
    "            parent_dir = os.path.basename(os.path.dirname(midi_file))\n",
    "            base_name = os.path.join(parent_dir, base_name)\n",
    "        if base_name in audio_map:\n",
    "            pairs.append(MidiPair(midi_file, audio_map[base_name]))\n",
    "\n",
    "    return pairs\n",
    "\n",
    "\n",
    "scoretube_pairs = load_pairs(scoretube_dir)\n",
    "print(f\"Loaded {len(scoretube_pairs)} scoretube pairs\")\n",
    "\n",
    "# play random pair\n",
    "random_pair = random.choice(scoretube_pairs)\n",
    "print(random_pair)\n",
    "# random_pair.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mmd_dir = \"/app2/suno/data/victor/mew/mmd_chunks\"\n",
    "mmd_pairs = load_pairs(mmd_dir)\n",
    "print(f\"Loaded {len(mmd_pairs)} mmd pairs\")\n",
    "\n",
    "# play random pair\n",
    "random_pair = random.choice(mmd_pairs)\n",
    "print(random_pair)\n",
    "# random_pair.play()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## mootube"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "mootube_dir = \"/app2/suno/data/victor/mootube\"\n",
    "\n",
    "mootube_metas = []\n",
    "for line in open(\"/app2/suno/data/victor/mootube/score_ytm_clean.jsonl\"):\n",
    "    meta = json.loads(line)\n",
    "    mootube_metas.append(meta)\n",
    "\n",
    "print(f\"Loaded {len(mootube_metas)} moo tube metas\")\n",
    "mootube_metas[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mootube_metas[0][\"matches\"][0][\"videoId\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mootube_pairs = []\n",
    "for line in open(\"/app2/suno/data/victor/mootube/score_ytm_clean.jsonl\"):\n",
    "    meta = json.loads(line)\n",
    "    midi_file = f\"/app2/suno/data/victor/mootube/midi/{meta['id']}.mid\"\n",
    "    audio_file = (\n",
    "        f\"/app2/suno/data/victor/mootube/audio/{meta['matches'][0]['videoId']}.webm\"\n",
    "    )\n",
    "    mootube_pairs.append(MidiPair(midi_file, audio_file))\n",
    "\n",
    "print(f\"Loaded {len(mootube_pairs)} moo tube pairs\")\n",
    "\n",
    "# play random pair\n",
    "for _ in range(1):\n",
    "    random_pair = random.choice(mootube_pairs)\n",
    "    print(random_pair)\n",
    "    # random_pair.play()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## maestro\n",
    "200 hours of piano"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "maestro_dir = \"/app2/suno/data/victor/maestro/maestro-v3.0.0\"\n",
    "maestro_pairs = load_pairs(maestro_dir, audio_ext=\"wav\", midi_ext=\"midi\")\n",
    "print(f\"Loaded {len(maestro_pairs)} maestro pairs\")\n",
    "\n",
    "# play random pair\n",
    "random_pair = random.choice(maestro_pairs)\n",
    "print(random_pair)\n",
    "# random_pair.play()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## trombone champ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trombone_champ_dir = \"/app2/suno/data/victor/trombone_champ\"\n",
    "trombone_champ_pairs = load_pairs(trombone_champ_dir, audio_ext=\"opus\")\n",
    "print(f\"Loaded {len(trombone_champ_pairs)} trombone champ pairs\")\n",
    "\n",
    "# play random pair\n",
    "random_pair = random.choice(trombone_champ_pairs)\n",
    "print(random_pair)\n",
    "# random_pair.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trombone_champ_vocals_dir = \"/app2/suno/data/victor/trombone_champ_vocals\"\n",
    "trombone_champ_vocals_pairs = load_pairs(trombone_champ_vocals_dir, audio_ext=\"opus\")\n",
    "print(f\"Loaded {len(trombone_champ_vocals_pairs)} trombone champ vocals pairs\")\n",
    "\n",
    "# play random pair\n",
    "random_pair = random.choice(trombone_champ_vocals_pairs)\n",
    "print(random_pair)\n",
    "# random_pair.play()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## hook theory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hook_theory_dir = \"/app2/suno/data/victor/hooktheory\"\n",
    "hooktheory_pairs = load_pairs(hook_theory_dir, audio_ext=\"opus\")\n",
    "print(f\"Loaded {len(hooktheory_pairs)} hook theory pairs\")\n",
    "\n",
    "# play random pair\n",
    "random_pair = random.choice(hooktheory_pairs)\n",
    "print(random_pair)\n",
    "# random_pair.play()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## synthetic "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from tqdm import tqdm\n",
    "\n",
    "synthetic_dir = \"/app2/suno/data/victor/gmd/MIDIs/\"\n",
    "synthetic_pairs = []\n",
    "for i in tqdm(\n",
    "    [\"0\", \"1\", \"2\", \"3\", \"4\", \"5\", \"6\", \"7\", \"8\", \"9\", \"a\", \"b\", \"c\", \"d\", \"e\", \"f\"]\n",
    "):\n",
    "    synthetic_pairs.extend(load_pairs(synthetic_dir + i, audio_ext=\"opus\"))\n",
    "    print(f\"Loaded {len(synthetic_pairs)} synthetic pairs\")\n",
    "\n",
    "# play random pair\n",
    "random_pair = random.choice(synthetic_pairs)\n",
    "print(random_pair)\n",
    "# random_pair.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "random_pair = random.choice(synthetic_pairs)\n",
    "print(random_pair)\n",
    "random_pair.play()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sara stems"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sara_dir = \"/app2/suno/data/victor/sara_midi\"\n",
    "\n",
    "sara_pairs = load_pairs(sara_dir, audio_ext=\"opus\", include_parent_dir=True)\n",
    "print(f\"Loaded {len(sara_pairs)} sara pairs\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# play random pair\n",
    "random_pair = random.choice(sara_pairs)\n",
    "print(random_pair)\n",
    "random_pair.play()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## EDA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Combine all pairs for EDA\n",
    "all_pairs = {\n",
    "    \"maestro\": maestro_pairs,\n",
    "    \"trombone_champ\": trombone_champ_pairs,\n",
    "    \"trombone_champ_vocals\": trombone_champ_vocals_pairs,\n",
    "    \"hooktheory\": hooktheory_pairs,\n",
    "    \"mootube\": mootube_pairs,\n",
    "    \"mmd\": mmd_pairs,\n",
    "    \"scoretube\": scoretube_pairs,\n",
    "    \"synthetic\": synthetic_pairs,\n",
    "}\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Get counts for each dataset\n",
    "dataset_names = list(all_pairs.keys())\n",
    "counts = [len(pairs) for pairs in all_pairs.values()]\n",
    "\n",
    "# Create bar plot\n",
    "plt.figure(figsize=(10, 6))\n",
    "plt.bar(dataset_names, counts)\n",
    "plt.title(\"Dataset Pair Counts\")\n",
    "plt.xlabel(\"Dataset\")\n",
    "plt.ylabel(\"Number of Pairs\")\n",
    "plt.xticks(rotation=45)\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "for dataset, pairs in all_pairs.items():\n",
    "    print(f\"Loaded {len(pairs)} {dataset} pairs\")\n",
    "\n",
    "print(f\"Total pairs: {sum(len(pairs) for pairs in all_pairs.values())}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate duration statistics for each dataset\n",
    "import random\n",
    "from collections import defaultdict\n",
    "\n",
    "duration_stats = defaultdict(list)\n",
    "\n",
    "for dataset_name, pairs in all_pairs.items():\n",
    "    if len(pairs) == 0:\n",
    "        print(f\"Skipping {dataset_name} - no pairs available\")\n",
    "        continue\n",
    "\n",
    "    # Sample up to 10 pairs from each dataset\n",
    "    sample_size = min(10, len(pairs))\n",
    "    sampled_pairs = random.sample(pairs, sample_size)\n",
    "\n",
    "    for i, pair in enumerate(sampled_pairs):\n",
    "        try:\n",
    "            # Get duration from the MIDI file\n",
    "            pair.load_midi()\n",
    "            duration = pair.midi.duration_s()\n",
    "            duration_stats[dataset_name].append(duration)\n",
    "        except Exception as e:\n",
    "            print(f\"  Sample {i + 1}: Error getting duration - {e}\")\n",
    "\n",
    "# Calculate and display summary statistics\n",
    "print(\"\\n\" + \"=\" * 50)\n",
    "print(\"DURATION SUMMARY STATISTICS\")\n",
    "print(\"=\" * 50)\n",
    "\n",
    "estimated_total_hours = {}\n",
    "\n",
    "for dataset_name, durations in duration_stats.items():\n",
    "    if durations:\n",
    "        avg_duration = sum(durations) / len(durations)\n",
    "        min_duration = min(durations)\n",
    "        max_duration = max(durations)\n",
    "\n",
    "        print(f\"\\n{dataset_name.upper()}:\")\n",
    "        print(\n",
    "            f\"  Average duration: {avg_duration:.2f} seconds ({avg_duration / 60:.2f} minutes)\"\n",
    "        )\n",
    "        # print(f\"  Min duration: {min_duration:.2f} seconds ({min_duration/60:.2f} minutes)\")\n",
    "        # print(f\"  Max duration: {max_duration:.2f} seconds ({max_duration/60:.2f} minutes)\")\n",
    "\n",
    "        # Estimate total dataset duration\n",
    "        total_pairs = len(all_pairs[dataset_name])\n",
    "        total_hours = (avg_duration * total_pairs) / 3600\n",
    "        estimated_total_hours[dataset_name] = total_hours\n",
    "        print(f\"  Estimated total dataset duration: {total_hours:.2f} hours\")\n",
    "\n",
    "# Plot total durations\n",
    "plt.figure(figsize=(10, 6))\n",
    "dataset_names = list(estimated_total_hours.keys())\n",
    "total_hours = list(estimated_total_hours.values())\n",
    "\n",
    "plt.bar(dataset_names, total_hours)\n",
    "plt.title(\"Estimated Total Dataset Durations\")\n",
    "plt.xlabel(\"Dataset\")\n",
    "plt.ylabel(\"Total Duration (Hours)\")\n",
    "plt.xticks(rotation=45)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# write a meta jsonl file with the pairs\n",
    "\n",
    "import json\n",
    "from tqdm import tqdm\n",
    "from concurrent.futures import ProcessPoolExecutor, as_completed\n",
    "import multiprocessing as mp\n",
    "\n",
    "\n",
    "def process_pair(pair_info):\n",
    "    dataset, pair = pair_info\n",
    "    try:\n",
    "        pair.load_midi()\n",
    "        midi = pair.midi\n",
    "        meta = {\n",
    "            \"midi_file\": pair.midi_file,\n",
    "            \"audio_file\": pair.audio_file,\n",
    "            \"dataset\": dataset,\n",
    "            \"duration_s\": midi.duration_s(),\n",
    "            \"num_notes\": midi.num_notes,\n",
    "            \"num_instruments\": midi.num_instruments,\n",
    "        }\n",
    "        return json.dumps(meta)\n",
    "    except Exception as e:\n",
    "        print(f\"Error processing pair {pair.midi_file}: {e}\")\n",
    "        return None\n",
    "\n",
    "\n",
    "import random\n",
    "\n",
    "train_meta_file = \"/app2/suno/data/victor/midi_pairs_meta_train.jsonl\"\n",
    "val_meta_file = \"/app2/suno/data/victor/midi_pairs_meta_val.jsonl\"\n",
    "\n",
    "# Prepare all pairs for processing\n",
    "all_pair_infos = []\n",
    "for dataset, pairs in all_pairs.items():\n",
    "    for pair in pairs:\n",
    "        all_pair_infos.append((dataset, pair))\n",
    "\n",
    "# Shuffle and split into train/val (90/10 split)\n",
    "random.shuffle(all_pair_infos)\n",
    "split_idx = int(0.98 * len(all_pair_infos))\n",
    "train_pair_infos = all_pair_infos[:split_idx]\n",
    "val_pair_infos = all_pair_infos[split_idx:]\n",
    "\n",
    "print(f\"Train pairs: {len(train_pair_infos)}\")\n",
    "print(f\"Val pairs: {len(val_pair_infos)}\")\n",
    "\n",
    "# Process pairs in parallel\n",
    "num_workers = min(\n",
    "    mp.cpu_count(), 64\n",
    ")  # Limit to 64 workers to avoid overwhelming the system\n",
    "\n",
    "# Process train set\n",
    "with open(train_meta_file, \"w\") as f:\n",
    "    with ProcessPoolExecutor(max_workers=num_workers) as executor:\n",
    "        # Submit all tasks\n",
    "        future_to_pair = {\n",
    "            executor.submit(process_pair, pair_info): pair_info\n",
    "            for pair_info in train_pair_infos\n",
    "        }\n",
    "\n",
    "        # Process completed tasks\n",
    "        for future in tqdm(\n",
    "            as_completed(future_to_pair),\n",
    "            total=len(train_pair_infos),\n",
    "            desc=\"Writing train meta\",\n",
    "        ):\n",
    "            result = future.result()\n",
    "            if result is not None:\n",
    "                f.write(result + \"\\n\")\n",
    "\n",
    "# Process val set\n",
    "with open(val_meta_file, \"w\") as f:\n",
    "    with ProcessPoolExecutor(max_workers=num_workers) as executor:\n",
    "        # Submit all tasks\n",
    "        future_to_pair = {\n",
    "            executor.submit(process_pair, pair_info): pair_info\n",
    "            for pair_info in val_pair_infos\n",
    "        }\n",
    "\n",
    "        # Process completed tasks\n",
    "        for future in tqdm(\n",
    "            as_completed(future_to_pair),\n",
    "            total=len(val_pair_infos),\n",
    "            desc=\"Writing val meta\",\n",
    "        ):\n",
    "            result = future.result()\n",
    "            if result is not None:\n",
    "                f.write(result + \"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!head /app2/suno/data/victor/midi_pairs_meta.jsonl"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Extract stems"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"7\"\n",
    "\n",
    "procid = int(os.environ.get(\"SLURM_PROCID\", 0))\n",
    "localid = int(os.environ.get(\"SLURM_LOCALID\", 0))\n",
    "world_size = int(os.environ.get(\"SLURM_JOB_NUM_NODES\", 1)) * int(\n",
    "    os.environ.get(\"SLURM_NTASKS_PER_NODE\", 1)\n",
    ")\n",
    "\n",
    "assert world_size > 0, \"WORLD_SIZE is 0\"\n",
    "\n",
    "print(f\"PROCID: {procid}, LOCALID: {localid}, WORLD_SIZE: {world_size}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.diffusion import generation as diffusion_gen\n",
    "from suno_utils.tasks.upsample_engine import UpsampleEngine, Request\n",
    "from suno_utils.tasks.dac_vae_fixed_25hz import (\n",
    "    decode_stream_to_full_audio,\n",
    "    encode,\n",
    "    decode,\n",
    ")\n",
    "\n",
    "import torch\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "from suno_utils.audio import Audio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out_stem_dir = \"/app2/suno/data/victor/midi_stems\"\n",
    "os.makedirs(out_stem_dir, exist_ok=True)\n",
    "\n",
    "pair_audio_paths = []\n",
    "pairs_to_process = {\n",
    "    \"scoretube\": scoretube_pairs,\n",
    "    \"trombone_champ\": trombone_champ_pairs,\n",
    "    \"hooktheory\": hooktheory_pairs,\n",
    "    \"mootube\": mootube_pairs,\n",
    "    \"mmd\": mmd_pairs,\n",
    "}  # ignore maestro\n",
    "\n",
    "for dataset, pairs in pairs_to_process.items():\n",
    "    for pair in pairs:\n",
    "        name = pair.audio_file.split(\"/\")[-1].split(\".\")[0]\n",
    "        pair_audio_paths.append((pair.audio_file, f\"{out_stem_dir}/{dataset}/{name}\"))\n",
    "print(len(pair_audio_paths))\n",
    "\n",
    "pair_audio_paths = pair_audio_paths[procid::world_size]\n",
    "print(len(pair_audio_paths))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "# diffusion_gen.preload_models(\n",
    "#     dit_model_filepath=\"/app/suno/checkpoints/2025-05-26_22-49-34_s8646/last_ckpt_infer.pt\",  # 12 stems\n",
    "#     codec_filepath=\"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\",\n",
    "#     compile=True,\n",
    "# )\n",
    "\n",
    "# engine = UpsampleEngine(min_chunk_size=25 * 30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gen_stem(\n",
    "    audio: Audio,\n",
    "    stem_type_cfg_scale=1.0,\n",
    "    tags=\"extract [split_karaoke]\",\n",
    "    steps=4,\n",
    "    seed=3,\n",
    "    codec_scale_factor=0.4,\n",
    "    scale_ctx_vector=True,\n",
    "    noise_ctx_level=0.0,\n",
    "    infill_prefix_latents=None,\n",
    "    infill_suffix_latents=None,\n",
    "):\n",
    "    vae = encode(audio)\n",
    "    gen_cfg = diffusion_gen.DiffusionGenerationConfig(\n",
    "        lyrics=tags,\n",
    "        steps=steps,\n",
    "        seed=seed,\n",
    "        codec_scale_factor=codec_scale_factor,\n",
    "        scale_ctx_vector=scale_ctx_vector,\n",
    "        noise_ctx_level=noise_ctx_level,\n",
    "        text_cfg_coef=stem_type_cfg_scale,\n",
    "        infill_prefix_latents=infill_prefix_latents,\n",
    "        infill_suffix_latents=infill_suffix_latents,\n",
    "        drop_semantic_tokens=True,\n",
    "    )\n",
    "    # print(gen_cfg)\n",
    "\n",
    "    request = Request(\n",
    "        id=\"dummy\",\n",
    "        generation_config=gen_cfg,\n",
    "        tokens=np.zeros((vae.shape[0], 1)),\n",
    "        input_tokens_finished=True,\n",
    "        stem_ctx_latents=vae,\n",
    "    )\n",
    "\n",
    "    result = engine.run_request(request, tqdm_enabled=False)\n",
    "    vae_latents = torch.concat(result.vae_latents)\n",
    "    # print(f\"vae_latents: {vae_latents.shape}\")\n",
    "    audios = []\n",
    "    for i in tqdm(range(vae_latents.shape[1]), desc=\"Decoding stems\"):\n",
    "        # audios.append(decode(vae_latents[:, i]))\n",
    "        audios.append(\n",
    "            decode_stream_to_full_audio(vae_latents[:, i], n_stride_tokens=25 * 10)\n",
    "        )\n",
    "    return audios"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_stem(args):\n",
    "    out_root, category, stem = args\n",
    "    if stem.loudness < -45:\n",
    "        return None\n",
    "    out_path = f\"{out_root}_{category}.opus\"\n",
    "    os.makedirs(os.path.dirname(out_path), exist_ok=True)\n",
    "    # print(f\"Writing {out_path}\")\n",
    "    stem.write_opus(out_path)\n",
    "    return category\n",
    "\n",
    "\n",
    "import threading\n",
    "\n",
    "categories = [\n",
    "    \"Vocals\",\n",
    "    \"Backing_Vocals\",\n",
    "    \"Drums\",\n",
    "    \"Bass\",\n",
    "    \"Guitar\",\n",
    "    \"Keyboard\",\n",
    "    \"Percussion\",\n",
    "    \"Strings\",\n",
    "    \"Synth\",\n",
    "    \"FX\",\n",
    "    \"Brass\",\n",
    "    \"Woodwinds\",\n",
    "]\n",
    "\n",
    "\n",
    "def process_id(pair):\n",
    "    input_path, output_path = pair\n",
    "    audio = Audio.from_file(input_path, n_channels=2)\n",
    "    stems = gen_stem(audio, steps=8)\n",
    "\n",
    "    # Prepare arguments for parallel processing\n",
    "    write_args = [\n",
    "        (output_path, category, stem) for category, stem in zip(categories, stems)\n",
    "    ]\n",
    "\n",
    "    # Use threading to write stems in parallel\n",
    "    results = [None] * len(write_args)\n",
    "    threads = []\n",
    "\n",
    "    def write_stem_thread(i, args):\n",
    "        results[i] = write_stem(args)\n",
    "\n",
    "    for i, args in enumerate(write_args):\n",
    "        thread = threading.Thread(target=write_stem_thread, args=(i, args))\n",
    "        threads.append(thread)\n",
    "        thread.start()\n",
    "\n",
    "    for thread in threads:\n",
    "        thread.join()\n",
    "\n",
    "    # Filter out None results\n",
    "    found_categories = [result for result in results if result is not None]\n",
    "    return found_categories\n",
    "\n",
    "\n",
    "# process_id(random.choice(pair_audio_paths))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for pair in tqdm(pair_audio_paths, mininterval=600):\n",
    "#     try:\n",
    "#         process_id(pair)\n",
    "#     except Exception as e:\n",
    "#         print(f\"Error processing {pair}: {e}\")\n",
    "#         continue"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Check stems"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !du -sh /app2/suno/data/victor/midi_stems/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# count number of stems in each dataset\n",
    "\n",
    "for dataset in pairs_to_process.keys():\n",
    "    # count number of stems in each dataset in the dir\n",
    "    dir = f\"{out_stem_dir}/{dataset}\"\n",
    "    # count number of files in the dir\n",
    "    num_files = len(os.listdir(dir))\n",
    "    print(f\"{dataset}: {num_files} stemmed\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_clean",
   "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.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}