{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Time-Dependent CFG Infill Example\n", "\n", "Runs on `sa/cfg` branch of glockenspiel. New features:\n", "* Adds additional CFG streams (future, artist, cover, lyrics)\n", "* ability to change weights of a stream over time\n", "* linear interpolation between weight endpoints" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "\n", "import torch\n", "import numpy as np\n", "from suno_utils.audio import Audio\n", "from suno_utils.gpt.chirp_v2_5 import (\n", " _get_model_if_needed,\n", " GenerationConfig,\n", ")\n", "from suno_utils.gpt.generation import CfgGenerationConfig\n", "from suno_utils.gpt.generation_engine import (\n", " align_codes,\n", " make_request,\n", ")\n", "from suno_utils.gpt.engine import Engine, Request, Job\n", "from suno_utils.gpt.prompt import Prompt\n", "from suno_utils.tasks.dac_2c_12cb import (\n", " preload_models as preload_codec_models,\n", " encode as codec_encode,\n", " decode as codec_decode,\n", ")\n", "from suno_utils.tasks.mert_25 import (\n", " preload_models as preload_semantic_models,\n", " encode as semantic_encode,\n", ")\n", "\n", "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"5\"" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Setup Engine" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def _interleave(semantic_arr, n_factor=1):\n", " new_semantic_arr = (\n", " np.zeros(\n", " (semantic_arr.shape[0] * n_factor, semantic_arr.shape[-1]),\n", " dtype=semantic_arr.dtype,\n", " )\n", " + cfg.semantic_vocab_size\n", " - 1\n", " )\n", " new_semantic_arr[::n_factor] = semantic_arr\n", " return new_semantic_arr\n", "\n", "\n", "def process_audio(audio, cfg, n_factor=1):\n", " audio = audio.normalize_volume(-16)\n", " sem_arr = semantic_encode(audio, device=\"cpu\")\n", " if n_factor > 1:\n", " sem_arr = _interleave(sem_arr, n_factor=n_factor)\n", " coarse_arr = codec_encode(audio)\n", " n_frames = min(sem_arr.shape[0], coarse_arr.shape[0])\n", " sem_arr = sem_arr[:n_frames, : cfg.semantic_n_codebooks]\n", " coarse_arr = coarse_arr[:n_frames, : cfg.coarse_n_codebooks]\n", "\n", " a_arr = np.concatenate([sem_arr, coarse_arr], axis=-1)\n", " return a_arr\n", "\n", "\n", "def load_audio(fp):\n", " return Audio.from_file(fp, n_channels=2, sample_rate=48_000, byte_width=2)\n", "\n", "\n", "N_BATCH = 2\n", "MAX_STREAMS = N_BATCH * 4\n", "\n", "# preload codec\n", "_ = preload_codec_models(\"/app/suno/models/chirp_v2/dac_2c_25x12.pt\")\n", "\n", "# preload mert\n", "_ = preload_semantic_models(\n", " checkpoint_filepath=\"/app/suno/models/chirp_v2/mert_25.pt\",\n", " centroids_filepath=\"/app/suno/models/chirp_v2/mert_25_2x4k.npy\",\n", " device=\"cpu\",\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "USE_COMPILE = False\n", "gpt_ckpt_path = _get_model_if_needed(\"/app/suno/data/dpo/models/model_30b_ft_t3.pt\")\n", "engine = Engine(\n", " gpt_ckpt_path,\n", " \"/app/suno/models/chirp_v2/tokenizer_60k.json\",\n", " max_sequences=MAX_STREAMS,\n", " compile=USE_COMPILE,\n", ")\n", "model = engine.model\n", "cfg = model.config\n", "tokenizer = engine.tokenizer" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Load Audio for Testing" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "audio = load_audio(\"../../samples/synthwave_fantasy.mp3\").get_segment(\n", " from_s=45, to_s=105\n", ")\n", "audio.play()\n", "in_infill_arr = process_audio(audio, cfg)\n", "in_history_arr = in_infill_arr[: 25 * 15, :].copy()\n", "in_future_arr = in_infill_arr[-25 * 15 :, :].copy()\n", "in_infill_arr = in_infill_arr[25 * 15 : (-25 * 15), :].copy()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Specify CFG Parameters Over Time\n", "\n", "Weight and max_steps must be the same length, they form tuples of endpoints. During inference we linearly interpolate the weight based on the next target. \n", "\n", "CFG streams need to be set in the new `cfg_streams` section of the config, this branch removes the default tag/neg_tag streams (for now). You can still pass a single `weight` and `max_steps`, or set `max_steps` to None. There are new CFG stream options: \"artist\", \"cover\", \"future\", \"lyrics.\" I've mostly tested with \"future\", using it to add more variety to an infill.\n", "\n", "I've found its best to start and end the CFG at 0, and leave it at the desired value for a section in the middle. In this example I didn't like the drum pattern in the original, and thought the section was too similar to the beginning of the song. The tag CFG stream adds in more drums, and by setting the future CFG to a negative value the infilled output differs more from the original. " ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "gconf = GenerationConfig(\n", " text=\"[instrumental]\",\n", " text_tags=\"full drum kit, electronic, synthwave\",\n", " history_arr=in_history_arr,\n", " future_arr=in_future_arr,\n", " n_batch=1,\n", " max_gen_duration_s=240,\n", " min_text_offset=0,\n", " eos_pad_duration_s=0,\n", " min_eos_p=0.05,\n", " text_start_control_tags=\"{start:15;duration:60}\",\n", " cfg_streams=[\n", " CfgGenerationConfig(\n", " stream_type=\"tag\", weight=[0.0, 2.5, 2.5, 0.0], max_steps=[0, 100, 600, 700]\n", " ),\n", " CfgGenerationConfig(\n", " stream_type=\"future\",\n", " weight=[0.0, -3.5, -3.5, 0.0],\n", " max_steps=[0, 100, 600, 700],\n", " ),\n", " ],\n", ")\n", "\n", "request = make_request(\"0\", gconf, cfg, tokenizer)\n", "in_arr = request.streams[0].prompt\n", "\n", "prompt = Prompt(gconf, cfg, tokenizer)\n", "\n", "assert N_BATCH * 4 <= 16" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Generate Infill\n", "\n", "(this only plays the infilled section, uncomment the if statement to get the history too)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "requests = [\n", " make_request(f\"{i}\", gconf.modify(n_batch=1), engine.model.config, engine.tokenizer)\n", " for i in range(N_BATCH)\n", "]\n", "for job in engine.run_request(requests, tqdm_enabled=True):\n", " stream = engine.token_generator(job)\n", " v = np.concatenate(\n", " [\n", " torch.stack(list(align_codes(stream, cfg))).detach().cpu().numpy()[:, -12:],\n", " in_future_arr[:, -12:],\n", " ],\n", " axis=0,\n", " )\n", " # if in_history_arr is not None:\n", " # v = np.concatenate(\n", " # [\n", " # in_history_arr[:, -12:],\n", " # v,\n", " # ],\n", " # axis=0,\n", " # )\n", " audio = codec_decode(v)\n", " audio.play()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### For Easy Comparison to the Original Section" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "audio = load_audio(\"../../samples/synthwave_fantasy.mp3\").get_segment(\n", " from_s=60, to_s=105\n", ")\n", "audio.play()" ] } ], "metadata": { "kernelspec": { "display_name": "suno_env", "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.14" } }, "nbformat": 4, "nbformat_minor": 2 }