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 tqdm import tqdm from joblib import Parallel, delayed 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, read_from_s3 from suno_utils.tasks.hoot import ( encode_filepaths, encode, clean_text, encode_and_align, _get_alignable_tokens, _assign_word_timings, get_aligned_lyrics, ctc_align, load_model_list, get_word_timing_from_audio_and_lyrics, _merge_into_lines, decode_logits, load_model, ) from suno_utils.utils.metrics import get_cer from suno_utils.tasks.ear import load_model from sklearn.metrics.pairwise import cosine_similarity 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.shimmerscore import shimmerscore from suno_utils.diffusion import generation as diffusion_gen from suno_utils.tasks.upsample_engine import UpsampleEngine, Request, Job 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") ) def load_audio_from_s3(s3_path): """Load and convert audio from S3 path.""" try: audio = Audio.from_s3(s3_path, n_channels=2) # .convert( # sample_rate=24_000, byte_width=2, n_channels=1 # ) return audio except Exception as e: print(f"Error loading audio from {s3_path}: {e}") return None class GenerateWorker: def __init__( self, output_path: str, ): self.output_path = output_path start_time = time.time() print("Start loading models") # load diffusion model num_gpus = torch.cuda.device_count() cuda_device = torch.cuda.current_device() print(f"Found {num_gpus} GPUs. Using GPU {cuda_device}.") 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" ) _ = preload_semantic_models(semantic_model_filepath, semantic_clusters_filepath) # load ear model ear_model_filepath = "s3://suno-data/christian/checkpoints/ear/ear_v2_s3080.pt" self.ear_model = load_model(ear_model_filepath, compile=True) print( f"Finish loading models. Took {round(time.time() - start_time, 2)} seconds" ) @staticmethod def download_models(dir_path=MOUNT_PATH): """Download diffusion models.""" print("Start downloading models") _ = 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 ) print("Finish downloading models") def generate(self, work_item): """Generate audio from a work item.""" work_item_idx, work_item_contents = work_item print(f"Processing work item {work_item_idx}") # Prepare all S3 paths first all_s3_paths = [] path_to_request_id = ( {} ) # Map to keep track of which paths belong to which request for item in work_item_contents: request_id, neg_id, pos_id = item neg_s3_path = f"s3://suno-data-uploads/studio/uploads/{neg_id}.mp3" pos_s3_path = f"s3://suno-data-uploads/studio/uploads/{pos_id}.mp3" all_s3_paths.extend([neg_s3_path, pos_s3_path]) path_to_request_id[neg_s3_path] = (request_id, "neg") path_to_request_id[pos_s3_path] = (request_id, "pos") print(f"Loading {len(all_s3_paths)} audio files in parallel") # Load all audio files in parallel audio_files = Parallel(n_jobs=-1)( delayed(load_audio_from_s3)(path) for path in all_s3_paths ) # Create a mapping of paths to loaded audio path_to_audio = dict(zip(all_s3_paths, audio_files)) results = {} for i, item in enumerate(work_item_contents): request_id, neg_id, pos_id = item neg_s3_path = f"s3://suno-data-uploads/studio/uploads/{neg_id}.mp3" pos_s3_path = f"s3://suno-data-uploads/studio/uploads/{pos_id}.mp3" neg_audio = path_to_audio[neg_s3_path] pos_audio = path_to_audio[pos_s3_path] if neg_audio is None or pos_audio is None: print(f"Skipping due to failed audio loading for request {request_id}") continue # shimmer score for neg and pos neg_shimmer_score = shimmerscore(neg_audio) pos_shimmer_score = shimmerscore(pos_audio) semantic = encode_semantic( [ neg_audio.convert(sample_rate=24_000, byte_width=2, n_channels=1), pos_audio.convert(sample_rate=24_000, byte_width=2, n_channels=1), ], do_clustering=False, ) neg_semantic = semantic[0] pos_semantic = semantic[1] # Compute cosine similarity for each token (each row in the matrix) # This gives a matrix of shape (2698, 2698), but we only want the diagonal pairwise_sim_matrix = cosine_similarity(neg_semantic, pos_semantic) # Get the similarity of corresponding tokens (i.e., diagonal) tokenwise_similarities = np.diag(pairwise_sim_matrix) # Take the mean of the token-wise similarities mean_similarity = tokenwise_similarities.mean() # convert to python float mean_similarity = float(mean_similarity) # print(f"{i}/{len(work_item_contents)}: Mean similarity: {mean_similarity}") # run ear model if neg_audio.duration_s > 5.0: neg_ear_score = self.ear_model.get_score(neg_audio) else: neg_ear_score = None if pos_audio.duration_s > 5.0: pos_ear_score = self.ear_model.get_score(pos_audio) else: pos_ear_score = None results[request_id] = { "mean_similarity": mean_similarity, "neg_id": neg_id, "pos_id": pos_id, "neg_ear_score": neg_ear_score, "pos_ear_score": pos_ear_score, "neg_shimmer_score": neg_shimmer_score, "pos_shimmer_score": pos_shimmer_score, } # save the results to a file with tempfile.TemporaryDirectory() as td: with open(os.path.join(td, f"{work_item_idx}.json"), "w") as f: json.dump(results, f) s3_client.upload_file( os.path.join(td, f"{work_item_idx}.json"), "suno-data", os.path.join(self.output_path, f"{work_item_idx}.json"), ) OUTPUT_STR = "interesting_clips_ahi_d3_20250528" GPU_TYPE = "A10G" # H100, A10G CHUNKSIZE = 128 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() image = base_image.run_function(download_model_wrapper_d, secrets=SECRETS) APP_NAME = f"batch-semantic-similarity-gpu" app = modal.App(APP_NAME, image=image, secrets=SECRETS) N_MAX_REPLICAS = 64 @app.cls( # gpu=modal.gpu.H100(count=1), gpu=modal.gpu.A10G(count=1) if GPU_TYPE == "A10G" else 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, output_path: str): import torch num_gpus = torch.cuda.device_count() print(f"Found {num_gpus} GPUs.") self.worker = GenerateWorker( 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 filepath = ( "/home/tony/Data/Preference/up_v2_d3/interesting_clips_ahi_d3_20250528.pkl" ) metas = pd.read_pickle(filepath) work_items = [] odd_ids = np.arange(0, len(metas), 2) pbar = tqdm(odd_ids) for odd_id in pbar: neg_request_id = metas.iloc[odd_id]["request_id"] pos_request_id = metas.iloc[odd_id + 1]["request_id"] assert neg_request_id == pos_request_id neg_id = metas.iloc[odd_id]["id"] pos_id = metas.iloc[odd_id + 1]["id"] work_items.append((neg_request_id, neg_id, pos_id)) print(f"Total work items: {len(work_items)}") work_items = list(funcy.chunks(CHUNKSIZE, work_items)) print(f"Chunksize: {CHUNKSIZE}, total chunks: {len(work_items)}") # add a work_item_idx to each work item work_items = [(i, work_item) for i, work_item in enumerate(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 = [int(f[0].split("/")[-1].replace(".json", "")) for f in existing_ids] # existing_ids = list(set(existing_ids)) # print(f"Total existing ids: {len(existing_ids)}") # now remove these from the work items # work_items = [f for f in work_items if f[0] not in existing_ids] print(f"Total work items remaining: {len(work_items)}") worker = GenerateStub(f"christian/outputs/{OUTPUT_STR}") # 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")