{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0d83737",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !nvidia-smi\n",
    "# !echo $HOSTNAME"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fed7c348",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"3\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51b8c00a-f1c4-49fb-9775-ee990a6ae103",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # create a model mix\n",
    "# mix_ratio = 3.0  # how much of raw to use\n",
    "# import torch\n",
    "# from collections import OrderedDict\n",
    "# cpd = \"/app/suno/checkpoints/\"\n",
    "# sd1 = torch.load(f\"{cpd}2025-05-21_22-39-00/last_ckpt_infer.pt\", map_location=\"cpu\", weights_only=True)\n",
    "# sd2 = torch.load(f\"{cpd}2025-05-19_13-46-02/last_ckpt_infer.pt\", map_location=\"cpu\", weights_only=True)\n",
    "# assert set(sd1.keys()) == set(sd2.keys())\n",
    "# sd1[\"model\"] = {k: sd1[\"model\"][k] * mix_ratio + sd2[\"model\"][k] * (1 - mix_ratio) for k in sd1[\"model\"].keys()}\n",
    "# torch.save(sd1, \"/home/tony/Data/checkpoints/gpt/test_add_30.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de4064c3-fdf5-4f29-bebd-8ba3c74f57c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model_ckpts = [\"2025-05-26_17-57-55\", \"2025-05-26_18-38-40\", \"2025-05-26_18-37-43\", \"2025-05-26_18-53-55\", \"2025-05-26_20-14-40\"]\n",
    "# import torch\n",
    "# from collections import OrderedDict\n",
    "# cpd = \"/app/suno/checkpoints/\"\n",
    "# sds = []\n",
    "# for model_path in model_ckpts:\n",
    "#     sds.append(torch.load(f\"{cpd}{model_path}/last_ckpt_infer.pt\", map_location=\"cpu\", weights_only=True))\n",
    "\n",
    "# sd1 = sds[0]\n",
    "# averaged_model = {}\n",
    "# for k in sd1[\"model\"].keys():\n",
    "#     averaged_model[k] = sum(sd_x[\"model\"][k] for sd_x in sds) / len(sds)\n",
    "# sd1[\"model\"] = averaged_model\n",
    "# torch.save(sd1, \"/home/tony/Data/checkpoints/gpt/model_45_6b_bjay_sft_0519_t1_r8_5ave.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0d2ac4c-cba0-4a1d-be09-7036a0b78d52",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import torch\n",
    "sys.path.insert(0, \"/home/tony/Work/neon/sunoDiff/\")\n",
    "\n",
    "diff_model_fp = \"/app2/suno/modal/models/tony/tmp/diff/v45_2b_step_2mil_ft_8k_infill_apr21_d4_v39.pt\"  # diff v2-4\n",
    "# diff_model_fp = \"/app/suno/checkpoints/2025-01-19_05-10-13_s1983/step_3000_infer.pt\"  # best\n",
    "\n",
    "from generation import preload_models as preload_diff_models, generate\n",
    "\n",
    "_ = preload_diff_models(\n",
    "    tokenizer_filepath=\"/app/suno/data/dpo/models/tokenizer_60k.json\",\n",
    "    semantic_model_filepath=\"/app/suno/data/dpo/models/mert_25.pt\",\n",
    "    semantic_clusters_filepath=\"/app/suno/data/dpo/models/mert_25_2x4k.npy\",\n",
    "    codec_filepath=\"/app/suno/data/dpo/models/dac_vae_tuned_25hz.pth\",\n",
    "    dit_model_filepath=diff_model_fp,\n",
    "    weights_precision=torch.bfloat16,\n",
    "    compile=True,\n",
    "    codec_scale_factor=0.4\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74d60c8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\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\n",
    "\n",
    "N_BATCH = 4\n",
    "\n",
    "engine = Engine(\n",
    "    #     \"/app/suno/checkpoints/2025-01-26_01-35-19/last_ckpt_infer.pt\",  # skip\n",
    "    # \"/home/tony/Data/checkpoints/gpt/2025-01-31_17-37-50_feb04.pt\",  # base long\n",
    "    # \"/app/suno/checkpoints/2025-02-24_01-06-45/last_ckpt_infer.pt\", # 6b dpo base 24\n",
    "    \"/app2/suno/checkpoints/2025-09-10_04-35-21/last_ckpt_infer.pt\", # d87\n",
    "    # \"/app2/suno/checkpoints/2025-08-19_20-48-40/last_ckpt_infer.pt\", # +8 to bluejay\n",
    "    # \"/home/tony/Data/checkpoints/gpt/model_45_6b_bjay_sft_0519_t1_r8_5ave.pt\",\n",
    "    #     \"/app/suno/checkpoints/2025-01-18_18-03-13/last_ckpt_infer.pt\",  # base\n",
    "    #     \"/app/suno/checkpoints/2024-12-05_22-01-59/last_ckpt_infer.pt\",  # (old chef)\n",
    "    \"/app/suno/models/chirp_v2/tokenizer_60k.json\",\n",
    "    max_sequences=4 * N_BATCH,\n",
    "    compile=False,\n",
    ")\n",
    "cfg = engine.model.config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "caa2446f-9692-42bb-8d60-db11478ed3c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_clips(text, tags, neg_tags=None, cfg_coef_tags_max_steps=None, control_tags=None):\n",
    "    n_skip_semantic = 1\n",
    "    # n_skip_semantic = 4\n",
    "    \n",
    "    gconf = GenerationConfig(\n",
    "        text=text,\n",
    "        text_tags=tags,\n",
    "        cfg_coef=1.0,\n",
    "        cfg_coef_tags=1.0,\n",
    "        # cfg_coef_max_steps=None,\n",
    "        cfg_coef_tags_max_steps=25 * 120 if cfg_coef_tags_max_steps is None 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",
    "        #     temp_coarse=0.9,\n",
    "        #     top_k_semantic=None,\n",
    "        #     top_k_coarse=None,\n",
    "        min_p_semantic=0.005,\n",
    "        # top_p_semantic=None,\n",
    "        # top_p_coarse=None,\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",
    "    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",
    "    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",
    "        # do stuff incase 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",
    "        # add\n",
    "        out_gpt.append(arr2)\n",
    "        \n",
    "    for in_sem_arr in out_gpt:\n",
    "        generate(\n",
    "            in_sem_arr,\n",
    "            lyrics=text,\n",
    "            tags=tags,\n",
    "            text_cfg_coef=2.0,\n",
    "            steps=10,\n",
    "            seed=0,\n",
    "            semantic_skip_factor=1,\n",
    "            downscale_ctx_vector=False,\n",
    "            noise_ctx_vector=0.5,\n",
    "            noise_ctx_vector_pad_size=25,\n",
    "        ).play()    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ec28b37",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # [Verse 1: female vocals, soft, beautiful]\n",
    "# # [Pre-Chorus: male, aggressive, shouting]\n",
    "# # [Chorus: female & male, together, harmonies]\n",
    "\n",
    "# text = \"\"\"\n",
    "# [Verse 1]\n",
    "# You are my fire\n",
    "# The one desire\n",
    "# Believe when I say\n",
    "# I want it that way\n",
    "\n",
    "# [Verse 2]\n",
    "# But we are two worlds apart\n",
    "# Can't reach to your heart\n",
    "# When you say\n",
    "# That I want it that way\n",
    "\n",
    "# [Chorus]\n",
    "# Tell me why\n",
    "# Ain't nothin' but a heartache\n",
    "# Tell me why\n",
    "# Ain't nothing but a mistake\n",
    "# Tell me why\n",
    "# I never wanna hear you say\n",
    "# I want it that way\n",
    "\n",
    "# [Verse 3]\n",
    "# Am I your fire?\n",
    "# Your one desire\n",
    "# Yes, I know it's too late\n",
    "# But I want it that way\n",
    "\n",
    "# [Chorus]\n",
    "# Tell me why\n",
    "# Ain't nothin' but a heartache\n",
    "# Tell me why\n",
    "# Ain't nothin' but a mistake\n",
    "# Tell me why\n",
    "# I never wanna hear you say\n",
    "# I want it that way\n",
    "\n",
    "# [Bridge]\n",
    "# Now I can see that we've fallen apart\n",
    "# From the way that it used to be, yeah\n",
    "# No matter the distance, I want you to know\n",
    "# That deep down inside of me\n",
    "\n",
    "# [Verse 4]\n",
    "# You are my fire\n",
    "# The one desire\n",
    "# You are (You are, you are, you are)\n",
    "# Don't wanna hear you say\n",
    "\n",
    "# [Break]\n",
    "# Ain't nothin' but a heartache (Hey, yeah)\n",
    "# Ain't nothin' but a mistake (Don't wanna hear you say)\n",
    "# I never wanna hear you say (Oh, yeah)\n",
    "# I want it that way\n",
    "# [Chorus: Brian, All, Nick, AJ]\n",
    "# Tell me why\n",
    "# Ain't nothin' but a heartache\n",
    "# Tell me why\n",
    "# Ain't nothin' but a mistake\n",
    "# Tell me why\n",
    "# I never wanna hear you say (Don't wanna hear you say)\n",
    "# I want it that way\n",
    "# Tell me why\n",
    "# Ain't nothin' but a heartache\n",
    "# Ain't nothin' but a mistake\n",
    "# Tell me why\n",
    "# I never wanna hear you say (Never wanna hear you say, yeah)\n",
    "# I want it that way\n",
    "\n",
    "# [Outro]\n",
    "# 'Cause I want it that way\n",
    "# \"\"\"\n",
    "\n",
    "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",
    "\"\"\"\n",
    "\n",
    "# tags = \"grunge alternative rock\"\n",
    "# tags = \"guitar violin drums\"\n",
    "# tags = \"A melancholic country song featuring prominent pedal steel guitar.\"\n",
    "\n",
    "# tags = \"Pop\"\n",
    "# tags = \"Gregorian Chant, Hymns, Christian, ethereal, female vocalist, medieval, atmospheric, spiritual, melodic, choral, peaceful, hypnotic, hymn, meditative, minimalistic, soothing, religious, soft, longing, dense, melancholic, mysterious, repetitive, acoustic\"\n",
    "# tags = \"Two-track instrumental featuring ambient nature sounds and rhythmic clanging, followed by sinister synthesizer music and buzzing electric circuit.\"\n",
    "# tags = \"Country, southern accent, deep south, raspy vocals\"\n",
    "# tags = \"Gospel\"\n",
    "# tags = \"Americana, Country, Contemporary Folk, lonely, sombre, introspective, longing, acoustic, melodic, bittersweet, pastoral, poetic, melancholic, male vocalist, depressive, pessimistic, dark, passionate\"\n",
    "# tags = \"Contemporary R&B, Pop Rap, Alternative R&B, Trap, nocturnal, female vocalist, aquatic, alcohol, male vocalist, sexual\"\n",
    "# tags = \"Pop Soul, Neo-Soul, Jazz Pop, Brill Building, Twist, drugs, female vocalist, rebellious, alcohol, rhythmic, depressive, introspective, passionate, sarcastic, melodic, hedonistic, anthemic\"\n",
    "# tags = \"power ballad, female vocals, male vocals, pop, rap, hip-hop, emotive, powerful, haunting\"\n",
    "# tags = \"Melodic Techno, Future Rave, Melodic House, rhythmic, atmospheric\"\n",
    "# tags = \"Dance-Pop, Electro House, Festival Progressive House, Electropop, Dutch House, passionate, female vocalist, triumphant, introspective, energetic, anthemic, rhythmic, sentimental\"\n",
    "# tags = \"passionate, raspy vocals, powerful delivery, female, pop\"\n",
    "# tags = \"Pop, female vocals, orchestra\"\n",
    "# tags = \"female vocals; This track start off quiet and slow featuring string instruments including a violin. It then builds into an energetic chorus with drums and trumpets.\"\n",
    "\n",
    "# text = \"\"\n",
    "# tags = \"orchestra, symphony\"\n",
    "# tags = \"orchestra, powerful trailer music, epic, dark, violin, female vocalizations\"\n",
    "# tags = \"The piece that opens with a soft, flute-led melody gently awakening the senses, conjuring the warmth of a sunrise over tranquil hills. The orchestra gradually swells with smooth strings and light percussion, evoking a blossoming sense of optimism and renewed energy. A graceful interplay of woodwinds and strings concludes on a delicate, uplifting note that captures the fresh promise of morning.\"\n",
    "# tags = \"A stirring, festive march led by bold brass fanfares and triumphant rhythms that command attention right from the start. Let snappy snare drums and buoyant woodwinds propel the music forward with an infectious energy. The piece culminates in a grand, celebratory finale that invites enthusiastic clapping and cheering.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3133ec0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_clips(text, tags=\"epic film orchestral\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2064fd87-541a-465c-8a47-8510e9123f47",
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate_clips(text, tags=\"k-pop\")"
   ]
  },
  {
   "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 blieve\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, tags=\"gregorian chant\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c044be24-be21-4810-82b1-bd30d05ffc31",
   "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",
    "generate_clips(text, tags=\"gregorian chant\", control_tags=\"tag_strength:1;\" )\n",
    "# generate_clips(text, tags=\"neo soul\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4707f57c-3269-43fd-a189-eb8b51898046",
   "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",
    "generate_clips(text, tags=\"gregorian chant\", control_tags=\"tag_strength:9\" )\n",
    "# generate_clips(text, tags=\"neo soul\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c9652cc-7ec3-405e-80fc-ffb2b58e5186",
   "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",
    "generate_clips(text, tags=\"gregorian chant, no instrumental\", control_tags=\"tag_strength:9\" )\n",
    "# generate_clips(text, tags=\"neo soul\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b0c144b-96f7-4b8b-ab9b-8db235eb78d0",
   "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 blieve\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, tags=\"gregorian chant, gregorian chant, gregorian chant\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ad05205-272a-4fc6-bd8e-6193111cb141",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[bagpipes, verse]\n",
    "A is for the amazing grace that we recieve\n",
    "B is for the blessings, every day we blieve\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, tags=\"bagpipes, scottish funeral march\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f64ec057",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"\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, tags=\"chinese pop woman\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5cb96fc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_clips(\"\", tags=\"late 19th century austrian romantic symphony, dramatic, large orchestra\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "741261a7-70bb-46a1-94f5-2547c17e6484",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "generate_clips(\"[Dramatic slow opening, drums, strong strings, emotional, intense] [Woodwinds comes in with constrast, delicate, sad, sorrow] [Back strings, fast tempo, echo the theme]\", \n",
    "               tags=\"\"\"\n",
    "beethoven, beethoven, beethoven, ludwig van beethoven, ludwig van beethoven as the composer,\n",
    "dg, philharmonic orchestra, DG, symphony 1st movement\"\"\",\n",
    "               neg_tags=\"vocals, vocals, pianos, classical, baroque, pop, mozart, opera, mozart, haydn, brahms, chopin, overture, overture, march, march\",\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": [
    "generate_clips(\"\", tags=\"JS bach, JSbach, fugue, piano, baroque\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16e662b1-6515-4e5c-90b0-4d429600edfc",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_clips(\"\", tags=\"fugue, piano, baroque, slow tempo\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9769cd6-0b69-4dcb-bb24-cc72108fd0c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"\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, tags=\"female bluegrass, fast tempo\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0964b676-c54b-4463-8cd2-8bbb80823835",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_clips(\"plastic 1, plastic 2, plastic 3, plastic 4, plastic 5\", tags=\"pop\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbe4f87d",
   "metadata": {},
   "source": [
    "### Over/Underpaint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6566827e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "from suno_utils.audio import Audio\n",
    "from generation import encode_semantic\n",
    "from suno_utils.tasks.demucs import split_vocals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aaf94996",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from suno_utils.utils.text import read_jsonl\n",
    "\n",
    "mm = np.memmap(\"/app/suno/data/chirp_v5/v1/data_val.bin\", dtype=np.uint16, mode=\"r\")\n",
    "metas = read_jsonl(\"/app/suno/data/chirp_v5/v1/metas_val.jsonl\")\n",
    "overpaint_metas = [\n",
    "    m\n",
    "    for m in metas\n",
    "    if m[\"dataset\"] == \"musdb_stems_overpaint\" and m[\"type\"] == \"instrumental\"\n",
    "]\n",
    "underpaint_metas = [\n",
    "    m\n",
    "    for m in metas\n",
    "    if m[\"dataset\"] == \"musdb_stems_underpaint\" and m[\"type\"] == \"vocals\"\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f98dcf7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"\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",
    "I found a way to let you in\n",
    "But I never really had a doubt\n",
    "Standin' in the light of your halo\n",
    "I got my angel now\n",
    "\n",
    "[Pre-Chorus]\n",
    "It's like I've been awakened\n",
    "Every rule, I had you breakin'\n",
    "It's the risk that I'm takin'\n",
    "I ain't ever gonna shut you out\n",
    "\n",
    "[Chorus]\n",
    "Everywhere I'm lookin' now\n",
    "I'm surrounded by your embrace\n",
    "Baby, I can see your halo\n",
    "You know you're my savin' grace\n",
    "You're everything I need and more\n",
    "It's written all over your face\n",
    "Baby, I can feel your halo\n",
    "Pray it won't fade away\n",
    "\n",
    "[Post-Chorus]\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo, ooh\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",
    "I swore I'd never fall again\n",
    "But this don't even feel like fallin'\n",
    "Gravity can't begin\n",
    "To pull me back to the ground again\n",
    "\n",
    "[Pre-Chorus]\n",
    "It's like I've been awakened\n",
    "Every rule, I had you breakin'\n",
    "The risk that I'm takin'\n",
    "I'm never gonna shut you out\n",
    "\n",
    "[Chorus]\n",
    "Everywhere I'm lookin' now\n",
    "I'm surrounded by your embrace\n",
    "Baby, I can see your halo\n",
    "You know you're my savin' grace\n",
    "You're everything I need and more\n",
    "It's written all over your face\n",
    "Baby, I can feel your halo\n",
    "Pray it won't fade away\n",
    "\n",
    "[Post-Chorus]\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo, ooh\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo, ooh\n",
    "\n",
    "[Bridge]\n",
    "Halo, ooh\n",
    "Halo, ooh\n",
    "Ooh\n",
    "\n",
    "[Chorus]\n",
    "Everywhere I'm lookin' now\n",
    "I'm surrounded by your embrace\n",
    "Baby, I can see your halo\n",
    "You know you're my savin' grace\n",
    "You're everything I need and more\n",
    "It's written all over your face\n",
    "Baby, I can feel your halo\n",
    "Pray it won't fade away\n",
    "\n",
    "[Post-Chorus]\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo, ooh\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo, ooh\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa3cd8c8",
   "metadata": {},
   "source": [
    "#### Underpaint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28bf850c",
   "metadata": {},
   "outputs": [],
   "source": [
    "a = Audio.from_file(\"../samples/halo.wav\", sample_rate=44_100, n_channels=2)\n",
    "# # a = Audio.from_file(\"../samples/calm_down.mp3\", sample_rate=44_100, n_channels=2)\n",
    "\n",
    "a_vocals, a_other = split_vocals(a.convert(44_100, 2, 2))\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())[:,:1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10e6d9b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# m = random.choice(underpaint_metas)\n",
    "# vocals_arr = mm[m[\"offset_idx\"]:m[\"offset_idx\"]+m[\"n_tokens\"]][:60*4][None].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1023d4b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "tags = \"Pop\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fa2fa71",
   "metadata": {},
   "outputs": [],
   "source": [
    "gconf = GenerationConfig(\n",
    "    text=text,\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": [
    "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",
    "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": [
    "for in_sem_arr in out_gpt:\n",
    "    generate(\n",
    "        in_sem_arr,\n",
    "        lyrics=text,\n",
    "        tags=tags,\n",
    "        text_cfg_coef=1.0,\n",
    "        steps=16,\n",
    "        seed=0,\n",
    "    ).play()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f687e4c",
   "metadata": {},
   "source": [
    "#### Overpaint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a779ace0",
   "metadata": {},
   "outputs": [],
   "source": [
    "a = Audio.from_file(\"../samples/halo.wav\", sample_rate=44_100, n_channels=2)\n",
    "# a = Audio.from_file(\"../samples/calm_down.mp3\", sample_rate=44_100, n_channels=2)\n",
    "\n",
    "a_vocals, a_other = split_vocals(a.convert(44_100, 2, 2))\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",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "233a72dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# a_vocals = Audio.from_file(\"../samples/martin_vocals_up.mp3\", sample_rate=44_100, n_channels=2)\n",
    "# vocals_arr = encode_semantic(a_vocals.convert(44_100, 2, 2).normalize_volume())[:,:1]\n",
    "\n",
    "# a_instrumental = Audio.from_file(\"../samples/martin_instrumental.m4a\", sample_rate=44_100, n_channels=2)\n",
    "# instrumental_arr = encode_semantic(a_instrumental.convert(44_100, 2, 2).normalize_volume())[:,:1]\n",
    "\n",
    "# m = random.choice(overpaint_metas)\n",
    "# instrumental_arr = mm[m[\"offset_idx\"]:m[\"offset_idx\"]+m[\"n_tokens\"]][:60*4][None].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fc2a900",
   "metadata": {},
   "outputs": [],
   "source": [
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "    #     text_tags=\"pop, sentimental, folk\",\n",
    "    overpaint_arr=instrumental_arr,\n",
    "    cfg_coef=1.2,\n",
    "    #     cfg_coef_tags=2.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": "d4462f43",
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "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": [
    "for in_sem_arr in out_gpt:\n",
    "    generate(\n",
    "        in_sem_arr,\n",
    "        lyrics=text,\n",
    "        #         tags=tags,\n",
    "        text_cfg_coef=1.0,\n",
    "        steps=16,\n",
    "        seed=0,\n",
    "    ).play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b1d7b0b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6267ad78",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abc1d47c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c61e6914",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "c95fb777",
   "metadata": {},
   "source": [
    "### Cover"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f21c1f65",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.audio import Audio\n",
    "from generation import encode_semantic\n",
    "from suno_utils.tasks.demucs import split_vocals\n",
    "\n",
    "# audio_filepath = \"../samples/martin2.m4a\"\n",
    "audio_filepath = \"../samples/halo.wav\"\n",
    "# audio_filepath = \"../samples/corn_chase.mp3\"\n",
    "# audio_filepath = \"../samples/martin_vocals.m4a\"\n",
    "\n",
    "a = Audio.from_file(audio_filepath, sample_rate=44_100, n_channels=2)\n",
    "cover_arr = encode_semantic(a.normalize_volume())[:, :1]\n",
    "a.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75677dcc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# text = \"\"\"\n",
    "# [Verse]\n",
    "# Raindrops falling away\n",
    "# Dancing in the grey\n",
    "# Clouds are here to stay\n",
    "# Love on a rainy day\n",
    "\n",
    "# [Verse 2]\n",
    "# Puddles on the ground\n",
    "# Feet splashing 'round\n",
    "# You and me we found\n",
    "# Paradise in the sound\n",
    "\n",
    "# [Chorus]\n",
    "# Love on a rainy day\n",
    "# Washing doubts away\n",
    "# Hold me and let's sway\n",
    "# It's our perfect cliche\n",
    "\n",
    "# Love on a rainy day\n",
    "# Washing doubts away\n",
    "# Hold me and let's sway\n",
    "# It's our perfect cliche\n",
    "# \"\"\"\n",
    "\n",
    "# text = \"\"\"\n",
    "# Well, you only need the light when it's burning low\n",
    "# Only miss the sun when it starts to snow\n",
    "# Only know you love her when you let her go\n",
    "# Only know you've been high when you're feeling low\n",
    "# Only hate the road when you're missing home\n",
    "# Only know you love her when you let her go\n",
    "\n",
    "# And you let her go\n",
    "# \"\"\"\n",
    "\n",
    "text = \"\"\"\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",
    "I found a way to let you in\n",
    "But I never really had a doubt\n",
    "Standin' in the light of your halo\n",
    "I got my angel now\n",
    "\n",
    "[Pre-Chorus]\n",
    "It's like I've been awakened\n",
    "Every rule, I had you breakin'\n",
    "It's the risk that I'm takin'\n",
    "I ain't ever gonna shut you out\n",
    "\n",
    "[Chorus]\n",
    "Everywhere I'm lookin' now\n",
    "I'm surrounded by your embrace\n",
    "Baby, I can see your halo\n",
    "You know you're my savin' grace\n",
    "You're everything I need and more\n",
    "It's written all over your face\n",
    "Baby, I can feel your halo\n",
    "Pray it won't fade away\n",
    "\n",
    "[Post-Chorus]\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo, ooh\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",
    "I swore I'd never fall again\n",
    "But this don't even feel like fallin'\n",
    "Gravity can't begin\n",
    "To pull me back to the ground again\n",
    "\n",
    "[Pre-Chorus]\n",
    "It's like I've been awakened\n",
    "Every rule, I had you breakin'\n",
    "The risk that I'm takin'\n",
    "I'm never gonna shut you out\n",
    "\n",
    "[Chorus]\n",
    "Everywhere I'm lookin' now\n",
    "I'm surrounded by your embrace\n",
    "Baby, I can see your halo\n",
    "You know you're my savin' grace\n",
    "You're everything I need and more\n",
    "It's written all over your face\n",
    "Baby, I can feel your halo\n",
    "Pray it won't fade away\n",
    "\n",
    "[Post-Chorus]\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo, ooh\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo, ooh\n",
    "\n",
    "[Bridge]\n",
    "Halo, ooh\n",
    "Halo, ooh\n",
    "Ooh\n",
    "\n",
    "[Chorus]\n",
    "Everywhere I'm lookin' now\n",
    "I'm surrounded by your embrace\n",
    "Baby, I can see your halo\n",
    "You know you're my savin' grace\n",
    "You're everything I need and more\n",
    "It's written all over your face\n",
    "Baby, I can feel your halo\n",
    "Pray it won't fade away\n",
    "\n",
    "[Post-Chorus]\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo, ooh\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo, ooh\n",
    "\"\"\"\n",
    "\n",
    "tags = \"folk music, haunting, sad, female vocals, emotive, guitar\"\n",
    "# tags = \"R&B, Pop, Contemporary R&B, Ballad, male vocals\"\n",
    "\n",
    "# text = \"\"\n",
    "# tags = \"Big Room Trance\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99745ae8",
   "metadata": {},
   "outputs": [],
   "source": [
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "    text_tags=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": [
    "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",
    "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": [
    "for in_sem_arr in out_gpt:\n",
    "    generate(\n",
    "        in_sem_arr,\n",
    "        lyrics=text,\n",
    "        tags=tags,\n",
    "        text_cfg_coef=1.0,\n",
    "        steps=16,\n",
    "        seed=0,\n",
    "    ).play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3452380a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be956547",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "19570e8d",
   "metadata": {},
   "source": [
    "### Visualize prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e5ff457",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.gpt.prompt import Prompt\n",
    "\n",
    "in_arr = requests[0].streams[0].prompt\n",
    "prompt = Prompt(\"\", engine.model.config)\n",
    "prompt.visualize(in_arr, compress=True)\n",
    "prompt.visualize(in_arr, compress=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d042392a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "699bda31",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "def convert_to_type(path):\n",
    "    d = torch.load(path)\n",
    "    d[\"best_val_loss\"] = float(d[\"best_val_loss\"])\n",
    "    torch.save(d, \"converted.pt\")\n",
    "    print(\"DONE\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "312a749f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7f5cce1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03a473d7",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbb3d211",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "addcbe70",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b290ea0b",
   "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
}
