import os import glob 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.tasks.ear import load_model 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( "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.5.1", "torchaudio==2.5.1", "transformers==4.44.0", "numpy==1.26.4", ) ) class EarWorker: def __init__( self, ear_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}.") self.model = load_model(ear_model_filepath, compile=True) print( f"Finish loading models. Took {round(time.time() - start_time, 2)} seconds" ) def score(self, work_items): print(f"Scoring {len(work_items)} work items") for filepath_group in work_items: from joblib import Parallel, delayed def process_filepath(filepath): s3_filepath = f"s3://suno-data/{filepath}" base_s3_id = filepath.split("_upsample_")[0] actual_base_s3_id = base_s3_id.split("/")[-1] upsample_id = ( filepath.split("_upsample_")[1].split(".mp3")[0].split("_")[0] ) audio_tensor, sr = read_from_s3(s3_filepath, read_f=torchaudio.load) return upsample_id, audio_tensor, sr, actual_base_s3_id # Process all filepaths in parallel results = Parallel(n_jobs=-1)( delayed(process_filepath)(filepath) for filepath in filepath_group ) # Create a dictionary of audio tensors indexed by upsample_id example_scores = {} for upsample_id, audio_tensor, sample_rate, base_id in results: scores, mean_score = self.model.get_score( audio_tensor, sample_rate=sample_rate, return_scores=True ) example_scores[upsample_id] = { "mean_score": mean_score, "scores": scores, } print(f"{upsample_id}: {mean_score}") print(f"Saving scores for {base_id}") # save the example_scores to a json file with tempfile.TemporaryDirectory() as temp_dir: temp_dir = pathlib.Path(temp_dir) json_filepath = os.path.join(temp_dir, f"{base_id}_ear_scores.json") with open(json_filepath, "w") as f: json.dump(example_scores, f) # upload the json file to s3 upload_s3_files( json_filepath, f"s3://suno-data/{self.output_path}/ear_scores/{base_id}_ear_scores.json", ) # 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" 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") # # EarWorker.download_models(EAR_MODEL_FILEPATH) image = base_image.run_function(download_model_wrapper_d, secrets=SECRETS) APP_NAME = f"batch-score-ear" app = modal.App(APP_NAME, image=image, secrets=SECRETS) N_MAX_REPLICAS = 64 EAR_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/ear/ear_v2_s3080.pt" @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, ear_ckpt: str, output_path: str): import torch num_gpus = torch.cuda.device_count() print(f"Found {num_gpus} GPUs.") self.worker = EarWorker(ear_ckpt, output_path) @modal.method() def generate(self, work_item: list[dict]): return self.worker.score(work_item) @app.local_entrypoint() def main(): filepath = "/home/christian/code/christian/metadata/sft/pos_interesting_clips_up_u_1_20241201_full_v2_grouped_filepaths.json" with open(filepath, "r") as f: grouped_filepaths = json.load(f) print(f"number of groups: {len(grouped_filepaths)}") # remove filepaths that have already been scored scored_filepaths = list_s3_dir( "s3://suno-data/christian/outputs/v45_2b_step_2_600_000/pos_interesting_clips_up_u_1_20241201_full_v2/ear_scores" ) scored_base_s3_ids = [ filepath[0] .split("_upsample_")[0] .replace("_ear_scores.json", "") .replace("ear_scores/", "") for filepath in scored_filepaths ] # print(scored_base_s3_ids[0]) # print(list(grouped_filepaths.keys())[0]) print(f"number of scored filepaths: {len(scored_base_s3_ids)}") grouped_filepaths = { k: v for k, v in grouped_filepaths.items() if k not in scored_base_s3_ids } # convert the grouped_filepaths to a list of lists grouped_filepaths = list(grouped_filepaths.values()) print(f"number of groups after removing scored filepaths: {len(grouped_filepaths)}") chunksize = 32 # num of prompts per worker # split the filepaths into chunks of chunksize work_items = [ grouped_filepaths[i : i + chunksize] for i in range(0, len(grouped_filepaths), chunksize) ] print(len(work_items), "work items") worker = GenerateStub( EAR_MODEL_FILEPATH, "christian/outputs/v45_2b_step_2_600_000/pos_interesting_clips_up_u_1_20241201_full_v2", ) 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") if True: 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")