{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a9ac25c9-4432-48f8-b403-b9618acf6206",
   "metadata": {},
   "source": [
    "# Audio Generation Inference Notebook\n",
    "\n",
    "This notebook demonstrates music generation using diffusion models and GPT-based semantic models.\n",
    "\n",
    "## System Check"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68e3b4ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment to check GPU status\n",
    "# !nvidia-smi\n",
    "# !echo $HOSTNAME"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35a6eaef-d02a-48e3-82aa-dd7b0493a74c",
   "metadata": {},
   "source": [
    "## Environment Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd591ed4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Configure GPU\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"2\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a54602f2-ce29-4cbe-abbb-f104b254f797",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3606dd14",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Standard library imports\n",
    "import json\n",
    "from typing import Optional\n",
    "\n",
    "# Third-party imports\n",
    "import torch\n",
    "import numpy as np\n",
    "\n",
    "# Suno utilities - Audio and text\n",
    "from suno_utils.audio import Audio\n",
    "from suno_utils.utils.text import read_jsonl\n",
    "from suno_utils.utils.s3 import list_s3_dir, read_from_s3\n",
    "\n",
    "# Suno utilities - Diffusion\n",
    "from suno_utils.diffusion.generation import (\n",
    "    preload_dit_model,\n",
    "    preload_tokenizer,\n",
    "    TOKENIZER_FILEPATH,\n",
    "    SEMANTIC_MODEL_FILEPATH,\n",
    "    SEMANTIC_CLUSTERS_FILEPATH,\n",
    "    _retrieve_models,\n",
    ")\n",
    "\n",
    "# Suno utilities - Semantic encoding\n",
    "from suno_utils.tasks.mert_25 import (\n",
    "    preload_models as preload_semantic_models,\n",
    "    encode as encode_semantic,\n",
    ")\n",
    "\n",
    "# Suno utilities - Codec\n",
    "from suno_utils.tasks.dac_vae_fixed_25hz import (\n",
    "    preload_models as preload_codec_models,\n",
    "    decode as codec_decode,\n",
    "    decode_stream_to_full_audio,\n",
    ")\n",
    "\n",
    "# Suno utilities - Generation\n",
    "from suno_utils.diffusion import generation as diffusion_gen\n",
    "from suno_utils.tasks.upsample_engine import UpsampleEngine, Request, Job\n",
    "from suno_utils.gpt.generation import GenerationConfig\n",
    "from suno_utils.gpt.engine import Engine\n",
    "from suno_utils.gpt.generation_engine import make_request"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e04b2bc-f037-4fae-86a8-6c8a10b44a94",
   "metadata": {},
   "source": [
    "## Configuration Constants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8e37fa5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model paths\n",
    "DIT_MODEL_FILEPATH = \"/app2/suno/modal/models/tony/tmp/diff/v45_2b_step_2mil_ft_8k_infill_apr21_d4_v39.pt\"\n",
    "GPT_MODEL_FILEPATH = \"/app2/suno/checkpoints/2025-11-20_02-37-25/last_ckpt_infer.pt\" # dodo auk sft0 # d70 2025-10-27_00-05-48\n",
    "TOKENIZER_FILEPATH = \"s3://suno-data/georg/models/tokenizers/tokenizer_60k.json\"\n",
    "SEMANTIC_MODEL_FILEPATH = \"s3://suno-data/georg/models/semantic/mert_25.pt\"\n",
    "SEMANTIC_CLUSTERS_FILEPATH = \"s3://suno-data/georg/models/semantic/mert_25_2x4k.npy\"\n",
    "CODEC_FILEPATH = \"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\"\n",
    "\n",
    "# Generation parameters\n",
    "N_BATCH = 4\n",
    "CODEC_SCALE_FACTOR = 0.4\n",
    "SCALE_CTX_VECTOR = True\n",
    "MIN_CHUNK_SIZE = 25 * 30  # 30 seconds at 25Hz"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef9bfab8-b567-4fb4-91aa-6bf6eb53636d",
   "metadata": {},
   "source": [
    "## Load Diffusion Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04547845",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check GPU availability\n",
    "num_gpus = torch.cuda.device_count()\n",
    "cuda_device = torch.cuda.current_device()\n",
    "print(f\"Found {num_gpus} GPUs. Using GPU {cuda_device}.\")\n",
    "\n",
    "# Load diffusion model\n",
    "_ = diffusion_gen.preload_dit_model(\n",
    "    dit_model_filepath=DIT_MODEL_FILEPATH,\n",
    "    use_ema_if_exists=True,\n",
    "    compile=False,\n",
    "    weights_precision=torch.bfloat16,\n",
    ")\n",
    "\n",
    "# Load supporting models\n",
    "_ = preload_tokenizer(TOKENIZER_FILEPATH)\n",
    "_ = preload_semantic_models(SEMANTIC_MODEL_FILEPATH, SEMANTIC_CLUSTERS_FILEPATH)\n",
    "_ = preload_codec_models(CODEC_FILEPATH)\n",
    "\n",
    "# Initialize diffusion engine\n",
    "diffusion_engine = UpsampleEngine(min_chunk_size=MIN_CHUNK_SIZE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8696cde6-9ea2-453b-8c95-7693cd02c205",
   "metadata": {},
   "source": [
    "## Load GPT Engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89085dd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize GPT engine\n",
    "engine = Engine(\n",
    "    GPT_MODEL_FILEPATH,\n",
    "    \"/app/suno/models/chirp_v2/tokenizer_60k.json\",\n",
    "    max_sequences=4 * N_BATCH,\n",
    "    compile=False,\n",
    ")\n",
    "cfg = engine.model.config\n",
    "\n",
    "print(f\"GPT engine loaded with config: {cfg}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bd0ad3c-7219-4315-9796-afef491f5b43",
   "metadata": {},
   "source": [
    "## Generation Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9eaa269c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_clips(\n",
    "    text: str,\n",
    "    tags: str,\n",
    "    neg_tags: Optional[str] = None,\n",
    "    cfg_coef_tags_max_steps: Optional[int] = None,\n",
    "    control_tags: Optional[str] = None,\n",
    "    n_skip_semantic: int = 1,\n",
    "    steps: int = 10,\n",
    "    text_cfg_coef: float = 2.0,\n",
    "    noise_ctx_level: float = 0.75,\n",
    ") -> None:\n",
    "    \"\"\"Generate audio clips from text and tags using GPT and diffusion models.\n",
    "\n",
    "    Args:\n",
    "        text: Lyrics or descriptive text for generation\n",
    "        tags: Genre/style tags for the audio\n",
    "        neg_tags: Negative tags to avoid (default: \"repetitive, loop\")\n",
    "        cfg_coef_tags_max_steps: Max steps for tags CFG (default: 25 * 120)\n",
    "        control_tags: Control tags for generation\n",
    "        n_skip_semantic: Semantic token skip rate (default: 1)\n",
    "        steps: Number of diffusion steps (default: 10)\n",
    "        text_cfg_coef: Text classifier-free guidance coefficient (default: 2.0)\n",
    "        noise_ctx_level: Context noise level (default: 0.75)\n",
    "    \"\"\"\n",
    "    # Configure GPT generation\n",
    "    gconf = GenerationConfig(\n",
    "        text=text,\n",
    "        text_tags=tags,\n",
    "        cfg_coef=1.0,\n",
    "        cfg_coef_tags=1.0,\n",
    "        cfg_coef_tags_max_steps=25 * 120\n",
    "        if cfg_coef_tags_max_steps is None\n",
    "        else cfg_coef_tags_max_steps,\n",
    "        n_repeat_tags=1,\n",
    "        n_skip_semantic=n_skip_semantic,\n",
    "        text_start_control_tags=f\"{control_tags}\" if control_tags is not None else None,\n",
    "        cfg_coef_neg_tags=-1,\n",
    "        text_neg_tags=\"repetitive, loop\" if neg_tags is None else neg_tags,\n",
    "        temp_semantic=0.90,\n",
    "        min_p_semantic=0.005,\n",
    "        n_batch=1,\n",
    "        min_eos_p=0.1,\n",
    "        min_text_offset=0,\n",
    "        eos_pad_duration_s=0,\n",
    "        max_gen_duration_s=int(4 * 60 / n_skip_semantic),\n",
    "        random_seed=0,\n",
    "    )\n",
    "\n",
    "    # Generate semantic tokens with GPT\n",
    "    requests = [\n",
    "        make_request(f\"{i}\", gconf, engine.model.config, engine.tokenizer)\n",
    "        for i in range(N_BATCH)\n",
    "    ]\n",
    "    jobs = engine.run_request(requests, tqdm_enabled=True)\n",
    "\n",
    "    out_gpt = []\n",
    "    for n, job in enumerate(jobs):\n",
    "        stream = engine.token_generator(job)\n",
    "        arr = torch.stack(list(stream))[:, 1]\n",
    "        if arr[-1] == 4000:\n",
    "            arr = arr[:-1]\n",
    "        print(f\"{round(arr.shape[-1]/25*n_skip_semantic)}s for track {n}\")\n",
    "\n",
    "        # Handle semantic skip\n",
    "        arr2 = (\n",
    "            torch.zeros(arr.shape[0] * n_skip_semantic, dtype=arr.dtype)\n",
    "            + cfg.semantic_pad_token\n",
    "        )\n",
    "        arr2[::n_skip_semantic] = arr\n",
    "        out_gpt.append(arr2)\n",
    "\n",
    "    # Generate audio from semantic tokens\n",
    "    for in_sem_arr in out_gpt:\n",
    "        gen_cfg = diffusion_gen.DiffusionGenerationConfig(\n",
    "            steps=steps,\n",
    "            lyrics=text,\n",
    "            tags=tags,\n",
    "            text_cfg_coef=text_cfg_coef,\n",
    "            ctx_cfg_coef=1.0,\n",
    "            codec_scale_factor=CODEC_SCALE_FACTOR,\n",
    "            scale_ctx_vector=SCALE_CTX_VECTOR,\n",
    "            noise_ctx_level=noise_ctx_level,\n",
    "            noise_ctx_pad_len=0,\n",
    "            drop_semantic_tokens=False,\n",
    "            seed=42,\n",
    "            rho=1.0,\n",
    "            sigma_min=0.5,\n",
    "            sigma_max=50.0,\n",
    "            semantic_mask_ratio=0.0,\n",
    "        )\n",
    "\n",
    "        request = Request(\n",
    "            id=\"dummy\",\n",
    "            generation_config=gen_cfg,\n",
    "            tokens=in_sem_arr,\n",
    "            input_tokens_finished=True,\n",
    "        )\n",
    "\n",
    "        result = diffusion_engine.run_request(request)\n",
    "\n",
    "        # Decode and play audio\n",
    "        vae_latents = torch.concat([vae_latent for vae_latent in result.vae_latents])\n",
    "        upsampled_audio = decode_stream_to_full_audio(vae_latents)\n",
    "        upsampled_audio.play()\n",
    "\n",
    "\n",
    "def generate_from_semantic(\n",
    "    in_sem_arr: torch.Tensor,\n",
    "    lyrics: str,\n",
    "    tags: str,\n",
    "    text_cfg_coef: float = 1.0,\n",
    "    steps: int = 16,\n",
    "    seed: int = 0,\n",
    ") -> Audio:\n",
    "    \"\"\"Generate audio from semantic tokens.\n",
    "\n",
    "    Args:\n",
    "        in_sem_arr: Input semantic token array\n",
    "        lyrics: Lyrics text\n",
    "        tags: Genre/style tags\n",
    "        text_cfg_coef: Text classifier-free guidance coefficient\n",
    "        steps: Number of diffusion steps\n",
    "        seed: Random seed\n",
    "\n",
    "    Returns:\n",
    "        Generated Audio object\n",
    "    \"\"\"\n",
    "    gen_cfg = diffusion_gen.DiffusionGenerationConfig(\n",
    "        steps=steps,\n",
    "        lyrics=lyrics,\n",
    "        tags=tags,\n",
    "        text_cfg_coef=text_cfg_coef,\n",
    "        ctx_cfg_coef=1.0,\n",
    "        codec_scale_factor=CODEC_SCALE_FACTOR,\n",
    "        scale_ctx_vector=SCALE_CTX_VECTOR,\n",
    "        noise_ctx_level=0.75,\n",
    "        noise_ctx_pad_len=0,\n",
    "        drop_semantic_tokens=False,\n",
    "        seed=seed,\n",
    "        rho=1.0,\n",
    "        sigma_min=0.5,\n",
    "        sigma_max=50.0,\n",
    "        semantic_mask_ratio=0.0,\n",
    "    )\n",
    "\n",
    "    request = Request(\n",
    "        id=\"dummy\",\n",
    "        generation_config=gen_cfg,\n",
    "        tokens=in_sem_arr,\n",
    "        input_tokens_finished=True,\n",
    "    )\n",
    "\n",
    "    result = diffusion_engine.run_request(request)\n",
    "    vae_latents = torch.concat([vae_latent for vae_latent in result.vae_latents])\n",
    "    return decode_stream_to_full_audio(vae_latents)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f66b5d7-8f36-410f-9c1a-ab78db47b45f",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "# Basic Text-to-Music Generation Examples\n",
    "\n",
    "## Example 1: Pop Song Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a751bcd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[intro]\n",
    "\n",
    "[verse]\n",
    "Walking down the street, feeling so alive\n",
    "Got my head in the clouds, got a gleam in my eye\n",
    "Every step I take, it's like a brand new start\n",
    "No matter where I'm going, I'll always find my part\n",
    "(oh-oh-oh)\n",
    "\n",
    "[chorus]\n",
    "Life is like a high-wire act, we're dancing in the sky\n",
    "No need to worry, no need to ask why\n",
    "With a little bit of courage, we can chase our dreams\n",
    "No matter what comes our way, we'll always be a team\n",
    "(we're unstoppable, yeah)\n",
    "\n",
    "[outro]\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3133ec0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_clips(text, tags=\"epic film orchestral\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9567524d-3238-4e03-b65d-031e68566ef1",
   "metadata": {},
   "source": [
    "## Example 2: Gregorian Chant Tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2517468-63bb-451a-9cad-12fabe5a8902",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[verse]\n",
    "A is for the amazing grace that we recieve\n",
    "B is for the blessings, every day we believe\n",
    "C is for the chorus, we sing it loud and clear\n",
    "D is for the devotion that we hold dear\n",
    "\n",
    "[verse]\n",
    "E is for the everlasting love that sets us free\n",
    "F is for the faith that guides us on this journey\n",
    "G is for the goodness that we share each day\n",
    "H is for the hope that never fades away\n",
    "\n",
    "[Chorus]\n",
    "B is for Buttocks, ripe and slightly damp\n",
    "Yum Yum, boy oh boy do I like God\n",
    "Every day is a gift, when you won a skateboard ramp.\n",
    "Yum yum, boy oh boy do I like God.\n",
    "\n",
    "[outro]\n",
    "\"\"\"\n",
    "\n",
    "# Basic gregorian chant\n",
    "generate_clips(text, tags=\"gregorian chant\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c044be24-be21-4810-82b1-bd30d05ffc31",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # With control tag strength 1\n",
    "# generate_clips(text, tags=\"gregorian chant\", control_tags=\"tag_strength:1;\")\n",
    "# # With control tag strength 9\n",
    "# generate_clips(text, tags=\"gregorian chant\", control_tags=\"tag_strength:9\")\n",
    "# # With negative tags for instrumental\n",
    "# generate_clips(\n",
    "#     text, tags=\"gregorian chant, no instrumental\", control_tags=\"tag_strength:9\"\n",
    "# )\n",
    "# # With repeated tags for emphasis\n",
    "# generate_clips(text, tags=\"gregorian chant, gregorian chant, gregorian chant\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79410531-c96a-418d-b79e-7fa55ea8b679",
   "metadata": {},
   "source": [
    "## Example 3: Other Styles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "141bcefc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Bagpipes test\n",
    "text_bagpipes = \"\"\"\n",
    "[bagpipes, verse]\n",
    "A is for the amazing grace that we recieve\n",
    "B is for the blessings, every day we believe\n",
    "C is for the chorus, we sing it loud and clear\n",
    "D is for the devotion that we hold dear\n",
    "\n",
    "[verse]\n",
    "E is for the everlasting love that sets us free\n",
    "F is for the faith that guides us on this journey\n",
    "G is for the goodness that we share each day\n",
    "H is for the hope that never fades away\n",
    "\n",
    "[Chorus]\n",
    "B is for Buttocks, ripe and slightly damp\n",
    "Yum Yum, boy oh boy do I like God\n",
    "Every day is a gift, when you won a skateboard ramp.\n",
    "Yum yum, boy oh boy do I like God.\n",
    "\n",
    "[outro]\n",
    "\"\"\"\n",
    "generate_clips(text_bagpipes, tags=\"bagpipes, scottish funeral march\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f64ec057",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Chinese pop test\n",
    "text_chinese = \"\"\"\n",
    "[Verse]\n",
    "一盏离愁孤灯伫立在窗口\n",
    "我在门后假装你人还没走\n",
    "旧地如重游月圆更寂寞\n",
    "夜半清醒的烛火不忍苛责我\n",
    "\n",
    "[Verse]\n",
    "一壶漂泊浪迹天涯难入喉\n",
    "你走之后酒暖回忆思念瘦\n",
    "水向东流时间怎么偷\n",
    "花开就一次成熟我却错过\n",
    "\n",
    "[Chorus]\n",
    "谁在用琵琶弹奏一曲东风破\n",
    "岁月在墙上剥落看见小时候\n",
    "犹记得那年我们都还很年幼\n",
    "而如今琴声幽幽我的等候你没听过\n",
    "\n",
    "[outro]\n",
    "\"\"\"\n",
    "generate_clips(text_chinese, tags=\"chinese pop woman\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd52dc41-1430-446e-847e-b1532d2f7b2d",
   "metadata": {},
   "source": [
    "## Example 4: Classical/Orchestral"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01a98504",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Late 19th century Austrian romantic symphony\n",
    "generate_clips(\n",
    "    \"\", tags=\"late 19th century austrian romantic symphony, dramatic, large orchestra\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "741261a7-70bb-46a1-94f5-2547c17e6484",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Beethoven-style symphony\n",
    "beethoven_desc = \"[Dramatic slow opening, drums, strong strings, emotional, intense] [Woodwinds comes in with constrast, delicate, sad, sorrow] [Back strings, fast tempo, echo the theme]\"\n",
    "beethoven_tags = \"beethoven, beethoven, beethoven, ludwig van beethoven, ludwig van beethoven as the composer, dg, philharmonic orchestra, DG, symphony 1st movement\"\n",
    "beethoven_neg = \"vocals, vocals, pianos, classical, baroque, pop, mozart, opera, mozart, haydn, brahms, chopin, overture, overture, march, march\"\n",
    "generate_clips(\n",
    "    beethoven_desc,\n",
    "    tags=beethoven_tags,\n",
    "    neg_tags=beethoven_neg,\n",
    "    cfg_coef_tags_max_steps=25 * 4 * 60,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c278ed9-e9df-47f3-ad5e-be9f225d4ace",
   "metadata": {},
   "outputs": [],
   "source": [
    "# JS Bach fugue\n",
    "generate_clips(\"\", tags=\"JS bach, JSbach, fugue, piano, baroque\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f092408-a491-4f79-a701-a4b12973bb9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# JS Bach fugue\n",
    "generate_clips(\"\", tags=\"bach, bach, fugue, piano, baroque\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16e662b1-6515-4e5c-90b0-4d429600edfc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Baroque fugue (slow)\n",
    "generate_clips(\"\", tags=\"fugue, piano, baroque, slow tempo\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0746787b-311a-4b25-b7a3-8ab816a97b63",
   "metadata": {},
   "source": [
    "## Example 5: Country/Bluegrass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37874fb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "text_country = \"\"\"\n",
    "[Verse 1]\n",
    "Almost Heaven, West Virginia\n",
    "Blue Ridge Mountains, Shenandoah River\n",
    "Life is old there, older than the trees\n",
    "Younger than the mountains, growing like a breeze\n",
    "\n",
    "[Chorus]\n",
    "Country roads, take me home\n",
    "To the place I belong\n",
    "West Virginia, mountain mama\n",
    "Take me home, country roads\n",
    "\n",
    "[end][end][end]\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc47d650-4bec-4a91-b1ba-def5716e9f27",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_clips(text_country, tags=\"female bluegrass, fast tempo\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f85a4117-9fa3-4b08-afe7-dd69c71ab8e3",
   "metadata": {},
   "source": [
    "## Experimental Tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb7dda71-0320-458f-b58a-f9cd90eeb4cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Edge case test with unusual lyrics\n",
    "generate_clips(\"plastic 1, plastic 2, plastic 3, plastic 4, plastic 5\", tags=\"pop\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbe4f87d",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "# Advanced: Over/Underpaint Techniques\n",
    "\n",
    "These techniques allow you to generate music conditioned on existing audio (vocals or instrumentals)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6566827e",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Additional Imports for Over/Underpaint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48768c9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "from suno_utils.tasks.demucs import split_vocals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aaf94996",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Load Validation Dataset (Optional)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f98dcf7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Sample Lyrics for Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d1b64e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sample lyrics for underpaint/overpaint/cover experiments\n",
    "sample_lyrics = \"\"\"\n",
    "[Verse 1]\n",
    "Remember those walls I built?\n",
    "Well, baby, they're tumblin' down\n",
    "And they didn't even put up a fight\n",
    "They didn't even make a sound\n",
    "\n",
    "[Pre-Chorus]\n",
    "It's like I've been awakened\n",
    "Every rule, I had you breakin'\n",
    "\n",
    "[Chorus]\n",
    "Everywhere I'm lookin' now\n",
    "I'm surrounded by your embrace\n",
    "You're everything I need and more\n",
    "It's written all over your face\n",
    "\n",
    "[Verse 2]\n",
    "Hit me like a ray of sun\n",
    "Burnin' through my darkest night\n",
    "You're the only one that I want\n",
    "Think I'm addicted to your light\n",
    "\n",
    "[outro]\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa3cd8c8",
   "metadata": {},
   "source": [
    "## Underpaint: Generate Instrumental from Vocals\n",
    "\n",
    "Generate instrumental music conditioned on existing vocal tracks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98c2d4e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load and split audio file\n",
    "a = Audio.from_file(\"../audios/halo.wav\", sample_rate=44_100, n_channels=2)\n",
    "a_vocals, a_other = split_vocals(a.convert(44_100, 2, 2))\n",
    "\n",
    "# Encode vocals to semantic tokens\n",
    "vocals_arr = encode_semantic(a_vocals.convert(44_100, 2, 2).normalize_volume())[:, :1]\n",
    "\n",
    "print(f\"Vocals semantic shape: {vocals_arr.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1023d4b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure generation\n",
    "tags = \"Pop\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fa2fa71",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create generation config with underpaint\n",
    "gconf = GenerationConfig(\n",
    "    text=sample_lyrics,\n",
    "    text_tags=tags,\n",
    "    underpaint_arr=vocals_arr,\n",
    "    cfg_coef=1.0,\n",
    "    cfg_coef_tags=1.0,\n",
    "    n_repeat_tags=1,\n",
    "    n_batch=1,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    "    max_gen_duration_s=2 * 60,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "020e0ac5",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Generate semantic tokens with underpaint conditioning\n",
    "requests = [\n",
    "    make_request(f\"{i}\", gconf, engine.model.config, engine.tokenizer)\n",
    "    for i in range(N_BATCH)\n",
    "]\n",
    "jobs = engine.run_request(requests, tqdm_enabled=True)\n",
    "\n",
    "out_gpt = []\n",
    "for n, job in enumerate(jobs):\n",
    "    stream = engine.token_generator(job)\n",
    "    arr = torch.stack(list(stream))[:, 1]\n",
    "    if arr[-1] == 4000:\n",
    "        arr = arr[:-1]\n",
    "    print(f\"{round(arr.shape[-1]/25)}s for track {n}\")\n",
    "    out_gpt.append(arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf653ddf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate and play audio from semantic tokens\n",
    "for in_sem_arr in out_gpt:\n",
    "    audio = generate_from_semantic(\n",
    "        in_sem_arr,\n",
    "        lyrics=sample_lyrics,\n",
    "        tags=tags,\n",
    "        text_cfg_coef=1.0,\n",
    "        steps=16,\n",
    "        seed=0,\n",
    "    )\n",
    "    audio.play()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f687e4c",
   "metadata": {},
   "source": [
    "## Overpaint: Generate Vocals over Instrumental\n",
    "\n",
    "Generate vocals conditioned on existing instrumental tracks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a779ace0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load and split audio file\n",
    "a = Audio.from_file(\"../audios/halo.wav\", sample_rate=44_100, n_channels=2)\n",
    "a_vocals, a_other = split_vocals(a.convert(44_100, 2, 2))\n",
    "\n",
    "# Encode both vocals and instrumental to semantic tokens\n",
    "vocals_arr = encode_semantic(a_vocals.convert(44_100, 2, 2).normalize_volume())[:, :1]\n",
    "instrumental_arr = encode_semantic(a_other.convert(44_100, 2, 2).normalize_volume())[\n",
    "    :, :1\n",
    "]\n",
    "\n",
    "print(f\"Instrumental semantic shape: {instrumental_arr.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fc2a900",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create generation config with overpaint\n",
    "gconf = GenerationConfig(\n",
    "    text=sample_lyrics,\n",
    "    overpaint_arr=instrumental_arr,\n",
    "    cfg_coef=1.2,\n",
    "    n_repeat_tags=1,\n",
    "    n_batch=1,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    "    max_gen_duration_s=2 * 60,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4462f43",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate semantic tokens with overpaint conditioning\n",
    "requests = [\n",
    "    make_request(f\"{i}\", gconf, engine.model.config, engine.tokenizer)\n",
    "    for i in range(N_BATCH)\n",
    "]\n",
    "jobs = engine.run_request(requests, tqdm_enabled=True)\n",
    "\n",
    "out_gpt = []\n",
    "for n, job in enumerate(jobs):\n",
    "    stream = engine.token_generator(job)\n",
    "    arr = torch.stack(list(stream))[:, 1]\n",
    "    if arr[-1] == 4000:\n",
    "        arr = arr[:-1]\n",
    "    print(f\"{round(arr.shape[-1]/25)}s for track {n}\")\n",
    "    out_gpt.append(arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c4e154e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate and play audio from semantic tokens\n",
    "for in_sem_arr in out_gpt:\n",
    "    audio = generate_from_semantic(\n",
    "        in_sem_arr,\n",
    "        lyrics=sample_lyrics,\n",
    "        tags=\"\",\n",
    "        text_cfg_coef=1.0,\n",
    "        steps=16,\n",
    "        seed=0,\n",
    "    )\n",
    "    audio.play()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c95fb777",
   "metadata": {},
   "source": [
    "## Cover: Style Transfer for Complete Audio\n",
    "\n",
    "Generate a cover version in a different style from an existing complete audio track."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f21c1f65",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load audio file for cover generation\n",
    "audio_filepath = \"../audios/halo.wav\"\n",
    "a = Audio.from_file(audio_filepath, sample_rate=44_100, n_channels=2)\n",
    "cover_arr = encode_semantic(a.normalize_volume())[:, :1]\n",
    "\n",
    "print(f\"Cover semantic shape: {cover_arr.shape}\")\n",
    "a.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75677dcc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define target style and lyrics\n",
    "cover_text = sample_lyrics\n",
    "cover_tags = \"folk music, haunting, sad, female vocals, emotive, guitar\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99745ae8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create generation config with cover conditioning\n",
    "gconf = GenerationConfig(\n",
    "    text=cover_text,\n",
    "    text_tags=cover_tags,\n",
    "    cover_arr=cover_arr,\n",
    "    cfg_coef=1.2,\n",
    "    cfg_coef_tags=2.5,\n",
    "    n_repeat_tags=1,\n",
    "    n_batch=1,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    "    max_gen_duration_s=2 * 60,\n",
    "    random_seed=0,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2a70da1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate semantic tokens with cover conditioning\n",
    "requests = [\n",
    "    make_request(f\"{i}\", gconf, engine.model.config, engine.tokenizer)\n",
    "    for i in range(N_BATCH)\n",
    "]\n",
    "jobs = engine.run_request(requests, tqdm_enabled=True)\n",
    "\n",
    "out_gpt = []\n",
    "for n, job in enumerate(jobs):\n",
    "    stream = engine.token_generator(job)\n",
    "    arr = torch.stack(list(stream))[:, 1]\n",
    "    if arr[-1] == 4000:\n",
    "        arr = arr[:-1]\n",
    "    print(f\"{round(arr.shape[-1]/25)}s for track {n}\")\n",
    "    out_gpt.append(arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d44152cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate and play audio from semantic tokens\n",
    "for in_sem_arr in out_gpt:\n",
    "    audio = generate_from_semantic(\n",
    "        in_sem_arr,\n",
    "        lyrics=cover_text,\n",
    "        tags=cover_tags,\n",
    "        text_cfg_coef=1.0,\n",
    "        steps=16,\n",
    "        seed=0,\n",
    "    )\n",
    "    audio.play()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19570e8d",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "# Debugging & Utilities"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ef90e5e",
   "metadata": {},
   "source": [
    "## Visualize Prompt Structure\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e5ff457",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize the prompt structure for debugging\n",
    "from suno_utils.gpt.prompt import Prompt\n",
    "\n",
    "if len(requests) > 0:\n",
    "    in_arr = requests[0].streams[0].prompt\n",
    "    prompt = Prompt(\"\", engine.model.config)\n",
    "    print(\"Compressed view:\")\n",
    "    prompt.visualize(in_arr, compress=True)\n",
    "    print(\"\\nFull view:\")\n",
    "    prompt.visualize(in_arr, compress=False)\n",
    "else:\n",
    "    print(\"No requests available to visualize\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "699bda31",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Utility: Convert Checkpoint Type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f71ce836",
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_checkpoint_type(path: str, output_path: str = \"converted.pt\") -> None:\n",
    "    \"\"\"Convert checkpoint by ensuring best_val_loss is float type.\n",
    "\n",
    "    Args:\n",
    "        path: Path to input checkpoint\n",
    "        output_path: Path to save converted checkpoint\n",
    "    \"\"\"\n",
    "    d = torch.load(path)\n",
    "    d[\"best_val_loss\"] = float(d[\"best_val_loss\"])\n",
    "    torch.save(d, output_path)\n",
    "    print(f\"Checkpoint converted and saved to {output_path}\")\n",
    "\n",
    "\n",
    "# Uncomment to use:\n",
    "# convert_checkpoint_type(\"/path/to/checkpoint.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "312a749f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "242acd6e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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": 5
}
