{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "\n", "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"6\"\n", "\n", "import json\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", "import torch\n", "import numpy as np\n", "\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", "\n", "from suno_utils.tasks.mert_25 import (\n", " preload_models as preload_semantic_models,\n", " encode as encode_semantic,\n", ")\n", "\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", "from suno_utils.diffusion import generation as diffusion_gen\n", "from suno_utils.tasks.upsample_engine import UpsampleEngine, Request, Job\n", "#from suno_utils.tasks.upsample_engine_old import UpsampleEngine, Request, Job\n", "\n", "import sys\n", "sys.path.insert(0, \"/home/m4burns/glockenspiel/suno_utils/notebooks/shimmerscore/\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# load diffusion model\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", "# dit models\n", "dit_model_filepath = \"/app2/suno/checkpoints/2025-07-14_15-27-20_s7143/last_ckpt_infer.pt\"\n", "\n", "# other models\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 = (\n", " \"s3://suno-data/georg/models/semantic/mert_25_2x4k.npy\"\n", ")\n", "codec_filepath = \"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\"\n", "CODEC_SCALE_FACTOR = 0.4\n", "SCALE_CTX_VECTOR = True\n", "\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", "_ = preload_tokenizer(tokenizer_filepath)\n", "_ = preload_semantic_models(semantic_model_filepath, semantic_clusters_filepath)\n", "_ = preload_codec_models(codec_filepath)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "diffusion_engine = UpsampleEngine(min_chunk_size=25*30) #, vae_version=\"v_vae_25_tuned_2\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# load the test clip here\n", "item_id = \"3b1830d7-5553-4785-be0b-ffeadf7756d8\"\n", "\n", "print(item_id)\n", "# load the semantic codes from s3\n", "s3_filepath = f\"s3://suno-data-uploads/studio/uploads/{item_id}.npz\"\n", "\n", "try:\n", " data = read_from_s3(s3_filepath, read_f=np.load)\n", "except Exception as e:\n", " print(f\"Error loading {s3_filepath}: {e}\")\n", "\n", "if \"v3.0_raw\" in data:\n", " codes = data[\"v3.0_raw\"]\n", "elif \"v3.5_raw\" in data:\n", " codes = data[\"v3.5_raw\"]\n", "elif \"v4.0_raw\" in data:\n", " codes = data[\"v4.0_raw\"]\n", "elif \"v5.0_raw\" in data:\n", " codes = data[\"v5.0_raw\"]\n", "else:\n", " raise ValueError(\"No codes found\")\n", "\n", "semantic_codes = torch.from_numpy(codes[:, 0]).long() # .cuda()\n", "# semantic_codes = semantic_codes[:3000]\n", "print(semantic_codes.shape)\n", "\n", "# load the vae\n", "vae_data_filepath = f\"s3://suno-data-uploads/studio/uploads/{item_id}_vae.npz\"\n", "vae_data = read_from_s3(vae_data_filepath, read_f=np.load)\n", "print(vae_data.keys())\n", "init_vae_latents = vae_data[\"vae_latents\"]\n", "print(init_vae_latents.shape)\n", "\n", "gen_seed = vae_data.get(\"seed\", None)\n", "print(gen_seed)\n", "\n", "text_data = read_from_s3(f\"s3://suno-data-uploads/studio/uploads/{item_id}_hoot.json\")\n", "aligned_lyrics = json.loads(text_data)\n", "# print(aligned_lyrics)\n", "\n", "tags = \"80s, new wave, female power\"\n", "\n", "lyrics = \"\"\n", "for elem in aligned_lyrics:\n", " if \"word\" in elem:\n", " lyrics += elem[\"word\"]\n", "# print(lyrics)\n", "mp3_filepath = f\"s3://suno-data-uploads/studio/uploads/{item_id}.mp3\"\n", "data = read_from_s3(s3_filepath, read_f=np.load)\n", "audio = Audio.from_s3(mp3_filepath, n_channels=2)\n", "audio.play()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "print()\n", "# Map slider value from 0–1 to a target range\n", "def map_slider(slider_value, min_out, max_out):\n", " \"\"\"Map slider from 0-1 range to min_out-max_out range, handling None.\"\"\"\n", " if slider_value is None:\n", " return None\n", " slider_value = max(0, min(1, slider_value)) # clamp to 0–1\n", " return min_out + (max_out - min_out) * slider_value\n", "\n", "# === Set global slider values from frontend input (0.0 to 1.0) ===\n", "slider_strength = 0.95 # example: strength from user input\n", "slider_quality = 0.5 # example: quality\n", "slider_tone = 0.5 # example: tone\n", "slider_stereo = 0.5 # example: stereo\n", "\n", "# === Map to model ranges ===\n", "strength_value = map_slider(slider_strength, 1.0, 4.0) # cfg\n", "quality_value = map_slider(slider_quality, 12, 28) # ear score\n", "tone_value = map_slider(slider_tone, -2.0, 2.0) # spectral centroid\n", "stereo_value = map_slider(slider_stereo, -2.0, 2.0) # stereo balance\n", "\n", "# === Format audio tags ===\n", "audio_tags_parts = []\n", "if quality_value is not None and quality_value != 20.0:\n", " audio_tags_parts.append(f\"quality: {int(quality_value)}\")\n", "if tone_value is not None and tone_value != 0.0:\n", " audio_tags_parts.append(f\"spectral_centroid: {tone_value:.1f}\")\n", "if stereo_value is not None and stereo_value != 0.0:\n", " #if stereo_value <= -1.5:\n", " # audio_tags_parts.append(f\"mono\")\n", " #else:\n", " audio_tags_parts.append(f\"stereo_width: {stereo_value:.1f}\")\n", "\n", "audio_tags = \", \".join(audio_tags_parts) if audio_tags_parts else \"\"\n", "print(f\"Prepared audio tags: {audio_tags}\")\n", "\n", "# === Example of using values in model logic ===\n", "text_cfg_coef = 2.5 if strength_value is None else strength_value\n", "\n", "steps = 10\n", "text_cfg_coef = 2.0\n", "noise_ctx_level = 0.75\n", "noise_ctx_pad_len = 0\n", "codec_scale_factor = 0.4\n", "num_seeds = 2\n", "\n", "seeds = list(range(3))\n", "upsampled_audios = []\n", "for idx, seed in enumerate(seeds):\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=True,\n", " noise_ctx_level=noise_ctx_level,\n", " noise_ctx_pad_len=noise_ctx_pad_len,\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=semantic_codes,\n", " input_tokens_finished=True,\n", " )\n", "\n", " result = diffusion_engine.run_request(request)\n", "\n", " vae_latents = []\n", " for vae_latent in result.vae_latents:\n", " #print(vae_latent.shape)\n", " mean = vae_latent.mean()\n", " std = vae_latent.std()\n", " vae_latents.append(vae_latent)\n", "\n", " vae_latents = torch.concat(vae_latents)\n", " upsampled_audio = decode_stream_to_full_audio(vae_latents)\n", " upsampled_audio.play()\n", " upsampled_audios.append(upsampled_audio)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "gen_id = \"ae137e1f-7562-48f6-a11c-b125e645282c\"\n", "gen_id_audio = \"cb105e06-6e20-4acc-9d1c-e7c9b36b24e2\"\n", "\n", "s3_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}.npz\"\n", "mp3_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id_audio}.mp3\"\n", "print(s3_filepath)\n", "data = read_from_s3(s3_filepath, read_f=np.load)\n", "audio = Audio.from_s3(mp3_filepath, n_channels=2)\n", "\n", "if \"v3.0_raw\" in data:\n", " codes = data[\"v3.0_raw\"]\n", "elif \"v3.5_raw\" in data:\n", " codes = data[\"v3.5_raw\"]\n", "elif \"v4.0_raw\" in data:\n", " codes = data[\"v4.0_raw\"]\n", "else:\n", " raise ValueError(\"No codes found\")\n", "\n", "semantic_codes = torch.from_numpy(codes[:, 0]).long() # .cuda()\n", "# semantic_codes = semantic_codes[:3000]\n", "print(semantic_codes.shape)\n", "\n", "text_data = read_from_s3(f\"s3://suno-data-uploads/studio/uploads/{gen_id}_hoot.json\")\n", "aligned_lyrics = json.loads(text_data)\n", "# print(aligned_lyrics)\n", "# lyrics = \"\".join(w[\"word\"] for w in aligned_lyrics if \"word\" in w)\n", "lyrics = \"walking down the streets feeling so alive i've got my head in the clouds got a gleam in my eye every step i take it's like a brand new start no matter where i'm going i'll always find my part life is like a hard wire act we're dancing in the sky no need to worry no need to ask why with a little bit of courage we can chase our dreams no matter what comes our way we'll always be a team we're unstoppable yeah\"\n", "tags = 'orchestral, orchestra, classical, epic movie'\n", "audio.normalize_volume().play()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "upsampled_audios = []\n", "for idx, seed in enumerate(seeds):\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=True,\n", " noise_ctx_level=noise_ctx_level,\n", " noise_ctx_pad_len=noise_ctx_pad_len,\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=semantic_codes,\n", " input_tokens_finished=True,\n", " )\n", "\n", " result = diffusion_engine.run_request(request)\n", "\n", " vae_latents = []\n", " for vae_latent in result.vae_latents:\n", " #print(vae_latent.shape)\n", " mean = vae_latent.mean()\n", " std = vae_latent.std()\n", " vae_latents.append(vae_latent)\n", "\n", " vae_latents = torch.concat(vae_latents)\n", " upsampled_audio = decode_stream_to_full_audio(vae_latents)\n", " upsampled_audio.play()\n", " upsampled_audios.append(upsampled_audio)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "BREAK" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Infill" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "list(vae_data.keys())\n", "ctx_vae_latents = vae_data[\"vae_latents\"]\n", "ctx_vae_latents.shape\n", "print(ctx_vae_latents.shape)\n", "decode_stream_to_full_audio(ctx_vae_latents).play()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "lyrics_30s = \"\"\"\n", "[Instrumental]\n", "\n", "\"\"\"" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "SEMANTIC_RATE_HZ = 25\n", "\n", "start_s = 10\n", "end_s = 20\n", "infill_dur_s = end_s - start_s\n", "context_dur_s = 30.0 - infill_dur_s\n", "\n", "# count total tokens of the context + infill\n", "infill_tokens = int(np.ceil(infill_dur_s * SEMANTIC_RATE_HZ))\n", "print(infill_dur_s)\n", "print(infill_tokens)\n", "\n", "prefix_start_s = start_s - (context_dur_s // 2)\n", "prefix_end_s = start_s\n", "suffix_start_s = end_s\n", "suffix_end_s = end_s + (context_dur_s // 2)\n", "print(\"prefix_start_s\", prefix_start_s)\n", "print(\"suffix_end_s\", suffix_end_s)\n", "\n", "prefix_start_idx = int(np.ceil(prefix_start_s * SEMANTIC_RATE_HZ))\n", "prefix_end_idx = int(np.ceil(prefix_end_s * SEMANTIC_RATE_HZ))\n", "prefix_tokens = prefix_end_idx - prefix_start_idx\n", "\n", "print(\"prefix_start_idx\", prefix_start_idx)\n", "print(\"prefix_end_idx\", prefix_end_idx)\n", "\n", "suffix_start_idx = int(np.ceil(suffix_start_s * SEMANTIC_RATE_HZ))\n", "suffix_end_idx = int(np.ceil(suffix_end_s * SEMANTIC_RATE_HZ))\n", "suffix_tokens = suffix_end_idx - suffix_start_idx\n", "\n", "print(\"suffix_start_idx\", suffix_start_idx)\n", "print(\"suffix_end_idx\", suffix_end_idx)\n", "\n", "total_context_tokens = prefix_tokens + suffix_tokens\n", "print(\"total_context_tokens\", total_context_tokens)\n", "\n", "total_tokens = total_context_tokens + infill_tokens\n", "if total_tokens < 750: \n", " # add more context to the right side\n", " context_to_add = 750 - total_tokens\n", " print(context_to_add)\n", " suffix_end_idx += context_to_add\n", "\n", "\n", "infill_prefix_latents = torch.tensor(ctx_vae_latents[prefix_start_idx: prefix_end_idx, :])\n", "infill_suffix_latents = torch.tensor(ctx_vae_latents[suffix_start_idx: suffix_end_idx, :])\n", "print(infill_prefix_latents.shape)\n", "print(infill_suffix_latents.shape)\n", "\n", "\n", "gen_cfg = diffusion_gen.DiffusionGenerationConfig(\n", " audio=None,\n", " steps=16,\n", " lyrics=lyrics_30s,\n", " tags=tags,\n", " text_cfg_coef=2.0,\n", " ctx_cfg_coef=1.0,\n", " codec_scale_factor=0.4,\n", " scale_ctx_vector=True,\n", " infill_prefix_latents=infill_prefix_latents,\n", " infill_suffix_latents=infill_suffix_latents,\n", " noise_ctx_level=0.0,\n", " drop_semantic_tokens=False,\n", " semantic_skip_factor=1,\n", " semantic_mask_ratio=1.0,\n", " is_diff_infill=True,\n", " #objective=\"rectified_flow\",\n", ")\n", "input_tokens = codes[prefix_start_idx : suffix_end_idx, 0]\n", "print(input_tokens.shape)\n", "request = Request(\n", " id=\"dummy\",\n", " generation_config=gen_cfg,\n", " tokens=input_tokens,\n", " input_tokens_finished=True,\n", ")\n", "result = diffusion_engine.run_request(request)\n", "\n", "vae_latents = torch.concat(result.vae_latents)\n", "print(vae_latents.shape)\n", "decode_stream_to_full_audio(vae_latents).play()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "decode_stream_to_full_audio(ctx_vae_latents).play()" ] } ], "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": 4 }