{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fae8f9c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:00:28.473525Z",
     "start_time": "2024-02-26T22:00:24.633479Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\"\n",
    "\n",
    "import numpy as np\n",
    "from suno_utils.audio import Audio\n",
    "from suno_utils.gpt.chirp_v2_5 import (\n",
    "    generate_audio,\n",
    "    generate_audio_stream,\n",
    "    preload_models,\n",
    "    codec_decode,\n",
    "    GenerationConfig,\n",
    "    semantic_encode,\n",
    "    codec_encode,\n",
    ")\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01a25410",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:00:28.621641Z",
     "start_time": "2024-02-26T22:00:28.475714Z"
    }
   },
   "outputs": [],
   "source": [
    "!ls /app/suno/checkpoints/2024-02-16_15-11-50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9108708b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:11.310953Z",
     "start_time": "2024-02-26T22:00:28.623186Z"
    }
   },
   "outputs": [],
   "source": [
    "_ = preload_models(\n",
    "    gpt_ckpt_path=\"/app/suno/checkpoints/2024-02-17_14-43-47/last_ckpt_infer.pt\", \n",
    "    tokenizer_path=\"/home/victor/data/models/chirp_v2/tokenizer_60k.json\",\n",
    "    semantic_ckpt_path=\"/home/victor/data/models/chirp_v2/mert_25.pt\",\n",
    "    semantic_centroids_path=\"/home/victor/data/models/chirp_v2/mert_25_2x4k.npy\",\n",
    "    codec_ckpt_path=\"/app/suno/tony/v3/dac_2c_25x12.pt\",\n",
    "    fasttext_ckpt_path=\"/home/victor/data/models/lid.176.bin\",\n",
    "    load_whisper=True,\n",
    "    use_gpu=True,\n",
    "    load_codec_device=\"cuda\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "072adcda",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:11.315241Z",
     "start_time": "2024-02-26T22:01:11.312604Z"
    }
   },
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[Verse 2]\n",
    "一壶漂泊浪迹天涯难入喉\n",
    "你走之后酒暖回忆思念瘦\n",
    "水向东流时间怎么偷\n",
    "花开就一次成熟我却错过\n",
    "[Chorus]\n",
    "谁在用琵琶弹奏一曲东风破\n",
    "岁月在墙上剥落看见小时候\n",
    "犹记得那年我们都还很年幼\n",
    "而如今琴声幽幽我的等候你没听过\n",
    "\n",
    "-\n",
    "\"\"\"\n",
    "# [Verse 1]\n",
    "# 一盏离愁孤灯伫立在窗口\n",
    "# 我在门后假装你人还没走\n",
    "# 旧地如重游月圆更寂寞\n",
    "# 夜半清醒的烛火不忍苛责我\n",
    "# [Chorus][Chorus][Chorus]\n",
    "# 谁在用琵琶弹奏一曲东风破\n",
    "# 枫叶将故事染色结局我看透\n",
    "# 篱笆外的古道我牵着你走过\n",
    "# 荒烟蔓草的年头就连分手都很沉默"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f5a71c6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.764924Z",
     "start_time": "2024-02-26T22:01:11.317440Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=text,\n",
    "        text_tags=\"pop\",\n",
    "        n_batch=4,\n",
    "        max_gen_duration_s=60,\n",
    "        # stream=False,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf3538f1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.768199Z",
     "start_time": "2024-02-26T22:01:16.768189Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bee63248",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.768992Z",
     "start_time": "2024-02-26T22:01:16.768983Z"
    }
   },
   "outputs": [],
   "source": [
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fa595a6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.769705Z",
     "start_time": "2024-02-26T22:01:16.769695Z"
    }
   },
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "268b742c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.770462Z",
     "start_time": "2024-02-26T22:01:16.770453Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=text,\n",
    "        text_tags=\"Chopin Piano\",\n",
    "        n_batch=2,\n",
    "        cfg_coef=1,\n",
    "        max_gen_duration_s=240,\n",
    "        text_start_control_tags=\"{start}\",\n",
    "        allow_eos=False,\n",
    "        # text_end_control_tags=\"{end}\",\n",
    "        # stream=False,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "388c74c0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.771155Z",
     "start_time": "2024-02-26T22:01:16.771146Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7a281ab",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.771866Z",
     "start_time": "2024-02-26T22:01:16.771858Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=text,\n",
    "        text_tags=\"Mozart Piano\",\n",
    "        n_batch=4,\n",
    "        cfg_coef=1,\n",
    "        max_gen_duration_s=240,\n",
    "        text_start_control_tags=\"{start}\",\n",
    "        allow_eos=False,\n",
    "        # text_end_control_tags=\"{end}\",\n",
    "        # stream=False,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88cf31a0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.772572Z",
     "start_time": "2024-02-26T22:01:16.772563Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08d15aa1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.773290Z",
     "start_time": "2024-02-26T22:01:16.773281Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=text,\n",
    "        text_tags=\"Bach fugue cello\",\n",
    "        n_batch=4,\n",
    "        cfg_coef=1,\n",
    "        max_gen_duration_s=240,\n",
    "        text_start_control_tags=\"{start}\",\n",
    "        allow_eos=False,\n",
    "        # stream=False,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3005f6a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.773838Z",
     "start_time": "2024-02-26T22:01:16.773831Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e869a6da",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.774566Z",
     "start_time": "2024-02-26T22:01:16.774558Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=text,\n",
    "        text_tags=\"Beethoven Piano Sonata 4th movement\",\n",
    "        n_batch=4,\n",
    "        cfg_coef=1,\n",
    "        max_gen_duration_s=240,\n",
    "        text_start_control_tags=\"{start}\",\n",
    "        allow_eos=False,\n",
    "        # stream=False,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c4db226",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.775348Z",
     "start_time": "2024-02-26T22:01:16.775341Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9592f21d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.776035Z",
     "start_time": "2024-02-26T22:01:16.776027Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=\"{start}\",\n",
    "        text_tags=\"ludwig van beethoven symphony\",\n",
    "        n_batch=4,\n",
    "        cfg_coef=1,\n",
    "        max_gen_duration_s=240,\n",
    "        text_start_control_tags=\"{start}\",\n",
    "        allow_eos=False,\n",
    "        # stream=False,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "282c044f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.776740Z",
     "start_time": "2024-02-26T22:01:16.776732Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08069fc1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.777398Z",
     "start_time": "2024-02-26T22:01:16.777390Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=text,\n",
    "        text_tags=\"gustav mahler Symphony\",\n",
    "        n_batch=4,\n",
    "        cfg_coef=1,\n",
    "        max_gen_duration_s=240,\n",
    "        text_start_control_tags=\"{start}\",\n",
    "        allow_eos=False,\n",
    "        # stream=False,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2deec59b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.778141Z",
     "start_time": "2024-02-26T22:01:16.778133Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ac60784",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.778844Z",
     "start_time": "2024-02-26T22:01:16.778835Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=\"{start}\",\n",
    "        text_tags=\"Mahler Piano sonata\",\n",
    "        n_batch=4,\n",
    "        cfg_coef=1,\n",
    "        max_gen_duration_s=240,\n",
    "        text_start_control_tags=\"{start}\",\n",
    "        allow_eos=False,\n",
    "        # stream=False,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72332d7a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.779517Z",
     "start_time": "2024-02-26T22:01:16.779510Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e3c1078",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.780194Z",
     "start_time": "2024-02-26T22:01:16.780186Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=\"{start}\",\n",
    "        text_tags=\"Anton Bruckner Symphony\",\n",
    "        n_batch=4,\n",
    "        cfg_coef=1,\n",
    "        max_gen_duration_s=240,\n",
    "        text_start_control_tags=\"{start}\",\n",
    "        allow_eos=False,\n",
    "        # stream=False,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa47b805",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.780952Z",
     "start_time": "2024-02-26T22:01:16.780944Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09bcff1d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.781579Z",
     "start_time": "2024-02-26T22:01:16.781571Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=\"{start}\",\n",
    "        text_tags=\"wolfgang amadeus mozart piano concerto\",\n",
    "        n_batch=2,\n",
    "        cfg_coef=1,\n",
    "        cfg_coef_tags_max_steps=None,\n",
    "        max_gen_duration_s=240,\n",
    "        text_start_control_tags=\"{start}\",\n",
    "        allow_eos=False,\n",
    "        # stream=False,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1e44877",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.782252Z",
     "start_time": "2024-02-26T22:01:16.782244Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de819320",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.782944Z",
     "start_time": "2024-02-26T22:01:16.782936Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=\"[intro] [shiny flute] [piano intermission] [violin chorus] [cello refrain] [end]\",\n",
    "        text_tags=\"wolfgang amadeus mozart, upbeat epic film orchestra\",\n",
    "        n_batch=2,\n",
    "        cfg_coef=1.5,\n",
    "        max_gen_duration_s=240,\n",
    "        # text_start_control_tags=\"{start}\",\n",
    "        # text_neg_tags=\"loopy noisy\",\n",
    "        # cfg_coef_neg_tags=-2.0,\n",
    "        allow_eos=True,\n",
    "        cfg_coef_max_steps=None,\n",
    "        # stream=False,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dc1bf29",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.783720Z",
     "start_time": "2024-02-26T22:01:16.783711Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34624075",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.784261Z",
     "start_time": "2024-02-26T22:01:16.784253Z"
    }
   },
   "outputs": [],
   "source": [
    "BREAK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f9e0366",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.784892Z",
     "start_time": "2024-02-26T22:01:16.784885Z"
    }
   },
   "outputs": [],
   "source": [
    "input_audio = Audio.from_file(\"audios/Pachelbel - Canon In D Major. Best version. [NlprozGcs80].wav\").get_segment(from_s=60, to_s=80)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f9d3323",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.785773Z",
     "start_time": "2024-02-26T22:01:16.785765Z"
    }
   },
   "outputs": [],
   "source": [
    "input_audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "001ec800",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.786482Z",
     "start_time": "2024-02-26T22:01:16.786473Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=\"\",\n",
    "        text_tags=\"canon\",\n",
    "        n_batch=4,\n",
    "        max_gen_duration_s=30,\n",
    "        # stream=False,\n",
    "    ),\n",
    "    history_audio=input_audio\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6309815d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.787222Z",
     "start_time": "2024-02-26T22:01:16.787214Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2270b189",
   "metadata": {},
   "source": [
    "# Oracle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "772a5d20",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.787923Z",
     "start_time": "2024-02-26T22:01:16.787915Z"
    }
   },
   "outputs": [],
   "source": [
    "input_audio = Audio.from_file(\"audios/test_swift.mp3\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29ecceb0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.788470Z",
     "start_time": "2024-02-26T22:01:16.788463Z"
    }
   },
   "outputs": [],
   "source": [
    "test_input_audio = input_audio.get_segment(from_s=180, to_s=200)\n",
    "test_input_audio.play(compress=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4b5d48c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.789163Z",
     "start_time": "2024-02-26T22:01:16.789155Z"
    }
   },
   "outputs": [],
   "source": [
    "oracle_semantic_arr = semantic_encode(test_input_audio, n_codebooks=1)\n",
    "oracle_coarse_arr = codec_encode(test_input_audio)[:oracle_semantic_arr.shape[0], :]\n",
    "print(oracle_semantic_arr.shape, oracle_coarse_arr.shape)\n",
    "oracle_arr = np.concatenate(\n",
    "    [oracle_semantic_arr, oracle_coarse_arr],\n",
    "    axis=-1,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "808f5130",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.789749Z",
     "start_time": "2024-02-26T22:01:16.789740Z"
    }
   },
   "outputs": [],
   "source": [
    "input_oracle_arr = oracle_arr[:, :1]\n",
    "print(input_oracle_arr.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "752ba22f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.790408Z",
     "start_time": "2024-02-26T22:01:16.790400Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=\" \" * 20,\n",
    "        text_tags=\"\",\n",
    "        n_batch=4,\n",
    "        max_gen_duration_s=30,\n",
    "        oracle_arr=input_oracle_arr,\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bba0d8c9",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.791224Z",
     "start_time": "2024-02-26T22:01:16.791215Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()\n",
    "    print(e.array_float[-100:].mean(axis=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad979c8b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.791899Z",
     "start_time": "2024-02-26T22:01:16.791892Z"
    }
   },
   "outputs": [],
   "source": [
    "input_pavarotti = Audio.from_file(\"audios/test_pavarotti.mp3\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f0f8785",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.792444Z",
     "start_time": "2024-02-26T22:01:16.792436Z"
    }
   },
   "outputs": [],
   "source": [
    "input_pavarotti_seg = input_pavarotti.get_segment(from_s=152, to_s=172)\n",
    "input_pavarotti_seg.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e91f1d1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.793183Z",
     "start_time": "2024-02-26T22:01:16.793176Z"
    }
   },
   "outputs": [],
   "source": [
    "pavarotti_oracle_semantic_arr = semantic_encode(input_pavarotti_seg, n_codebooks=1)\n",
    "pavarotti_oracle_coarse_arr = codec_encode(input_pavarotti_seg)[\n",
    "    : pavarotti_oracle_semantic_arr.shape[0], :\n",
    "]\n",
    "print(pavarotti_oracle_semantic_arr.shape, pavarotti_oracle_coarse_arr.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6502b74a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.793929Z",
     "start_time": "2024-02-26T22:01:16.793920Z"
    }
   },
   "outputs": [],
   "source": [
    "masked_oracle_semantic_arr = oracle_semantic_arr.copy().astype(np.int16)\n",
    "# masked_oracle_semantic_arr[0::5, :] = -1\n",
    "masked_oracle_coarse_arr = pavarotti_oracle_coarse_arr.copy().astype(np.int16)\n",
    "# masked_oracle_coarse_arr[20:, :] = -1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b97481d6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.794619Z",
     "start_time": "2024-02-26T22:01:16.794611Z"
    }
   },
   "outputs": [],
   "source": [
    "mixed_oracle_arr =  np.concatenate(\n",
    "    [masked_oracle_semantic_arr], #, masked_oracle_coarse_arr[:, 0:1]],\n",
    "    axis=-1,\n",
    ")\n",
    "print(mixed_oracle_arr.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76dcc73a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.795309Z",
     "start_time": "2024-02-26T22:01:16.795301Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=\" \" * 20,\n",
    "        text_tags=\"\",\n",
    "        n_batch=4,\n",
    "        max_gen_duration_s=30,\n",
    "        oracle_arr=mixed_oracle_arr,\n",
    "    ),\n",
    "    # history_audio=input_pavarotti_seg\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84e04fd4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.796014Z",
     "start_time": "2024-02-26T22:01:16.796007Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()\n",
    "    print(e.array_float[-100:].mean(axis=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b79d93dc",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "edc87efa",
   "metadata": {},
   "source": [
    "# continuation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ce0c278",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.796696Z",
     "start_time": "2024-02-26T22:01:16.796688Z"
    }
   },
   "outputs": [],
   "source": [
    "input_audio = Audio.from_file(\"audios/test_swift.mp3\")\n",
    "\n",
    "test_input_audio = input_audio.get_segment(from_s=180, to_s=200)\n",
    "test_input_audio.play(compress=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "984e8c9b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.797314Z",
     "start_time": "2024-02-26T22:01:16.797307Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=\"[chorus] purple rain\" * 20,\n",
    "        text_tags=\"\",\n",
    "        n_batch=4,\n",
    "        max_gen_duration_s=30,\n",
    "        # stream=False,\n",
    "    ),\n",
    "    history_audio=test_input_audio,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "296b34eb",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.798029Z",
     "start_time": "2024-02-26T22:01:16.798021Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df83ed93",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.798715Z",
     "start_time": "2024-02-26T22:01:16.798707Z"
    }
   },
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "On a vibrant summer morning, curiosity drew me to the innovative realm of Suno AI. Their latest\n",
    "masterpiece, Chirp, known for its unique ability to generate music from words, promised a bold\n",
    "step forward in music creation. I dived into their platform, typing in familiar lyrics. As I clicked\n",
    "'Create,' the ensuing symphony seamlessly blended human emotion and machine precision,\n",
    "heralding a new era in the convergence of music, words, and machine learning\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b5a2270",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.799403Z",
     "start_time": "2024-02-26T22:01:16.799395Z"
    }
   },
   "outputs": [],
   "source": [
    "out_audio = generate_audio(\n",
    "    GenerationConfig(\n",
    "        text=text,\n",
    "        text_tags=\"hifi rap\",\n",
    "        n_batch=2,\n",
    "        max_gen_duration_s=120,\n",
    "        # stream=False,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b5656bb",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.800159Z",
     "start_time": "2024-02-26T22:01:16.800151Z"
    }
   },
   "outputs": [],
   "source": [
    "for e in out_audio:\n",
    "    e.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84046437",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.800830Z",
     "start_time": "2024-02-26T22:01:16.800822Z"
    }
   },
   "outputs": [],
   "source": [
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f167a97",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.801501Z",
     "start_time": "2024-02-26T22:01:16.801493Z"
    }
   },
   "outputs": [],
   "source": [
    "random.randint(0, 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56cdc374",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-26T22:01:16.802038Z",
     "start_time": "2024-02-26T22:01:16.802030Z"
    }
   },
   "outputs": [],
   "source": [
    "50176 / 8 / 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a6b6e5b",
   "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.13"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
