{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.tasks.mert_25 import (\n",
    "    encode as mert_encode,\n",
    "    SAMPLE_RATE,\n",
    "    EMBEDDING_RATE,\n",
    "    decode,\n",
    ")\n",
    "from suno_utils.diffusion_stems import generation as diffusion_gen\n",
    "from suno_utils.audio import Audio\n",
    "from suno_utils.tasks.mert_25 import preload_models\n",
    "from suno_utils.tasks.data_loader import load_audio_opus\n",
    "\n",
    "_ = preload_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",
    ")\n",
    "\n",
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"3\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio_arr = load_audio_opus(\n",
    "    \"/home/sara/glockenspiel/suno_utils/suno_utils/scripts/gpt/NrkrK3W1tEA.opus\",\n",
    "    n_channels=2,\n",
    "    max_duration_s=120,\n",
    "    target_sample_rate=48_000,\n",
    "    normalize_volume=True,\n",
    "    debug=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test = Audio.from_array_float(audio_arr, sample_rate=48_000, auto_compress=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "encoded = mert_encode([test])\n",
    "decoded = mert_decode([encoded])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = diffusion_gen.preload_models(\n",
    "    tokenizer_filepath=\"s3://suno-data/georg/models/tokenizers/tokenizer_60k.json\",\n",
    "    codec_filepath=\"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\",\n",
    "    dit_model_filepath=\"s3://suno-data/sara/models/stems_v1p2.pt\",\n",
    "    compile=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio = Audio.from_file(\"/home/sara/samples/welcome_to_nowhere.wav\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arrays = [audio]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio.array_float.reshape(1, -1).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for audio_tensor in arrays:\n",
    "    gen_cfg = diffusion_gen.DiffusionGenerationConfig(\n",
    "        audio=Audio.from_array_float(\n",
    "            audio.array_float.reshape(1, -1), sample_rate=audio.sample_rate\n",
    "        ),\n",
    "        stem_type_id=0,\n",
    "        steps=10,\n",
    "    )\n",
    "\n",
    "    pred_audio, l = diffusion_gen.generate(gen_cfg)\n",
    "    pred_audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mert_encode([pred_audio])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.utils.text import read_jsonl, read_json, write_jsonl, write_json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test = read_jsonl(\"metas.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_json = read_json(\"/app/suno/data/chirp_v5/v1/info_val.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test[1663:1667]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(test_json[\"musdb_stems_underpaint\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for t in test:\n",
    "    if \"paint\" in t[\"task\"]:\n",
    "        print(t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = read_jsonl(\"/home/sara/glockenspiel/musdb_stems/mert_25_2x4k/metas/part_0.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "b = np.load(\"/home/sara/glockenspiel/musdb_stems/mert_25_2x4k/part_0.npz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "b.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for val in a:\n",
    "    print(val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metas = read_jsonl(\"metas.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metas[0][\"original_s3_filepath\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "paint_metas = []\n",
    "for m in metas:\n",
    "\n",
    "    def make_row(stem_type, row):\n",
    "        name = row[\"id\"]\n",
    "        return {\n",
    "            \"id\": f\"{name}_{stem_type}\",\n",
    "            \"s3_filepath\": row[\"s3_filepath\"],\n",
    "            \"duration_s\": row[\"duration_s\"],\n",
    "            \"type\": stem_type,\n",
    "        }\n",
    "\n",
    "    full = make_row(\"full\", m)\n",
    "    vocals = make_row(\"vocals\", m)\n",
    "    instrumentals = make_row(\"instrumentals\", m)\n",
    "\n",
    "    paint_metas.append(full)\n",
    "    paint_metas.append(instrumentals)\n",
    "    paint_metas.append(vocals)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "write_jsonl(paint_metas, \"metas_paint.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "paint_metas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = read_jsonl(\n",
    "    \"/app/suno/sara/cover_filter/covers_filtered_0.3_0.8_all.jsonl\", max_lines=500\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x[0].keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test = read_jsonl(\"musdb_stems/metas_v0.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test = read_jsonl(\"METAS_TEST.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "labelbox_cleaned = read_json(\n",
    "    \"/home/sara/glockenspiel/suno_utils/task_eval/labelbox/labelbox_sources_cleaned.json\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = {}\n",
    "\n",
    "for key, prompts in labelbox_cleaned.items():\n",
    "    gendered_prompts = []\n",
    "    for source in prompts:\n",
    "        if len(source[\"lyrics\"]) > 0:\n",
    "            female = source.copy()\n",
    "            female[\"tags_no_gender\"] = female[\"tags\"]\n",
    "            female[\"tags\"] = \"female singer, \" + female[\"tags\"]\n",
    "            female[\"vocalist\"] = \"female\"\n",
    "\n",
    "            male = source.copy()\n",
    "            male[\"tags_no_gender\"] = male[\"tags\"]\n",
    "            male[\"tags\"] = \"male singer, \" + male[\"tags\"]\n",
    "            male[\"vocalist\"] = \"male\"\n",
    "\n",
    "            gendered_prompts.append(female)\n",
    "            gendered_prompts.append(male)\n",
    "\n",
    "    random.shuffle(gendered_prompts)\n",
    "    print(key, len(gendered_prompts), len(prompts))\n",
    "    if len(gendered_prompts) > 0:\n",
    "        output[key] = gendered_prompts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\n",
    "    \"/home/sara/glockenspiel/suno_utils/task_eval/labelbox/labelbox_sources_gendered.json\",\n",
    "    \"w\",\n",
    "    encoding=\"utf-8\",\n",
    ") as json_file:\n",
    "    json.dump(output, json_file, indent=4, ensure_ascii=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "write_json(\n",
    "    output,\n",
    "    \"/home/sara/glockenspiel/suno_utils/task_eval/labelbox/labelbox_sources_gendered.json\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.audio.conversion import Audio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio = Audio.from_file(\n",
    "    \"/home/sara/glockenspiel/suno_utils/suno_utils/scripts/gpt/NrkrK3W1tEA.opus\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_auk = read_json(\n",
    "    \"/home/sara/glockenspiel/suno_utils/task_eval/modal_runs/genre_mappings_chirp-v5-sem-6b_2025_04_01-15_24_24.json\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for key, prompts in test_auk.items():\n",
    "    gendered_prompts = []\n",
    "    for source in prompts:\n",
    "        if source[\"tags\"].startswith(\"male singer, \"):\n",
    "            source[\"vocalist\"] = \"male\"\n",
    "            source[\"tags_no_gender\"] = source[\"tags\"][13:]\n",
    "        else:\n",
    "            source[\"vocalist\"] = \"female\"\n",
    "            source[\"tags_no_gender\"] = source[\"tags\"][15:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_auk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "write_json(\n",
    "    test_auk,\n",
    "    \"/home/sara/glockenspiel/suno_utils/task_eval/modal_runs/genre_mappings_chirp-v5-sem-6b_2025_04_01-15_24_24.json\",\n",
    ")"
   ]
  },
  {
   "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
}