"""Audio/text embedding application on modal.""" import os import modal from suno_utils.worker.schema import QueueItem import torch import json import shutil from suno_utils.audio import Audio from suno_utils.worker.loader import S3Loader from suno_utils.worker.utils import print_gpu_memory_usage from suno_utils.worker.settings import s3_client from suno_utils.gpt import chirp_v2_5 as chirp_v3 from suno_utils.models.ditto_v2.ditto_v2 import Ditto from suno_utils.worker.modal_base import get_modal_base_image ############## CHANGE THESE ############## DEPLOYMENT_TYPE = "dev" # dev, prod ########################################## ENCODER_CONCURRENCY_LIMITS = { "dev": 200, "prod": 250, } KEEP_WARM = { "dev": 1, "prod": 1, } DITTO_EMBEDDING_DIM = 128 # set number of cpus. VERBOSE_MESSAGE = DEPLOYMENT_TYPE == "dev" MOUNT_PATH = "/suno/models" aws_secret = modal.Secret.from_name("studio-aws") SECRETS = [ aws_secret, modal.Secret.from_name("openai-secret"), modal.Secret.from_name("api-callback-token"), ] # fix transformers for mert base_image = get_modal_base_image().pip_install("transformers==4.44.0") DITTO_S3_PATH = "s3://suno-data/minz/models/ditto_v2_epoch_57.pt" class DittoWorker(S3Loader): def __init__(self): S3Loader.__init__(self) print("Start loading models") self.ditto_path = os.path.join(MOUNT_PATH, "ditto_models", "ditto.pt") if not os.path.exists(self.ditto_path): self.ditto_path = chirp_v3._get_model_if_needed(DITTO_S3_PATH, cache_dir=MOUNT_PATH) print("found models locally.") self.ditto = Ditto( latent_dim=DITTO_EMBEDDING_DIM, model_path=self.ditto_path, is_flash=False, is_serving=True, ) self.ditto = self.ditto.eval().cuda() print("Finish loading models") @staticmethod def download_models(dir_path=MOUNT_PATH): print("Start downloading models") dl_path = chirp_v3._get_model_if_needed(DITTO_S3_PATH, cache_dir=dir_path) print(f"Downloaded model to {dl_path}") # Define source and destination paths source_ditto_path = dl_path dest_dir = os.path.join(dir_path, "ditto_models") # Create destination directory if it doesn't exist os.makedirs(dest_dir, exist_ok=True) # Move Ditto model file dest_ditto_path = os.path.join(dest_dir, "ditto.pt") shutil.move(source_ditto_path, dest_ditto_path) print(f"Moved Ditto model to {dest_ditto_path}") def download_model_wrapper_g(): # 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", DITTO_S3_PATH) DittoWorker.download_models() image = base_image.run_function(download_model_wrapper_g, secrets=SECRETS).add_local_python_source( "suno_utils", copy=False ) APP_NAME = f"ditto-v2-{DEPLOYMENT_TYPE}" app = modal.App(APP_NAME, image=image) @app.cls( gpu="T4", secrets=SECRETS, timeout=200, scaledown_window=400, retries=modal.Retries( max_retries=1, backoff_coefficient=2.0, initial_delay=5.0, ), max_containers=ENCODER_CONCURRENCY_LIMITS[DEPLOYMENT_TYPE], min_containers=0, # TODO: split this into multiple apps ) @modal.concurrent(max_inputs=32) # max 10 has ~ 5 GB at peak class DittoWorkerStub: def __init__(self): import torch num_gpus = torch.cuda.device_count() print(f"Found {num_gpus} GPUs.") self.worker = DittoWorker() @modal.method() def encode_audio( self, queue_item_json: str, ) -> None: if VERBOSE_MESSAGE: print_gpu_memory_usage(self.__class__.__name__) torch.cuda.reset_max_memory_allocated() queue_item = QueueItem.parse_raw(queue_item_json) dur = float(queue_item.metadata.get("dur", 30)) start = float(queue_item.metadata.get("start", 0)) if dur > 30: raise ValueError("Duration must be less than 30 seconds") task = queue_item.metadata.get("task", "self_sim") s3_id = queue_item.id s3_url = f"s3://suno-data-uploads/studio/uploads/{s3_id}.mp3" callback_url = queue_item.callback_url try: s3_client.get_object(Bucket="suno-data-uploads", Key=f"studio/uploads/{s3_id}.mp3") except: print(f"File {s3_id} doesn't exist on s3. Is it deleted?") return None audio = Audio.from_s3(s3_url, n_channels=1, sample_rate=24000) if audio.duration_s < 30: print(f"File {s3_id} is too short for ditto encoding.") return None print(f"Starting ditto job with {s3_id}, duration {audio.duration_s}.") audio = audio.get_segment(from_s=start, to_s=start + dur) wav = torch.tensor(audio.array_float).unsqueeze(0).cuda() emb = self.worker.ditto.music_to_latent(wav, task=task)[0].detach().cpu().numpy() if id and callback_url: payload = { **queue_item.metadata, "id": s3_id, "vector": emb.tolist(), "index_name": "song_vector_v2", } print(f"Notifying progress to {callback_url} with payload {payload}") queue_item.notify_progress(payload) print(f"finished ditto job with {s3_id}") return emb @modal.method() def encode_text(self, text: str) -> None: if VERBOSE_MESSAGE: print_gpu_memory_usage(self.__class__.__name__) torch.cuda.reset_max_memory_allocated() te = self.worker.ditto.text_to_latent("[CLS]" + text)[0].detach().cpu().numpy() return te @modal.method() def encode_texts(self, queue_item_json: str) -> list[dict]: queueItem = QueueItem(**json.loads(queue_item_json)) metadata = queueItem.metadata texts = metadata.get("encode_texts_in_ditto", []) embeddings = [] for text in texts: te = self.worker.ditto.text_to_latent("[CLS]" + text)[0].detach().cpu().numpy() embeddings.append({"text": text, "embedding": te.tolist()}) if queueItem.callback_url: for embedding in embeddings: queueItem.notify_progress( {**metadata, "id": queueItem.id, "embeddings": embedding["embedding"]} ) return embeddings @app.local_entrypoint() def main(): ditto_worker = DittoWorkerStub() queue_item = QueueItem(id="4a77dea7-19f3-46d2-8b0a-b2b7e9ea9a05", metadata={"task": "self_sim"}) queue_item_json = queue_item.model_dump_json() print(ditto_worker.encode_audio.remote(queue_item_json=queue_item_json)) print(ditto_worker.encode_text.remote("jazz")) queue_item = QueueItem( id="4a77dea7-19f3-46d2-8b0a-b2b7e9ea9a05", metadata={"encode_texts_in_ditto": ["jazz", "rock", "pop"]}, ) queue_item_json = queue_item.model_dump_json() print(ditto_worker.encode_texts.remote(queue_item_json=queue_item_json))