"""Audio conversion application on modal. Audio <--> code operations. Audio conversion to codes, codec cycle codes, codes to Audios. """ import json import logging import os import tempfile import time import modal import numpy as np from suno_utils.audio import Audio from suno_utils.gpt import chirp_v2 from suno_utils.gpt import chirp_v2_5 as chirp_v3 from suno_utils.tasks import ss_vad from suno_utils.worker.codec_glue import ( preload_vae_models_v1, preload_vae_models_v2, vae_version_to_decode_fn, vae_version_to_encode_fn, ) from suno_utils.worker.modal_model_volume import MODEL_STORE_VOLUME_PREFIX, CODEC_PATH_DICT from suno_utils.worker.loader import S3Loader, get_latents, retry_s3_upload from suno_utils.worker.modal_base import get_modal_base_image_with_flash_attention from suno_utils.worker.modal_model_configs import get_model_version, is_vae_model, VAEVersion from suno_utils.worker.schema import QueueItem from suno_utils.worker.settings import s3_client from suno_utils.worker.utils import retry_decorator from suno_utils.worker.modal_model_volume import MODEL_STORE_VOLUME_DIR, model_store_volume ############## CHANGE THESE ############## DEPLOYMENT_TYPE = "dev" # dev, prod ########################################## 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"), modal.Secret.from_name("datadog-metrics"), ] base_image = get_modal_base_image_with_flash_attention().pip_install("transformers==4.44.0") logger = logging.getLogger(__name__) fast_retry_s3_download = retry_decorator(3, wait_seconds=1)(s3_client.download_fileobj) class CycleWorker(S3Loader): def __init__(self, load_ss_vad: bool = False): """Init and load the modals to GPUs. Args: load_ss_vad: whether to load the ss_vad model Only for vocal / instrumental stems. """ S3Loader.__init__(self) # coarse models print("Start loading models") v2_codec_path = chirp_v2._get_model_if_needed(chirp_v2.CODEC_CKPT_PATH, cache_dir=MOUNT_PATH) chirp_v2.preload_codec_models(v2_codec_path) v3_codec_path = chirp_v3._get_model_if_needed(chirp_v3.CODEC_CKPT_PATH, cache_dir=MOUNT_PATH) chirp_v3.preload_codec_models(v3_codec_path) print("Start loading vae model") preload_vae_models_v1( f"{MODEL_STORE_VOLUME_PREFIX}{CODEC_PATH_DICT[VAEVersion.V_VAE_25_PEAQ_1.value]}", ) preload_vae_models_v2( f"{MODEL_STORE_VOLUME_PREFIX}{CODEC_PATH_DICT[VAEVersion.V_VAE_25_TUNED_2.value]}", ) print("Finish loading vae model") # semantic models semantic_ckpt_path = chirp_v3._get_model_if_needed( chirp_v3.SEMANTIC_CKPT_PATH, cache_dir=MOUNT_PATH ) semantic_centroids_path = chirp_v3._get_model_if_needed( chirp_v3.SEMANTIC_CENTROIDS_PATH, cache_dir=MOUNT_PATH ) chirp_v3.preload_semantic_models( checkpoint_filepath=semantic_ckpt_path, centroids_filepath=semantic_centroids_path, ) if load_ss_vad: ss_vad_config_path = chirp_v3._get_model_if_needed(ss_vad.YAML_PATH, cache_dir=MOUNT_PATH) ss_vad_model_path = chirp_v3._get_model_if_needed(ss_vad.MODEL_PATH, cache_dir=MOUNT_PATH) ss_vad.preload_models( checkpoint_filepath=ss_vad_model_path, config_path=ss_vad_config_path, ) print("Finish loading models") @staticmethod def download_models(dir_path=MOUNT_PATH): """Download the codec models for dac8, dac12.""" print("Start downloading models") # coarse models v2_codec_path = chirp_v2._get_model_if_needed(chirp_v2.CODEC_CKPT_PATH, cache_dir=dir_path) chirp_v2.preload_codec_models(v2_codec_path) v3_codec_path = chirp_v3._get_model_if_needed(chirp_v3.CODEC_CKPT_PATH, cache_dir=dir_path) chirp_v3.preload_codec_models(v3_codec_path) # semantic models semantic_ckpt_path = chirp_v3._get_model_if_needed( chirp_v3.SEMANTIC_CKPT_PATH, cache_dir=dir_path ) semantic_centroids_path = chirp_v3._get_model_if_needed( chirp_v3.SEMANTIC_CENTROIDS_PATH, cache_dir=dir_path ) chirp_v3.preload_semantic_models( checkpoint_filepath=semantic_ckpt_path, centroids_filepath=semantic_centroids_path, ) # get source seperatation model ss_vad_config_path = chirp_v3._get_model_if_needed(ss_vad.YAML_PATH, cache_dir=dir_path) ss_vad_model_path = chirp_v3._get_model_if_needed(ss_vad.MODEL_PATH, cache_dir=dir_path) ss_vad.preload_models( checkpoint_filepath=ss_vad_model_path, config_path=ss_vad_config_path, ) print("Finish downloading models.") def encode_audio_to_npz( self, audio: Audio, s3_npz_id: str = "", right_affine: bool = True, model_full_version: str = "3.0.0.0", encode_vae_version: str | None = None, ): """Given an input audio, encode the semantic and codec codes into npz file. This function can cylce any audio to the int codes we want. """ sem_arr = chirp_v3.semantic_encode(audio, right_affine=right_affine)[:, :1].T coarse_arr = chirp_v3.codec_encode(audio, right_affine=right_affine).T arr_len = min(sem_arr.shape[-1], coarse_arr.shape[-1]) sem_arr = sem_arr[:, sem_arr.shape[-1] - arr_len :] coarse_arr = coarse_arr[:, coarse_arr.shape[-1] - arr_len :] a_arr = np.concatenate([sem_arr, coarse_arr], axis=0).T # (T, C) a_arr = a_arr.astype(np.int32) print(f"Cycle {s3_npz_id}, audio: {audio.duration_s}s, a_arr: {a_arr.shape}") if s3_npz_id: print("Writing to s3") # Note that this writes as 3.0.0.0 self._write_npz(QueueItem(id=s3_npz_id, metadata={}), a_arr, APP_NAME, model_full_version) if encode_vae_version is not None: print(f"Cycle {s3_npz_id}: encode vae to {encode_vae_version}.") assert encode_vae_version in vae_version_to_encode_fn # this api takes a list of tensors... vae_arr = vae_version_to_encode_fn[encode_vae_version](audio) if not isinstance(vae_arr, np.ndarray): print(f"Cycle {s3_npz_id}: Failed to encode vae.") return temp_queue_item = QueueItem(id=s3_npz_id, metadata={}) self._write_vae_latents_npz( temp_queue_item, vae_arr, encode_vae_version, APP_NAME, n_sem_tokens=sem_arr.shape[0], ) return a_arr def download_model_wrapper_e(): # 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.") CycleWorker.download_models() image = base_image.run_function( download_model_wrapper_e, secrets=SECRETS, volumes={MODEL_STORE_VOLUME_DIR: model_store_volume} ).add_local_python_source("suno_utils", copy=False) APP_NAME = f"cycle-{DEPLOYMENT_TYPE}" app = modal.App(APP_NAME, image=image) @app.cls( gpu="A10G", secrets=SECRETS, timeout=400, # the wav encoding can take a while scaledown_window=480, memory=15000, retries=modal.Retries( max_retries=1, backoff_coefficient=2.0, initial_delay=5.0, ), min_containers=0, max_containers=1000, region="us-east", buffer_containers=0 if DEPLOYMENT_TYPE == "dev" else 1, volumes={MODEL_STORE_VOLUME_DIR: model_store_volume}, ) @modal.concurrent(max_inputs=2) class CodecCycleStub: def __init__(self): import torch num_gpus = torch.cuda.device_count() print(f"Found {num_gpus} GPUs.") self.worker = CycleWorker(load_ss_vad=False) def _cycle_code(self, item: QueueItem, upload_to_s3: bool = True): """Cycle the audio through the codec. Note that since v4, the audio is actually higher fidelity than codec decoded audios. So we don't need to decode from the codec tokens any more. This also generalized to future cases where the coarse tokens are not predicted by GPTs. Args: item: the queue item to cycle upload_to_s3: whether to upload the output to s3 """ model_full_version = "3.0.0.0" audio_s3_path = f"s3://suno-data-uploads/studio/uploads/{item.prompt_audio}.mp3" try: audio = Audio.from_s3(audio_s3_path.replace(".mp3", ".webm"), n_channels=2) except Exception as e: print(f"Cycle {item.id}, error getting audio from s3: {e}. Trying mp3.") audio = Audio.from_s3(audio_s3_path, n_channels=2) print(f"Cycle {item.id}, audio: {audio.duration_s}s") history_arr = self.worker.encode_audio_to_npz( audio, item.id, right_affine=False, model_full_version=model_full_version ) print(f"Cycle {item.id} Done!") if not upload_to_s3: return history_arr @modal.method() def cycle_code(self, queue_item: str, upload_to_s3: bool = True) -> None | np.ndarray: """Takes queue item (or a list of queueitems) and cycle them through semantic and codec.""" input_item = json.loads(queue_item) if isinstance(input_item, dict): item = QueueItem(**input_item) print(f"Start item: {item.id}") return self._cycle_code(item, upload_to_s3) elif isinstance(input_item, list): for sub_item in input_item: item = QueueItem(**sub_item) try: print(f"Start item: {item.id}") self._cycle_code(item, upload_to_s3) print(f"Finished item: {item.id}") except Exception as e: print(f"Error processing item: {item.id}, {str(e)}") else: print("Input format not supported.") @app.cls( gpu="A10G", secrets=SECRETS, timeout=400, # the wav encoding can take a while scaledown_window=480, memory=10000, retries=modal.Retries( max_retries=1, backoff_coefficient=2.0, initial_delay=5.0, ), max_containers=300, min_containers=1 if DEPLOYMENT_TYPE == "dev" else 2, region="us-east", buffer_containers=0 if DEPLOYMENT_TYPE == "dev" else 1, volumes={MODEL_STORE_VOLUME_DIR: model_store_volume}, ) @modal.concurrent(max_inputs=2) class CycleStub: def __init__(self): import torch num_gpus = torch.cuda.device_count() print(f"Found {num_gpus} GPUs.") self.worker = CycleWorker(load_ss_vad=False) @modal.method() def encode_audio( self, audio: Audio | str, s3_npz_id: str = "", right_affine: bool = True, encode_vae_version: str | None = None, ): if isinstance(audio, Audio): print("Warning, don't do this in prod. Sending audio wav is expensive") if isinstance(audio, str): audio = Audio.from_s3(audio, n_channels=2) print(f"CycleStub: {s3_npz_id}, audio: {audio.duration_s}s") # encode both with right_affine, to respect the boundary condition return self.worker.encode_audio_to_npz( audio, s3_npz_id, right_affine, encode_vae_version=encode_vae_version ) @app.cls( gpu="H100", secrets=SECRETS, timeout=400, # the wav encoding can take a while scaledown_window=480, memory=10000, retries=modal.Retries( max_retries=1, backoff_coefficient=2.0, initial_delay=5.0, ), max_containers=30, min_containers=1, region="us-east", buffer_containers=0, volumes={MODEL_STORE_VOLUME_DIR: model_store_volume}, ) @modal.concurrent(max_inputs=8) class VaeCycleStub: def __init__(self): import torch num_gpus = torch.cuda.device_count() print(f"Found {num_gpus} GPUs.") self.worker = CycleWorker(load_ss_vad=False) @modal.method() def encode_audio_to_vae_latents( self, audio_s3_path: str, s3_npz_id: str = "", encode_vae_version: str = "", return_vae_latents: bool = False, ): if isinstance(audio_s3_path, str): print(f"VAE cycle: {audio_s3_path}, get input request.") try: if audio_s3_path.endswith(".mp3"): try: audio = Audio.from_s3(audio_s3_path.replace(".mp3", ".webm"), n_channels=2) except Exception as e: print(f"VAE cycle: {audio_s3_path}, webm doesn't exist.") audio = Audio.from_s3(audio_s3_path, n_channels=2) else: audio = Audio.from_s3(audio_s3_path, n_channels=2) except Exception as e: print(f"VAE cycle: {audio_s3_path}, error getting audio from s3: {e}") return # encode both with right_affine, to respect the boundary condition print(f"VAE Cycle {s3_npz_id}: encode vae to {encode_vae_version}.") assert encode_vae_version in vae_version_to_encode_fn vae_arr = vae_version_to_encode_fn[encode_vae_version](audio) if not isinstance(vae_arr, np.ndarray): print(f"Cycle {s3_npz_id}: Failed to encode vae.") return # make sure it is the right datatype vae_arr = vae_arr.astype(np.float16) temp_queue_item = QueueItem(id=s3_npz_id, metadata={}) try: # if the file exists, we need to update it with the new vae latents in the npz file s3_client.get_object(Bucket="suno-data-uploads", Key=f"studio/uploads/{s3_npz_id}_vae.npz") with tempfile.TemporaryDirectory() as td: # download the audio file npz_filename = os.path.join(td, f"{s3_npz_id}_vae.npz") with open(npz_filename, "wb") as tmp_file: fast_retry_s3_download( "suno-data-uploads", f"studio/uploads/{s3_npz_id}_vae.npz", tmp_file ) vae_np = np.load(npz_filename) new_vae_npz = { **vae_np, f"vae_latents_{encode_vae_version}": vae_arr, } np.savez(npz_filename, **new_vae_npz) retry_s3_upload(npz_filename, "suno-data-uploads", f"studio/uploads/{s3_npz_id}_vae.npz") except s3_client.exceptions.NoSuchKey: print(f"{s3_npz_id} -- vae doesn't exist.") self.worker._write_vae_latents_npz(temp_queue_item, vae_arr, encode_vae_version, APP_NAME) if return_vae_latents: print(f"VAE Cycle {s3_npz_id}: return vae latents {vae_arr.shape}.") return vae_arr @app.cls( gpu="A10G", secrets=SECRETS, timeout=400, # the wav encoding can take a while scaledown_window=480, memory=10000, retries=modal.Retries( max_retries=1, backoff_coefficient=2.0, initial_delay=5.0, ), max_containers=300, min_containers=1 if DEPLOYMENT_TYPE == "dev" else 2, region="us-east", buffer_containers=0 if DEPLOYMENT_TYPE == "dev" else 1, volumes={MODEL_STORE_VOLUME_DIR: model_store_volume}, ) @modal.concurrent(max_inputs=3) class WavCycleStub: def __init__(self): import torch num_gpus = torch.cuda.device_count() print(f"Found {num_gpus} GPUs.") self.worker = CycleWorker(load_ss_vad=False) self.modal_f_encode_audio_to_vae = modal.Cls.from_name( f"cycle-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}", "VaeCycleStub", )().encode_audio_to_vae_latents @modal.method() def convert_to_wav(self, queue_item: str) -> str: """Convert request to wav. Upload to s3 and return the s3 address.""" start_time = time.time() item = QueueItem(**json.loads(queue_item)) audio_id = item.id print(f"Getting request: {queue_item}") extension = ".opus" if item.metadata.get("convert_to_opus", False) else ".wav" extension_without_dot = extension[1:] if "user_id" not in item.metadata or "clip_user_id" not in item.metadata: print(f"This is illegal access: {item.metadata}.") return f"s3://suno-data-uploads/studio/uploads/{audio_id}{extension}" try: s3_client.get_object(Bucket="suno-data-uploads", Key=f"studio/uploads/{audio_id}{extension}") item.notify_progress( { "id": item.id, "type": f"generate_{extension_without_dot}_file", }, ) print( f"Found existing {extension_without_dot} on s3. Total processing time {round(time.time() - start_time, 2)}." ) return f"s3://suno-data-uploads/studio/uploads/{audio_id}{extension}" except s3_client.exceptions.NoSuchKey: print(f"{item.id}: Converting file to {extension_without_dot}.") # for now use a simple check is_bot_request = isinstance(item.model_name, str) and "engine-b" in item.model_name # TODO: this is a hack is_vae = item.model_name is not None and is_vae_model(item.model_name) if is_bot_request: audio = Audio.from_s3( f"s3://suno-data-uploads/studio/uploads/{item.id}.mp3", n_channels=2 ) elif is_vae: print(f"{item.id}: Vae mode detected") try: with tempfile.TemporaryDirectory() as td: # this is a lazy load -- no vae conversion is done vae_latents, vae_latents_version = get_latents( os.path.join(td, f"{item.id}_vae.npz"), item.id, target_vae_version="", ) except Exception as e: print(f"Cycle {item.id}, error getting vae latents: {e}") vae_latents = None vae_latents_version = None n_default_stride_tokens = 60 * 25 # very long duration -- 1 min if vae_latents is None or vae_latents_version not in vae_version_to_decode_fn: logger.info(f"Missing VAE latents for id: {item.id}, encoding VAE on the fly") vae_latents_version = VAEVersion.V_VAE_25_PEAQ_1.value vae_f_call = self.modal_f_encode_audio_to_vae.spawn( audio_s3_path=f"s3://suno-data-uploads/studio/uploads/{item.id}.mp3", s3_npz_id=item.id, encode_vae_version=vae_latents_version, return_vae_latents=True, ) print(f"VAE Cycle {item.id}: encoding vae latents.") vae_latents = vae_f_call.get(timeout=60) print(f"VAE Cycle {item.id}: get callback.") assert vae_latents is not None audio = vae_version_to_decode_fn[vae_latents_version]( vae_latents, n_stride_tokens=min(n_default_stride_tokens, vae_latents.shape[0] - 5), ) print(f"{item.id}: decoded vae latents, {vae_latents_version}, {audio.duration_s}s") else: target_model_version = get_model_version(item.model_name if item.model_name else "v3") # if the input is a v2, we want to load it to v3 # otherwise, if they are 3.5, etc, it is the same codec for now. if target_model_version == "2.0.0.0": target_model_version = "3.0.0.0" if "continue_at" in item.metadata: item.metadata.pop("continue_at") if item.prompt_audio is None: # overwrite the prompt audio id to the audio_id for loading item.prompt_audio = audio_id audio_tokens = self.worker._load_audio_prompt(item, target_model_version) if audio_tokens is None: raise ValueError(f"Failed to load audio prompt for {audio_id}") codec_tokens = audio_tokens[:, 1:] n_default_stride_tokens = 60 * 25 # very long duration -- 1 min # node that this decode stream removes clicks audio = chirp_v3.codec_decode_stream_to_full_audio( codec_tokens, n_stride_tokens=min(n_default_stride_tokens, codec_tokens.shape[0] - 5), n_overlap_tokens=5, ) if item.metadata.get("convert_to_opus", False): # Note: today this path is only used for new studio, so we don't write webm with tempfile.NamedTemporaryFile(suffix=".opus") as temp_file: audio.write_opus(temp_file.name) retry_s3_upload( temp_file.name, "suno-data-uploads", f"studio/uploads/{audio_id}.opus", ExtraArgs={ "ContentType": "audio/ogg", }, ) item.notify_progress( { "id": item.id, "type": "generate_opus_file", }, ) print( f"Uploaded to s3 for item {audio_id}. Total processing time {round(time.time() - start_time, 2)}." ) return f"s3://suno-data-uploads/studio/uploads/{audio_id}.opus" else: with tempfile.NamedTemporaryFile(suffix=".wav") as temp_file: audio.write_wav(temp_file.name) retry_s3_upload( temp_file.name, "suno-data-uploads", f"studio/uploads/{audio_id}.wav", ExtraArgs={ "ContentType": "audio/wav", }, ) item.notify_progress( { "id": item.id, "type": "generate_wav_file", }, ) print( f"Uploaded to s3 for item {audio_id}. Total processing time {round(time.time() - start_time, 2)}." ) return f"s3://suno-data-uploads/studio/uploads/{audio_id}.wav" @app.cls( gpu="A10G", secrets=SECRETS, timeout=400, # the wav encoding can take a while scaledown_window=480, memory=10000, retries=modal.Retries( max_retries=1, backoff_coefficient=2.0, initial_delay=5.0, ), max_containers=300, min_containers=1 if DEPLOYMENT_TYPE == "dev" else 2, region="us-east", buffer_containers=0 if DEPLOYMENT_TYPE == "dev" else 1, volumes={MODEL_STORE_VOLUME_DIR: model_store_volume}, ) @modal.concurrent(max_inputs=4) class VocalDetectionStub: def __init__(self): import torch num_gpus = torch.cuda.device_count() print(f"Found {num_gpus} GPUs.") self.worker = CycleWorker(load_ss_vad=True) @modal.method() def detect_vocals_from_audio(self, audio_s3_path: str) -> bool: """Detect if vocals exist in the audio based on source separation.""" audio = Audio.from_s3(audio_s3_path, n_channels=2) audio_has_vocal = ss_vad.detect_vocals(audio) return audio_has_vocal @app.cls( cpu=4, gpu="A10G", secrets=SECRETS, timeout=400, # the wav encoding can take a while scaledown_window=480, memory=10000, retries=modal.Retries( max_retries=1, backoff_coefficient=2.0, initial_delay=5.0, ), min_containers=1, max_containers=300, region="us-east", buffer_containers=0 if DEPLOYMENT_TYPE == "dev" else 1, volumes={MODEL_STORE_VOLUME_DIR: model_store_volume}, ) @modal.concurrent(max_inputs=4) class StemStub: def __init__(self): import torch num_gpus = torch.cuda.device_count() print(f"Found {num_gpus} GPUs.") self.worker = CycleWorker(load_ss_vad=True) self.modal_f_video_generation = modal.Cls.from_name( f"videos-v2-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}", "DummyV0Stub", )().write_video self.modal_f_encode_audio = modal.Cls.from_name( f"cycle-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}", "CycleStub", )().encode_audio @modal.method() def split_vocals_and_instrumentals_from_audio(self, queue_item: str) -> None: item = QueueItem(**json.loads(queue_item)) print(f"StemStub received: {queue_item}") # for now use a simple check is_bot_request = isinstance(item.model_name, str) and "engine-b" in item.model_name assert item.ids is not None and len(item.ids) == 2 if not is_bot_request: try: is_vae = item.model_name is not None and ("v4" in item.model_name) if is_vae and item.prompt_audio is not None: print(f"{item.id}: Vae mode detected") with tempfile.TemporaryDirectory() as td: # this is a lazy load -- no vae conversion is done vae_latents, vae_latents_version = get_latents( os.path.join(td, f"{item.prompt_audio}_vae.npz"), item.prompt_audio, target_vae_version="", ) n_default_stride_tokens = 60 * 25 # very long duration -- 1 min assert vae_latents is not None # TODO: this is hacky design -- see the TODO on the loader for future loading audio = vae_version_to_decode_fn[vae_latents_version]( vae_latents, n_stride_tokens=min(n_default_stride_tokens, vae_latents.shape[0] - 5), ) print(f"{item.id}: decoded vae latents, {vae_latents_version}, {audio.duration_s}s") else: # default to v3 in case the model name is empty -- happens for audio uploads target_model_version = get_model_version( item.model_name if item.model_name else "v3" ) # if the input is a v2, we want to load it to v3 # otherwise, if they are 3.5, etc, it is the same codec for now. if target_model_version == "2.0.0.0": target_model_version = "3.0.0.0" audio_tokens = self.worker._load_audio_prompt(item, target_model_version) if audio_tokens is None: raise ValueError(f"Failed to load audio prompt for {item.prompt_audio}") codec_tokens = audio_tokens[:, 1:] n_default_stride_tokens = 60 * 25 # very long duration -- 1 min # node that this decode stream removes clicks audio = chirp_v3.codec_decode_stream_to_full_audio( codec_tokens, n_stride_tokens=min(n_default_stride_tokens, codec_tokens.shape[0] - 5), n_overlap_tokens=5, ) print(f"{item.id}: decoded rvq latents, {audio.duration_s}s") except Exception as e: print(f"Error loading audio: {e}") audio = Audio.from_s3( f"s3://suno-data-uploads/studio/uploads/{item.prompt_audio}.mp3", n_channels=2 ) print(f"Stem {item.id}: audio loaded, {audio.duration_s}s") vocals_array = ss_vad.encode(audio) assert isinstance(vocals_array, np.ndarray) if np.abs(vocals_array).max() > 1: vocals_array = vocals_array / np.abs(vocals_array).max() audio = audio.convert( sample_rate=ss_vad.SAMPLE_RATE, byte_width=audio.byte_width, n_channels=audio.n_channels ) print( f"Done with source separation: vocals shape = " f"{vocals_array.shape}, audio shape = " f"{audio.array_float.shape}" ) max_audio_length = min(vocals_array.shape[1], audio.array_float.shape[1]) instrumentals_array = ( audio.array_float[:, :max_audio_length] - vocals_array[:, :max_audio_length] ) # normalize since it is possible to overflow if np.abs(instrumentals_array).max() > 1: instrumentals_array = instrumentals_array / np.abs(instrumentals_array).max() vocals_audio = Audio.from_array_float(vocals_array, sample_rate=ss_vad.SAMPLE_RATE) instrumentals_audio = Audio.from_array_float( instrumentals_array, sample_rate=ss_vad.SAMPLE_RATE ) else: # for bots -- directly pull the audio from the s3 audio = Audio.from_s3( f"s3://suno-data-uploads/studio/uploads/{item.prompt_audio}.mp3", n_channels=2 ) # just normalize volume :) vocals_audio = audio.normalize_volume(target_db=-10) instrumentals_audio = audio.normalize_volume(target_db=-40) vocals_s3_id, instrumentals_s3_id = item.ids if "image_s3_id" not in item.metadata and isinstance(item.prompt_audio, str): item.metadata["image_s3_id"] = "image_" + item.prompt_audio.split(".")[0] # hide the callback url to avoid bugging mc for the video finish vocals_item = item.model_copy( deep=True, update={"id": vocals_s3_id, "ids": None, "callback_url": None} ) self.worker._write_audio_only(vocals_item, vocals_audio) self.worker.copy_and_upload_image(vocals_item) # hide the callback url to avoid bugging mc for the video finish instrumentals_item = item.model_copy( deep=True, update={"id": instrumentals_s3_id, "ids": None, "callback_url": None} ) self.worker._write_audio_only(instrumentals_item, instrumentals_audio) self.worker.copy_and_upload_image(instrumentals_item) _ = self.modal_f_video_generation.spawn( vocals_item.model_dump_json(), f"image_{vocals_item.id}.jpeg", ) # also encode the audio to vae latents self.modal_f_encode_audio.spawn( audio=f"s3://suno-data-uploads/studio/uploads/{vocals_s3_id}.mp3", s3_npz_id=vocals_s3_id, encode_vae_version=VAEVersion.V_VAE_25_PEAQ_1.value, ) self.modal_f_encode_audio.spawn( audio=f"s3://suno-data-uploads/studio/uploads/{instrumentals_s3_id}.mp3", s3_npz_id=instrumentals_s3_id, encode_vae_version=VAEVersion.V_VAE_25_PEAQ_1.value, ) # clear the instrumental's prompt instrumentals_item.metadata["prompt"] = "" _ = self.modal_f_video_generation.spawn( instrumentals_item.model_dump_json(), f"image_{instrumentals_item.id}.jpeg", ) if item.callback_url: item.notify_progress( { "id": item.id, "type": "edit_stems", "stem_ids": [vocals_s3_id, instrumentals_s3_id], "source_clip_id": item.metadata.get("source_clip_id", item.prompt_audio), }, ) print("Done with uploads separation") return @app.local_entrypoint() def main(): # for a codec cycle (int to int) test input = json.dumps( { "id": "49a07421-8f87-49cb-a8d5-b78ff880c35b_gen_cycle_0", "prompt_audio": "49a07421-8f87-49cb-a8d5-b78ff880c35b_gen_0", "prompt_text": "", "metadata": {}, }, ) model = CodecCycleStub() model.cycle_code.remote(input) # test vae encoding vae_model = VaeCycleStub() vae_f_call = vae_model.encode_audio_to_vae_latents.spawn( "s3://suno-data-uploads/studio/uploads/a17f7ec0-e7a4-437d-92c6-4d1faa6e99f5.mp3", s3_npz_id="test_vae_encode", encode_vae_version=VAEVersion.V_VAE_25_PEAQ_1.value, return_vae_latents=True, ) vae_latents = vae_f_call.get(timeout=60) print("Got callback, vae latents shape = ", vae_latents.shape) model = CycleStub() # test encode audio model.encode_audio.remote( Audio.from_beep(48000, duration_s=8, n_channels=2), s3_npz_id="test_encode", ) wav_model = WavCycleStub() # test wav generation # NOTE you will have to update the id to test the wav conversion wav_gen_input = json.dumps( { "id": "a17f7ec0-e7a4-437d-92c6-4d1faa6e99f5", "model_name": "chirp-v4-h-s-32", "metadata": {"user_id": 0, "clip_user_id": 0}, }, ) wave_file_path = wav_model.convert_to_wav.remote(wav_gen_input) print(wave_file_path) # test source separation track_separation_input = json.dumps( { "id": "0d4de30e-e796-4531-a472-6d1344508f95", "prompt_audio": "5e3a9ab5-6735-49c8-93e4-2d4bc68c05c3", "prompt_npz": None, "prompt_text": None, "metadata": { "image_s3_id": "image_a0eb8347-4d5c-40ca-a3ac-241e25634ee8", "source_clip_id": "5e3a9ab5-6735-49c8-93e4-2d4bc68c05c3", }, "gen_duration": 12, "callback_url": "https://studio-api.suno.ai/api/edit/webhook/finish-stems/", "model_name": "chirp-v3-engine-i", "title": None, "ids": ["3a726cc3-3e00-41a0-bb10-7c7d3eff23b6", "39554f59-49ad-4e06-b84d-6246b5d07e64"], }, ) stem_model = StemStub() stem_model.split_vocals_and_instrumentals_from_audio.remote(track_separation_input) # test vocal detection vocal_detection_model = VocalDetectionStub() expected_results = [True, True, True, False, False, False, True, False, True, False] for i, expected_result in enumerate(expected_results): start_time = time.time() audio_s3_path = f"s3://suno-data/tony/test_audio/user_uploads/{i + 1}.mp3" file_has_vocal = vocal_detection_model.detect_vocals_from_audio.remote(audio_s3_path) print( f"Tested file {i} has vocal: {file_has_vocal}, " f"expected {expected_result}, " f"{time.time() - start_time:.2f}s elapsed." ) if file_has_vocal != expected_result: print("WARNING! Vocal detection model is not working as expected.") print("Done")