import os import time import modal import torch import pathlib import funcy import pandas as pd import numpy as np import tempfile import uuid import json import torchaudio from suno_utils.audio import Audio from suno_utils.utils.s3 import read_from_s3 from suno_utils.utils.text import read_jsonl from suno_utils.worker.settings import s3_client from suno_utils.worker.modal_base import MODAL_MOUNTS from suno_utils.utils.s3 import read_from_s3, upload_s3_files, list_s3_dir from suno_utils.diffusion.generation import ( generate_offline, preload_models, get_model_if_needed, 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, EMBEDDING_RATE as SEMANTIC_HZ, ) from suno_utils.tasks.dac_vae_fixed_25hz import ( preload_models as preload_dac_vae_fixed_25hz, encode as codec_encode, decode as codec_decode, ) 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.debian_slim() .apt_install( "curl", "ffmpeg", "sox", "unzip", "libsox-fmt-mp3", "zlib1g-dev", "git" ) .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.4.0") .pip_install( "boto3", "transformers", "tokenizers", "encodec", "ctc_segmentation", "psutil", "redis", "gradio", "pydantic", "nnAudio", "rpyc", "biopython>=1.81", # TODO: don't love this depdendency, for hoot "pynvml", # for torch cuda utilization "torchsde", "funcy", "fasttext", "wandb", ) .pip_install_from_pyproject( str( "/home/christian/code/glockenspiel/suno_utils/pyproject.toml", ), # force_build=True, ) .pip_install("torch==2.4.0") # this is cause flash-attn can't work with 2.5 yet .run_commands( "FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE pip install flash-attn==2.6.3 --no-build-isolation", ) .pip_install( "torch==2.5.1", "torchaudio==2.5.1", "transformers==4.44.0", "numpy==1.26.4", ) ) class GenerateWorker: def __init__( self, dit_model_filepath: str, codec_model_filepath: str, output_path: str, ): self.output_path = output_path start_time = time.time() print("Start loading models") # 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}.") if ( "dac_vae_fixed" in codec_model_filepath or "dac_vae_tuned" in codec_model_filepath ): self.patch_size = 1 self.codec_scale_factor = 0.4 self.downscale_ctx_vector = False elif "convnext_vae_tuned" in codec_model_filepath: self.patch_size = 1 self.codec_scale_factor = 1.0 self.downscale_ctx_vector = False else: self.patch_size = 1 self.codec_scale_factor = 2.5 self.downscale_ctx_vector = True _ = preload_models( tokenizer_filepath=TOKENIZER_FILEPATH, semantic_model_filepath=SEMANTIC_MODEL_FILEPATH, semantic_clusters_filepath=SEMANTIC_CLUSTERS_FILEPATH, codec_filepath=codec_model_filepath, dit_model_filepath=dit_model_filepath, weights_precision=torch.bfloat16, model_type="prefix", codec_scale_factor=self.codec_scale_factor, ) print( f"Finish loading models. Took {round(time.time() - start_time, 2)} seconds" ) @staticmethod def download_models(dit_model_filepath, codec_model_filepath, dir_path=MOUNT_PATH): """Download diffusion models.""" print("Start downloading models") _ = get_model_if_needed(TOKENIZER_FILEPATH, cache_dir=dir_path) _ = get_model_if_needed(SEMANTIC_MODEL_FILEPATH, cache_dir=dir_path) _ = get_model_if_needed(SEMANTIC_CLUSTERS_FILEPATH, cache_dir=dir_path) _ = get_model_if_needed(codec_model_filepath, cache_dir=dir_path) _ = get_model_if_needed(dit_model_filepath, cache_dir=dir_path) print("Finish downloading models") def generate(self, work_items): print(f"Generating {len(work_items)} work items") for work_item in work_items: print(work_item) # get the id gen_id = work_item["id"] # check if the metadata exists on s3 # skip anything if the duration is longer than 240s if work_item["duration_s"] > 480: print(f"Skipping {gen_id} because it is longer than 480s") continue # check if the metadata.json file exists on s3 metadata_path = os.path.join(self.output_path, f"{gen_id}_metadata.json") if "s3_filepath" in work_item: try: # we will download the mp3 file from s3 and then semantic encode it s3_filepath = work_item["s3_filepath"] # s3_filepath is a full path, so we need remove the s3://suno-data/ prefix s3_filepath = s3_filepath.replace("s3://suno-data/", "") audio_path = os.path.join(tempfile.gettempdir(), f"{gen_id}.mp3") s3_client.download_file("suno-data", s3_filepath, audio_path) # load audio audio = Audio.from_file(audio_path) # also encode semantic codes semantic_codes = encode_semantic( audio.convert(sample_rate=24_000, byte_width=2, n_channels=1) ).astype(np.int64)[:, 0] # save semantic codes to s3 and npz file local_path = os.path.join(tempfile.gettempdir(), f"{gen_id}.npz") s3_output_path = os.path.join(self.output_path, f"{gen_id}.npz") np.savez(local_path, semantic_codes=semantic_codes) s3_client.upload_file( local_path, "suno-data", s3_output_path, ExtraArgs={ "ContentType": "application/octet-stream", }, ) except Exception as e: print(f"Error encoding audio: {e}") continue else: try: s3_filepath = f"s3://suno-data-uploads/studio/uploads/{gen_id}.npz" data = read_from_s3(s3_filepath, read_f=np.load) if "v3.0_raw" in data: codes = data["v3.0_raw"] elif "v3.5_raw" in data: codes = data["v3.5_raw"] elif "v4.0_raw" in data: codes = data["v4.0_raw"] else: raise ValueError("No codes found") semantic_codes = torch.from_numpy(codes[:, 0]).long() # .cuda() except Exception as e: print(f"Error loading semantic codes: {e}") continue # perform upsampling with multiple seeds # we want to save out # 3. mp3 audio file try: with tempfile.TemporaryDirectory() as td: base_filename = f"{work_item['id']}" for n in range(1): # create a uuid for this example # upsample_uuid = str(uuid.uuid4()) steps = 16 seed = np.random.randint(0, 1000000) text_cfg_coef = np.random.choice( [1.0, 1.25, 1.5, 2.0, 2.5, 3.0] ) noise_ctx_vector = np.random.choice([0.0, 0.1, 0.2, 0.5]) # check for valid inputs if work_item["text"] is None: text = "" else: text = work_item["text"] if work_item["tags"] is None: tags = "" else: tags = ", ".join(work_item["tags"]) # first codec cycle the audio audio = Audio.from_file( audio_path, n_channels=2 ) # reload audio encoded_audio = codec_encode( audio.convert( sample_rate=48_000, byte_width=2, n_channels=2 ) ) decoded_audio = codec_decode(encoded_audio) # save audio to disk then upload to s3 audio_path = os.path.join(td, f"{base_filename}_codec.mp3") decoded_audio.write_hq_mp3(audio_path) s3_filepath = os.path.join( self.output_path, f"{base_filename}_codec.mp3" ) s3_client.upload_file( audio_path, "suno-data", s3_filepath, ExtraArgs={"ContentType": "audio/mp3"}, ) # prepare semantic codes print("semantic codes", semantic_codes.shape) semantic_codes = torch.from_numpy(semantic_codes).long() # generate normally using ctx upsampled_audio, upsampled_latents = generate_offline( audio=semantic_codes, lyrics=text, tags=tags, text_cfg_coef=text_cfg_coef, steps=steps, sigma_max=50.0, downscale_ctx_vector=self.downscale_ctx_vector, noise_ctx_vector=noise_ctx_vector, sampler_eta=1.0, sampler_s_noise=1.0, return_latents=True, seed=seed, use_ctx_vector=True, ) # save audio to s3 upsampled_audio_path = os.path.join( td, f"{base_filename}_diff_ctx.mp3" ) upsampled_audio.write_hq_mp3(upsampled_audio_path) # save latents to s3 # upsampled_latents_path = os.path.join( # td, f"{base_filename}}_vae.npz" # ) # np.savez(upsampled_latents_path, upsampled_latents) s3_filepath = os.path.join( self.output_path, f"{base_filename}_diff_ctx.mp3" ) print(f"Uploading to {s3_filepath}") s3_client.upload_file( upsampled_audio_path, "suno-data", s3_filepath, ExtraArgs={ "ContentType": "audio/mp3", }, ) # now generate without using ctx upsampled_audio, upsampled_latents = generate_offline( audio=semantic_codes, lyrics=text, tags=tags, text_cfg_coef=text_cfg_coef, steps=steps, sigma_max=50.0, downscale_ctx_vector=self.downscale_ctx_vector, sampler_eta=1.0, sampler_s_noise=1.0, return_latents=True, noise_ctx_vector=noise_ctx_vector, seed=seed, use_ctx_vector=False, ) # save audio to s3 upsampled_audio_path = os.path.join( td, f"{base_filename}_diff_no_ctx.mp3" ) upsampled_audio.write_hq_mp3(upsampled_audio_path) # save latents to s3 # upsampled_latents_path = os.path.join( # td, f"{base_filename}}_vae.npz" # ) # np.savez(upsampled_latents_path, upsampled_latents) s3_filepath = os.path.join( self.output_path, f"{base_filename}_diff_no_ctx.mp3" ) print(f"Uploading to {s3_filepath}") s3_client.upload_file( upsampled_audio_path, "suno-data", s3_filepath, ExtraArgs={ "ContentType": "audio/mp3", }, ) # s3_filepath = os.path.join( # self.output_path, f"{base_filename}_{upsample_uuid}_vae.npz" # ) # print(f"Uploading to {s3_filepath}") # s3_client.upload_file( # upsampled_latents_path, # "suno-data", # s3_filepath, # ) # create a json file with the metadata metadata = { "clip_id": work_item["id"], "seed": seed, "text_cfg_coef": text_cfg_coef, "text": text, "tags": tags, "steps": steps, "noise_ctx_vector": noise_ctx_vector, } metadata_path = os.path.join(td, f"{base_filename}.json") with open(metadata_path, "w") as f: json.dump(metadata, f) s3_filepath = os.path.join( self.output_path, f"{base_filename}_metadata.json" ) print(f"Uploading to {s3_filepath}") s3_client.upload_file( metadata_path, "suno-data", s3_filepath, ExtraArgs={ "ContentType": "application/json", }, ) except Exception as e: print(f"Error generating: {e}") return # prod diffusion model # DIT_MODEL_FILEPATH = "s3://suno-data/georg/tmp/2b_prefix_ft.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/tony/tmp/diff/dit_v3_dpo_t10_3k_5e6_b100_t25.pt" DIT_MODEL_FILEPATH = ( "s3://suno-data/christian/checkpoints/diffusion/v45_2b_step_2_600_000.pt" ) CODEC_MODEL_FILEPATH = "s3://suno-data/minz/models/dac_vae_tuned_25hz.pth" 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(DIT_MODEL_FILEPATH, CODEC_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 = 32 @app.cls( gpu=modal.gpu.A10G(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, dit_ckpt: str, codec_ckpt: str, output_path: str): import torch num_gpus = torch.cuda.device_count() print(f"Found {num_gpus} GPUs.") self.worker = GenerateWorker(dit_ckpt, codec_ckpt, output_path) @modal.method() def generate(self, work_item: list[dict]): return self.worker.generate(work_item) @app.local_entrypoint() def main(): ckpt_name = os.path.basename(DIT_MODEL_FILEPATH).split(".")[0] print(ckpt_name) # load the positive prompts # work_items = read_from_s3( # "s3://suno-data/christian/sft/pos_interesting_clips_up_u_1_20241201_full.jsonl", # read_f=read_jsonl, # ) work_items = read_jsonl( "/home/christian/code/christian/metadata/discogs_subset_sampled_metas.jsonl" ) chunksize = 16 # num of prompts per worker print(len(work_items), "work items") worker = GenerateStub( DIT_MODEL_FILEPATH, CODEC_MODEL_FILEPATH, f"christian/outputs/{ckpt_name}/discogs_subset_sampled_metas", ) # look for existing metadata files existing_files = list_s3_dir( f"s3://suno-data/christian/outputs/{ckpt_name}/discogs_subset_sampled_metas" ) existing_files = [f[0] for f in existing_files if f[0].endswith("_metadata.json")] existing_ids = [ os.path.basename(f).split("_metadata.json")[0] for f in existing_files ] existing_ids = set(existing_ids) # now filter the work items work_items = [w for w in work_items if w["id"] not in existing_ids] # split into chunks work_items = [ work_items[i : i + chunksize] for i in range(0, len(work_items), chunksize) ] print(f"{len(work_items)} chunks of {chunksize} work items") if False: 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") # Calculate predicted time based on number of chunks and available GPUs num_chunks = len(work_items) - 1 # Excluding the test chunk time_per_chunk = time.time() - t0 predicted_time = ( (time_per_chunk * num_chunks) / (N_MAX_REPLICAS if N_MAX_REPLICAS > 0 else 1) / 60 / 60 ) print(f"predicted time: {predicted_time}h for batch generation") 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") # aws s3 ls s3://suno-data/christian/outputs/v45_2b_step_2_600_000/discogs_subset_sampled_metas/ # aws s3 sync s3://suno-data/christian/outputs/v45_2b_step_2_600_000/discogs_subset_sampled_metas/