{ "cells": [ { "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", ")\n", "from suno_utils.gpt.engine import Engine\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", "from suno_utils.diffusion import generation as diffusion_gen\n", "\n", "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"2\"\n", "\n", "torch._logging.set_logs(recompiles=True, recompiles_verbose=True) # , guards=True)" ] }, { "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 = 1\n", "MAX_STREAMS = 5\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=\"s3://suno-data/georg/models/semantic/mert_25.pt\",\n", " centroids_filepath=\"s3://suno-data/georg/models/semantic/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(\n", " \"/app/suno/checkpoints/2024-12-08_06-09-46/last_ckpt_infer.pt\"\n", ")\n", "engine = Engine(\n", " gpt_ckpt_path,\n", " \"s3://suno-data/georg/trained_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": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "audio = load_audio(\"/home/sara/samples/v4_echoes_orbit.mp3\").get_segment(\n", " from_s=0, to_s=10\n", ")\n", "audio.play()\n", "artist_arr = process_audio(audio, cfg)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "prompt_text = \"\"\"\n", "[Verse]\\nSick of intensifying dark\\nMinimum wage and zero awards\\nI got manipulated\\nLed here misdirected\\n\n", "All my cash wasted\\nI hate being poor\\n\\n[Verse]\\nCause I was drivin' to my night shift\\nIn a Hyundai accent\\n\n", "On a road that overlooks the sea\\nI pulled over to the breakdown lane\\nThere was something I could smell\\n\\n\n", "[Pre-Chorus]\\nA comfort\\nA fragrance rare\\nThere was no way I knew that I could get anywhere\\n\\n[Chorus]\\nIn the moonlit care\n", "\\nIn the dream your are\\nWhen it's blue out there\\nWhen it's blue up there\\nWhen it's cold around\\nI can shut my eyes\\nI can \n", "get you there\\nThere's a world out there\\n\\n[Verse]\\nSick of intensifying low\\n'At home' was never quite where that feeling \n", "comes\\nI chase fleeting sparks that\\nKeep me tethered running\\nDown the road I'm jumping\\nTill I light one up\\n\\n[Verse]\\nAnd \n", "now we're flying through the void\\nIn a ship I've voided\\nBut at least it always knows the pace\\nIt brings me right back to \n", "the place I know so well\n", "\"\"\"\n", "\n", "style_tags = (\n", " \"stripped back ambient minimal techno, groovy bass, modular synthesizers, reverb\"\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from suno_utils.gpt.generation import CfgGenerationConfig\n", "from suno_utils.gpt.engine import GenerationConfig\n", "from suno_utils.gpt.generation_engine import align_codes, make_request\n", "from suno_utils.gpt.prompt import Prompt\n", "\n", "gconf = GenerationConfig(\n", " text=\"I can sing, welcome to Suno\",\n", " text_tags=\"edm\",\n", " text_neg_tags=\"noise\",\n", " cfg_coef=1.5,\n", " cfg_coef_max_steps=100,\n", " cfg_coef_tags=0.0,\n", " cfg_coef_neg_tags=0.0,\n", " cfg_coef_tags_max_steps=200,\n", " n_batch=1,\n", ")\n", "\n", "request = make_request(\"0\", gconf, cfg, tokenizer)\n", "in_arr = request.streams[0].prompt\n", "for stream in request.streams:\n", " print(stream, stream.weight, stream.max_n_steps)\n", "\n", "prompt = Prompt(gconf, cfg, tokenizer)" ] }, { "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", "\n", "gpt_outputs = []\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", " ],\n", " axis=0,\n", " )\n", " gpt_outputs.append(v.copy())\n", " audio = codec_decode(v)\n", " audio.play()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from suno_utils.tasks.upsample_engine import (\n", " UpsampleEngine,\n", " Request,\n", " DiffusionGenerationConfig,\n", ")\n", "\n", "diffusion_gen.preload_models(compile=False)\n", "diff_engine = UpsampleEngine(min_chunk_size=25 * 30)\n", "diffusion_generation_config = DiffusionGenerationConfig(\n", " lyrics=prompt_text, tags=style_tags, steps=5, seed=42\n", ")\n", "\n", "for x in gpt_outputs:\n", " audio_semantic_codes = torch.from_numpy(x[:, 0]).reshape(1, -1).long()\n", " request = Request(\n", " id=\"0\",\n", " generation_config=diffusion_generation_config,\n", " tokens=[c for c in audio_semantic_codes[0]],\n", " input_tokens_finished=True,\n", " )\n", " job = diff_engine.run_request(request, tqdm_enabled=True)\n", " full_audio = Audio.concatenate(job.generated_audios)\n", " full_audio.play()" ] } ], "metadata": { "kernelspec": { "display_name": "suno_clean", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.15" } }, "nbformat": 4, "nbformat_minor": 2 }