import os import sys import time import modal import torch import pathlib import funcy import json import pandas as pd import numpy as np import tempfile from suno_utils.audio import Audio from suno_utils.gpt.generation import GenerationConfig from suno_utils.gpt.engine import Engine from suno_utils.gpt.generation_engine import make_request from suno_utils.worker.settings import s3_client from suno_utils.worker.modal_base import MODAL_MOUNTS from suno_utils.gpt import chirp_v2_5 as chirp_v2 from suno_utils.utils.text import read_jsonl from suno_utils.utils.s3 import list_s3_dir from suno_utils.diffusion.generation import ( preload_dit_model, preload_tokenizer, TOKENIZER_FILEPATH, SEMANTIC_MODEL_FILEPATH, SEMANTIC_CLUSTERS_FILEPATH, _retrieve_models, ) from suno_utils.tasks.mert_25 import ( preload_models as preload_semantic_models, encode as encode_semantic, ) from suno_utils.tasks.dac_vae_fixed_25hz import ( preload_models as preload_codec_models, decode as codec_decode, ) from suno_utils.diffusion import generation as diffusion_gen from suno_utils.tasks.upsample_engine import UpsampleEngine, Request, Job from suno_utils.tasks.dac_vae_fixed_25hz import decode_stream_to_full_audio MOUNT_PATH = "/suno/models" aws_secret = modal.Secret.from_name("studio-aws") SECRETS = [ aws_secret, modal.Secret.from_dict( { "SUNO_ASSETS_PATH": "/suno/models/assets", "XDG_CACHE_HOME": "/suno/models/", } ), modal.Secret.from_name("api-callback-token"), modal.Secret.from_name("datadog-metrics"), ] base_image = ( modal.Image.from_registry("nvidia/cuda:12.4.0-devel-ubuntu22.04", add_python="3.10") .apt_install( "curl", "ffmpeg", "sox", "unzip", "libsox-fmt-mp3", "zlib1g-dev", "git", "clang" ) .run_commands( [ 'curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip"', "unzip -q awscliv2.zip", "./aws/install", ] ) .dockerfile_commands( [ "COPY --from=datadog/serverless-init:1.2.1 /datadog-init /app/datadog-init", 'ENTRYPOINT ["/app/datadog-init"]', ] ) .pip_install("torch==2.5.1", "torchaudio==2.5.1") .pip_install( "flashinfer-python", index_url="https://flashinfer.ai/whl/cu124/torch2.5/" ) .pip_install_private_repos( "github.com/suno-ai/glockenspiel.git@a5dba4e50#subdirectory=descript-audio-codec&egg=descript-audio-codec", git_user="mcamac", secrets=[modal.Secret.from_name("victor-modal-github-token")], ) .pip_install_private_repos( "github.com/suno-ai/neon.git@1c83548#subdirectory=hoot", git_user="mcamac", secrets=[modal.Secret.from_name("victor-modal-github-token")], ) .pip_install( "boto3", "transformers", "tokenizers", "encodec", "ctc_segmentation", "psutil", "redis", "pydantic", "nnAudio", "rpyc", "biopython>=1.81", # TODO: don't love this depdendency, for hoot "pynvml", # for torch cuda utilization "torchsde", "ninja", "wheel", ) .pip_install_from_pyproject( "/home/christian/code/glockenspiel/suno_utils/pyproject.toml", ) .run_commands( # This is really slow "git clone https://github.com/Dao-AILab/flash-attention.git", "cd flash-attention/hopper && python setup.py install", gpu="h100", ) .apt_install( "libogg0", "libopus0", "opus-tools", ) .pip_install("transformers==4.44.0", "wandb") ) class GenerateWorker: def __init__(self, gpt_ckpt: str, dit_model_filepath: str, output_path: str): self.output_path = output_path start_time = time.time() print("Start loading models") tokenizer_filepath = "s3://suno-data/georg/models/tokenizers/tokenizer_60k.json" gpt_ckpt_path = chirp_v2._get_model_if_needed(gpt_ckpt, cache_dir=MOUNT_PATH) tokenizer_path = chirp_v2._get_model_if_needed( tokenizer_filepath, cache_dir=MOUNT_PATH ) print(tokenizer_path) N_BATCH = 2 engine = Engine( gpt_ckpt_path, tokenizer_path=tokenizer_path, max_sequences=4 * N_BATCH, compile=False, ) cfg = engine.model.config self.engine = engine self.cfg = cfg # load gpt model num_gpus = torch.cuda.device_count() cuda_device = torch.cuda.current_device() print(f"Found {num_gpus} GPUs. Using GPU {cuda_device}.") tokenizer_filepath = "s3://suno-data/georg/models/tokenizers/tokenizer_60k.json" semantic_model_filepath = "s3://suno-data/georg/models/semantic/mert_25.pt" semantic_clusters_filepath = ( "s3://suno-data/georg/models/semantic/mert_25_2x4k.npy" ) codec_filepath = "s3://suno-data/minz/models/dac_vae_tuned_25hz.pth" _ = diffusion_gen.preload_dit_model( dit_model_filepath=dit_model_filepath, use_ema_if_exists=True, compile=False, weights_precision=torch.bfloat16, ) _ = preload_tokenizer(tokenizer_filepath) _ = preload_semantic_models(semantic_model_filepath, semantic_clusters_filepath) _ = preload_codec_models(codec_filepath) self.diffusion_engine = UpsampleEngine(min_chunk_size=25 * 30) print( f"Finish loading models. Took {round(time.time() - start_time, 2)} seconds" ) @staticmethod def download_models(gpt_ckpt, dit_model_filepath, dir_path=MOUNT_PATH): """Download diffusion models.""" print("Start downloading models") _ = diffusion_gen.get_model_if_needed( codec_filepath="s3://suno-data/minz/models/dac_vae_tuned_25hz.pth", cache_dir=dir_path, ) _ = diffusion_gen.get_model_if_needed( diffusion_gen.SEMANTIC_MODEL_FILEPATH, cache_dir=dir_path ) _ = diffusion_gen.get_model_if_needed( diffusion_gen.SEMANTIC_CLUSTERS_FILEPATH, cache_dir=dir_path ) _ = diffusion_gen.get_model_if_needed( diffusion_gen.CODEC_FILEPATH, cache_dir=dir_path ) _ = diffusion_gen.get_model_if_needed(dit_model_filepath, cache_dir=dir_path) _ = chirp_v2._get_model_if_needed(gpt_ckpt, cache_dir=dir_path) _ = chirp_v2._get_model_if_needed(chirp_v2.TOKENIZER_PATH, cache_dir=dir_path) print("Finish downloading models") def generate(self, work_item): global models """Generate audio from a work item.""" # when saving to s3 we will use a structure like this: # self.output_path/ # item_id/ # original_semantic.npz # 0_generated_audio.mp3 # 0_generated_semantic.npz # 0_metadata.json # 1_generated_audio.mp3 # 1_generated_semantic.npz # 1_metadata.json # ... # metadata.json for item in work_item: item_id = item["id"] # start by semantic encoding the original audio audio = Audio.from_s3(item["s3_filepath"], n_channels=2).convert( sample_rate=24_000, byte_width=2, n_channels=1 ) semantic_codes = encode_semantic(audio)[:, 0] # take just first codebook print(semantic_codes.shape) # save to temporary directory as npz with tempfile.TemporaryDirectory() as td: npz_filepath = os.path.join(td, f"{item_id}_original_semantic.npz") np.savez(npz_filepath, semantic_codes=semantic_codes) # move to s3 s3_client.upload_file( npz_filepath, "suno-data", f"{self.output_path}/{item_id}/{item_id}_original_semantic.npz", ) # sample inference parameters n_skip_semantic = 1 cfg_coef = np.random.choice([1.0, 1.5, 2.0, 2.5, 3.0]) cfg_coef_tags = np.random.choice([1.0, 1.5, 2.0, 2.5, 3.0]) n_repeat_tags = np.random.randint(1, 3) temp_semantic = np.random.uniform(0.8, 1.0) gpt_seed = np.random.randint(0, 1000000) tags_str = ", ".join(item["tags"]) # generate text gconf = GenerationConfig( text=item["text"], text_tags=tags_str, cfg_coef=cfg_coef, cfg_coef_tags=cfg_coef_tags, cfg_coef_max_steps=None, cfg_coef_tags_max_steps=None, n_repeat_tags=n_repeat_tags, n_skip_semantic=n_skip_semantic, # text_start_control_tags="{duration:210}", temp_semantic=temp_semantic, n_batch=1, min_text_offset=0, eos_pad_duration_s=0, max_gen_duration_s=int(8 * 60 / n_skip_semantic), random_seed=gpt_seed, ) requests = [ make_request( f"{i}", gconf, self.engine.model.config, self.engine.tokenizer ) for i in range(1) ] jobs = self.engine.run_request(requests, tqdm_enabled=True) out_gpt = [] for n, job in enumerate(jobs): stream = self.engine.token_generator(job) arr = torch.stack(list(stream))[:, 1] if arr[-1] == 4000: arr = arr[:-1] print(f"{round(arr.shape[-1]/25*n_skip_semantic)}s for track {n}") # do stuff incase skip arr2 = ( torch.zeros(arr.shape[0] * n_skip_semantic, dtype=arr.dtype) + self.cfg.semantic_pad_token ) arr2[::n_skip_semantic] = arr # add out_gpt.append(arr2) # diffusion parameters diffusion_seed = np.random.randint(0, 1000000) diffusion_steps = np.random.choice([8, 10, 12, 14, 16, 18, 20]) diffusion_text_cfg_coef = np.random.choice([1.0, 1.5, 1.75, 2.0, 2.5]) noise_ctx_level = np.random.choice([0.0, 0.25, 0.5, 0.75, 1.0]) # diffusion_ctx_cfg_coef = 1.0 gen_cfg = diffusion_gen.DiffusionGenerationConfig( steps=diffusion_steps, lyrics=item["text"], tags=tags_str, text_cfg_coef=diffusion_text_cfg_coef, codec_scale_factor=0.4, scale_ctx_vector=True, noise_ctx_level=noise_ctx_level, noise_ctx_pad_len=25, drop_semantic_tokens=False, seed=diffusion_seed, ) request = Request( id="dummy", generation_config=gen_cfg, tokens=out_gpt[0], input_tokens_finished=True, ) result = self.diffusion_engine.run_request(request) vae_latents = torch.concat(result.vae_latents) upsampled_audio = decode_stream_to_full_audio(vae_latents) metadata = { "original_audio": item["s3_filepath"], "text": item["text"], "tags": tags_str, "gpt": { "cfg_coef": float(cfg_coef), "cfg_coef_tags": float(cfg_coef_tags), "n_repeat_tags": int(n_repeat_tags), "n_skip_semantic": int(n_skip_semantic), "temp_semantic": float(temp_semantic), "seed": int(gpt_seed), }, "diffusion": { "steps": int(diffusion_steps), "seed": int(diffusion_seed), "text_cfg_coef": float(diffusion_text_cfg_coef), "noise_ctx_level": float(noise_ctx_level), }, } # we want to save out # audio file of the final audio # npz of the estimated semantics # metadata json with the original prompt and tags # copy some stuff # we want to copy the original audio and the npz of the original semantics with tempfile.TemporaryDirectory() as td: # save the gpt semantics gpt_semantic_path = os.path.join( td, f"{item_id}_generated_semantic.npz" ) np.savez(gpt_semantic_path, semantic_codes=out_gpt[0]) # save the metadata metadata_path = os.path.join(td, f"{item_id}_metadata.json") with open(metadata_path, "w") as f: json.dump(metadata, f) # save audio to s3 upsampled_audio_path = os.path.join(td, f"{item_id}.mp3") upsampled_audio.write_hq_mp3(upsampled_audio_path) s3_filepath = os.path.join( self.output_path, f"{item_id}", f"{item_id}.mp3" ) print(f"Uploading to {s3_filepath}") s3_client.upload_file( upsampled_audio_path, "suno-data", s3_filepath, ExtraArgs={ "ContentType": "audio/mpeg", }, ) # move the gpt semantic to s3 s3_filepath = os.path.join( self.output_path, f"{item_id}", f"{item_id}_generated_semantic.npz" ) s3_client.upload_file( gpt_semantic_path, "suno-data", s3_filepath, ExtraArgs={ "ContentType": "application/octet-stream", }, ) # move the metadata to s3 s3_filepath = os.path.join( self.output_path, f"{item_id}", f"{item_id}_metadata.json" ) s3_client.upload_file( metadata_path, "suno-data", s3_filepath, ExtraArgs={ "ContentType": "application/json", }, ) # base model 6b GPT_CKPT = ( "s3://suno-data/christian/checkpoints/gpt/2025-02-04_21-04-31-last_ckpt_infer.pt" ) # 6b v2 finetune # GPT_CKPT = ( # "s3://suno-data/christian/checkpoints/gpt/2025-02-10_16-52-41-step_9000_infer.pt" # ) # GPT_CKPT = "s3://suno-data/christian/checkpoints/gpt/2025-02-13_11-27-29-step_10_000.pt" # prod diffusion model # DIT_MODEL_FILEPATH = "s3://suno-data/tony/tmp/diff/dit_v3_dpo_t10_3k_5e6_b100_t25.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/georg/tmp/2b_prefix_ft.pt" DIT_MODEL_FILEPATH = ( "s3://suno-data/christian/checkpoints/diffusion/v45_2b_step_2_600_000.pt" ) def download_model_wrapper_d(): # this print is necessary to have modal rerun this when MODEL changes # Modal tracks referenced global variables # Change the name of the function to force a rerun print("Downloading model for history encoder") GenerateWorker.download_models(GPT_CKPT, DIT_MODEL_FILEPATH) image = base_image.run_function(download_model_wrapper_d, secrets=SECRETS) APP_NAME = f"batch-generate-gpu" app = modal.App(APP_NAME, image=image, secrets=SECRETS) N_MAX_REPLICAS = 16 @app.cls( gpu=modal.gpu.H100(count=1), cpu=4, secrets=SECRETS, timeout=2 * 60 * 60, container_idle_timeout=240, mounts=MODAL_MOUNTS, memory=15000, concurrency_limit=N_MAX_REPLICAS, ) class GenerateStub: def __init__(self, gpt_ckpt: str, dit_ckpt: str, output_path: str): import torch num_gpus = torch.cuda.device_count() print(f"Found {num_gpus} GPUs.") self.worker = GenerateWorker(gpt_ckpt, dit_ckpt, output_path) @modal.method() def generate(self, work_item: list[dict]): return self.worker.generate(work_item) @app.local_entrypoint() def main(): # checkpoints in s3://suno-data/christian/checkpoints # prompts in s3://suno-data/christian/prompts # outputs in s3://suno-data/christian/outputs output_str = "mountain-whisper-eclipse" work_items = read_jsonl( "/home/christian/code/christian/metadata/discogs_subset_sampled_metas.jsonl" ) print(f"Total work items: {len(work_items)}") # first check for existing ids in the output path existing_ids = list_s3_dir( f"s3://suno-data/christian/outputs/{output_str}/", ) existing_ids = [os.path.dirname(f[0]).split("/")[-1] for f in existing_ids] existing_ids = list(set(existing_ids)) print(f"Total existing ids: {len(existing_ids)}") print(existing_ids[:3]) # now remove these from the work items work_items = [f for f in work_items if f["id"] not in existing_ids] print(f"Total work items remaining: {len(work_items)}") chunksize = 16 # num of prompts per worker worker = GenerateStub( GPT_CKPT, DIT_MODEL_FILEPATH, f"christian/outputs/{output_str}" ) work_items = list(funcy.chunks(chunksize, work_items)) print(f"Chunksize: {chunksize}, total chunks: {len(work_items)}") # print("Testing inference...") # t0 = time.time() # for work_item in work_items[:1]: # _ = worker.generate.remote(work_item) # print(f"{int(round(time.time()-t0))}s for test") print("Running batch inference...") t0 = time.time() _ = list(worker.generate.map(work_items)) print(round((time.time() - t0) / 60 / 60), "h total for batch generation")