"""Decode and merge audios application on modal.""" import json import os import tempfile import time import traceback import uuid from typing import Any import ffmpeg import modal import numpy as np from suno_utils.audio import Audio from suno_utils.audio.conversion import change_audio_speed, fade_out_audio from suno_utils.gpt import chirp_v2, chirp_v2_5 from suno_utils.worker.codec_glue import ( preload_vae_models_v1, preload_vae_models_v2, vae_version_to_decode_fn, ) from suno_utils.worker.loader import VERSION_GPT_MAPPING, S3Loader, get_latents, get_tokens from suno_utils.worker.modal_base import get_modal_base_image from suno_utils.worker.schema import QueueItem from suno_utils.worker.settings import s3_client from suno_utils.worker.utils import download_models_to_dir, retry_decorator from suno_utils.worker.modal_model_configs import get_vae_version_gain_adjust ############## CHANGE THESE ############## DEPLOYMENT_TYPE = "dev" # dev, prod ########################################## CONCURRENCY_LIMITS = { "dev": 4, "prod": 100, } KEEP_WARM = { "dev": 1, "prod": 4, } CHIRP_V2_MODELS = dict( codec_path="georg/models/codec/dac_2c_25x8.pt", ) CHIRP_V3_MODELS = dict( codec_path="georg/models/codec/dac_2c_25x12.pt", ) VAE_MODELS = dict( codec_path_v1="christian/25hz_vae_peaq_kl_0.005.pth", codec_path_v2="minz/models/dac_vae_tuned_25hz.pth", ) MOUNT_PATH = "/suno/models" SECRETS = [ modal.Secret.from_name("studio-aws"), modal.Secret.from_name("openai-secret"), modal.Secret.from_dict( { "DD_SITE": "datadoghq.com", "DD_ENV": DEPLOYMENT_TYPE, "DD_SERVICE": DEPLOYMENT_TYPE, "DD_LOGS_ENABLED": "false", "DD_TRACE_ENABLED": "true" if DEPLOYMENT_TYPE == "dev" else "false", # memory leak }, ), modal.Secret.from_name("datadog-metrics"), modal.Secret.from_name("api-callback-token"), ] retry_s3_download = retry_decorator(3, wait_seconds=20)(s3_client.download_fileobj) fast_retry_s3_download = retry_decorator(3, wait_seconds=1)(s3_client.download_fileobj) def _get_clip_vae_version(clip_id: str, default_vae_version: str = "v_vae_25_peaq_1") -> str: try: with tempfile.NamedTemporaryFile(suffix=".npz") as tmp_file: fast_retry_s3_download("suno-data-uploads", f"studio/uploads/{clip_id}_vae.npz", tmp_file) latent_npz = np.load(tmp_file.name) return str(latent_npz["v_vae"].item()) except Exception as e: print(f"{clip_id}: failed to get vae latents version: {e}") return default_vae_version def get_dimensions_and_fps(file_path): for stream in ffmpeg.probe(file_path)["streams"]: if "width" in stream and "height" in stream: return stream["width"], stream["height"], int(stream["r_frame_rate"].split("/")[0]) raise ValueError("Can't find the stream") class ConcatWorker(S3Loader): def __init__(self): super().__init__() self.modal_f_video_generation = modal.Cls.lookup( f"videos-v2-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}", "DummyV0Stub", )().write_video def preload(self): start_time = time.time() chirp_v2.preload_codec_models( f"{MOUNT_PATH}/{CHIRP_V2_MODELS['codec_path']}", ) chirp_v2_5.preload_codec_models( f"{MOUNT_PATH}/{CHIRP_V3_MODELS['codec_path']}", ) preload_vae_models_v1( f"{MOUNT_PATH}/{VAE_MODELS['codec_path_v1']}", ) preload_vae_models_v2( f"{MOUNT_PATH}/{VAE_MODELS['codec_path_v2']}", ) finish_time = time.time() print(f"Preloading took {finish_time - start_time}s") @staticmethod def download_models(dir_path=MOUNT_PATH): """Use AWS CLI to download models if they don't exist.""" files = ( list(CHIRP_V3_MODELS.values()) + list(CHIRP_V2_MODELS.values()) + list(VAE_MODELS.values()) ) download_models_to_dir(files, dir_path) def concat_videos( self, new_id: str, ids: list[str], continue_at_list: list[float | None] | None = None, include_history_list: list[float | None] | None = None, ): """ new_id is the id of the output clip ids is a list of ids of the input clips continue_at_list is a list of floats, continue_at time in seconds include_history_list is a list of floats, include_history time in seconds """ start_time = time.time() print(f"{new_id}: Request new_id") print(f"{new_id}: {ids}, {continue_at_list}, {include_history_list}") assert len(ids) >= 2 # Old backend could use concat_videos function without continue_at_list. # Otherwise, continue_at_list should be the same length as ids if not continue_at_list or (len(continue_at_list) != len(ids)): print("continue_at and ids length mismatch") continue_at_list = [None] * len(ids) if not include_history_list or (len(include_history_list) != len(ids)): print("include_history and ids length mismatch") include_history_list = [None] * len(ids) with tempfile.TemporaryDirectory() as td: # download all the .npz files and get compatible trimmed tokens # need to get the last one first to understand the model version last_id = ids[-1] filename = os.path.join(td, f"{last_id}.npz") with open(filename, "wb") as tmp_file: retry_s3_download("suno-data-uploads", f"studio/uploads/{last_id}.npz", tmp_file) npz = np.load(filename) # super fucking hacky way to check if the last file is from hybrid is_last_hybrid = False if "inference_runner" in npz and "chirpv4" in str(npz["inference_runner"].item()): is_last_hybrid = True if "model_version" in npz: full_model_version = npz["model_version"].item() arch, major, _, _ = full_model_version.split(".") target_version = f"{arch}.{major}" tokens = npz[f"v{target_version}_raw"] else: full_model_version = "2.0.0.0" target_version = "2.0" tokens = npz["v1_raw"] tokens_list = [ get_tokens(os.path.join(td, f"{id}.npz"), id, target_version, to_s=continue_at) for id, continue_at in zip(ids[:-1], continue_at_list[:-1]) ] tokens_list += [tokens] if target_version == "2.0": new_npz = { "v1_raw": np.concatenate(tokens_list, axis=0), } else: new_npz = { f"v{target_version}_raw": np.concatenate(tokens_list, axis=0), "concat_runner": APP_NAME, "model_version": full_model_version, } full_npz_path = os.path.join(td, "concat.npz") np.savez(full_npz_path, **new_npz) print( f"{new_id}: Processing time:{round(time.time() - start_time, 2)} seconds, download all npz files" ) all_latents = None if is_last_hybrid: print(f"{new_id}: loading vae latents") # we need to load the latents for each clip, and crop... vae_filename = os.path.join(td, f"{last_id}.npz") with open(vae_filename, "wb") as tmp_file: retry_s3_download("suno-data-uploads", f"studio/uploads/{last_id}_vae.npz", tmp_file) vae_npz = np.load(vae_filename) target_vae_version = str(vae_npz["v_vae"].item()) latents_list = [ get_latents( os.path.join(td, f"{id}_vae.npz"), id, to_s=continue_at, target_vae_version=target_vae_version, ) for id, continue_at in zip(ids, continue_at_list) ] if any(latents is None for latents, _ in latents_list): print( f"{new_id}: No valid vae latents found for some of the clips. Can't concat vae latents." ) else: all_latents = np.concatenate([latents for latents, _ in latents_list], axis=0) full_latents_npz_path = None if all_latents is not None: full_latents_npz_path = os.path.join(td, "concat_vae.npz") new_latents_npz = { "vae_latents": all_latents, "concat_runner": APP_NAME, "v_vae": target_vae_version, } np.savez(full_latents_npz_path, **new_latents_npz) concated_codec_tokens = np.concatenate(tokens_list, axis=0)[:, 1:] estimated_duration = concated_codec_tokens.shape[0] / 25 print(f"{new_id}: Estimated audio duration {estimated_duration}s") if estimated_duration < 1: raise ValueError("The estimated duration is less than 1 second.") used_vae_latents_version = None if estimated_duration < 60 * 20: # greater than 20 mins, we will still do it # we increase the n_stride_tokens so we have less overlap and faster decode n_default_stride_tokens = 60 * 25 # very long duration -- 1 min if all_latents is not None: audio = vae_version_to_decode_fn[target_vae_version]( all_latents, n_stride_tokens=min(n_default_stride_tokens, all_latents.shape[0] - 5), ) used_vae_latents_version = target_vae_version print(f"{new_id}: decoded vae latents, {round(audio.duration_s, 2)}s") else: audio = VERSION_GPT_MAPPING.get( f"{target_version}" ).codec_decode_stream_to_full_audio( concated_codec_tokens, n_stride_tokens=min(n_default_stride_tokens, concated_codec_tokens.shape[0] - 5), n_overlap_tokens=5, ) print(f"{new_id}: decoded codec tokens, {round(audio.duration_s, 2)}s") else: # we dont' want to decode again...but we also don't want to fail # we will just the mp3s and concat them. This is not ideal but it's better than failing. print(f"{new_id}: Will use mp3s to get the long audio.") audios_files = [f"{id}.mp3" for id in ids] audios_from_mp3 = [] for fname, continue_at in zip(audios_files, continue_at_list): full_path = os.path.join(td, fname) with open(full_path, "wb") as tmp_file: retry_s3_download( "suno-data-uploads", f"studio/uploads/{fname}", tmp_file, ) audio_from_mp3 = Audio.from_file(full_path, n_channels=2) if continue_at: audio_from_mp3 = audio_from_mp3.get_segment(to_s=continue_at) audios_from_mp3.append(audio_from_mp3) audio = Audio.concatenate(audios_from_mp3) print( f"{new_id}: Processing time:{round(time.time() - start_time, 2)} seconds, decoded audio {round(audio.duration_s, 2)}." ) temp_item = QueueItem(id=new_id, metadata={}) self._write_audio_only( temp_item, audio, gain_adjust=get_vae_version_gain_adjust(used_vae_latents_version) ) s3_client.upload_file(full_npz_path, "suno-data-uploads", f"studio/uploads/{new_id}.npz") if full_latents_npz_path is not None: s3_client.upload_file( full_latents_npz_path, "suno-data-uploads", f"studio/uploads/{new_id}_vae.npz" ) print( f"{new_id}: Processing time:{round(time.time() - start_time, 2)} seconds, finish uploads" ) return audio.duration_s def download_model_wrapper_6(): print("Downloading models....") ConcatWorker.download_models() image = ( get_modal_base_image() .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") .pip_install("flashinfer", index_url="https://flashinfer.ai/whl/cu121/torch2.4/") .add_local_python_source("suno_utils", copy=True) .run_function(download_model_wrapper_6, secrets=[modal.Secret.from_name("studio-aws")]) ) APP_NAME = f"concat-v3-{DEPLOYMENT_TYPE}" app = modal.App(APP_NAME, image=image) @app.cls( cpu=4.0, gpu="A10G", # T4s have slower ffmpegs... secrets=SECRETS, timeout=400, # ppl are crazy... scaledown_window=240, retries=modal.Retries( max_retries=2, backoff_coefficient=2.0, initial_delay=5.0, ), max_containers=CONCURRENCY_LIMITS[DEPLOYMENT_TYPE], min_containers=1, cloud="aws", region="us-east", ) @modal.concurrent(max_inputs=4) class ConcatInfillingStub: def __init__(self): self.worker = ConcatWorker() self.worker.preload() self.modal_f_video_generation = modal.Cls.lookup( f"videos-v2-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}", "DummyV0Stub", )().write_video self.modal_f_cycle = modal.Cls.lookup( f"cycle-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}", "CycleStub", )().encode_audio @modal.method() def concat_infilling_with_queue_item(self, queue_item: str, callback_url: str | None = None): """For inflling we do something special.""" start_time = time.time() if DEPLOYMENT_TYPE == "dev": print(f"Recieved {queue_item}") item = QueueItem(**json.loads(queue_item)) if callback_url: item.callback_url = callback_url # this is technically not a real concat! # we just need to decode out the full array # and regenerate the video audio_id = item.prompt_audio if not isinstance(audio_id, str): print(f"ConcatInfilling {item.id}: Can't find audio id {audio_id}. {queue_item}") item.notify_progress( { "id": item.id, "type": "concat_infilling", "ok": 0, "error_type": "generation_failure", "error_message": "Can't retrieve audio info for infilling.", } ) raise ValueError("Can't retrieve audio info for infilling.") with tempfile.TemporaryDirectory() as td: # download all the .npz files and get compatible trimmed tokens # need to get the last one first to understand the model version filename = os.path.join(td, f"{audio_id}.npz") with open(filename, "wb") as tmp_file: retry_s3_download("suno-data-uploads", f"studio/uploads/{audio_id}.npz", tmp_file) npz = np.load(filename) tokens = npz.get("full_arr") infill_inference_runner = str(npz["inference_runner"].item()) # TODO: this is not the best way -- QueueItem should have an info for this is_diffusion_infill = "cycle" in infill_inference_runner # TODO: for diffusion we need to concat the semantictokens -- this is not done yet if tokens is None and not is_diffusion_infill: raise ValueError(f"Can't load item {item.id}, audio {audio_id}") if is_diffusion_infill: # TODO: for diffusion we need to concat the semantic tokens if we don't want to cycle... pass else: concated_codec_tokens = tokens[:, 1:] estimated_duration = tokens.shape[0] / 25 print("estimated audio duration", estimated_duration) if estimated_duration < 1: raise ValueError("The estimated duration is less than 1 second.") # we increase the n_stride_tokens so we have less overlap and faster decode n_default_stride_tokens = 60 * 25 # very long duration -- 1 min target_version = "4.0" # this is fixed for now total_latents = None target_vae_version = None if "chirpv4" in infill_inference_runner or is_diffusion_infill: try: print(f"{item.id}: loading vae latents") # we need to load the latents for each clip, and crop... with open(filename, "wb") as tmp_file: fast_retry_s3_download( "suno-data-uploads", f"studio/uploads/{audio_id}_vae.npz", tmp_file ) latent_npz = np.load(filename) infill_with_pad_latents = latent_npz["vae_latents"] infill_parent_id = str(latent_npz["parent_clip_id"].item()).replace(".mp3", "") target_vae_version = str(latent_npz["v_vae"].item()) parent_latents, _ = get_latents( os.path.join(td, f"{infill_parent_id}_vae.npz"), infill_parent_id, target_vae_version=target_vae_version, ) assert isinstance(parent_latents, np.ndarray) # note these are loaded based on semantic codes # in a typical case, vae is < semantic codes by 5 tokens # parent semantic_l is x # parent vae_l is x - 5 # on top of that, the infill is also < semantic codes by 5 tokens... # so we don't need to shift....X.x parent_from_start_index = latent_npz.get("parent_from_start_index", 0) parent_from_end_index = latent_npz.get("parent_from_end_index", 0) # calculate the stupid parent offset parent_codes = get_tokens( os.path.join(td, f"{infill_parent_id}.npz"), infill_parent_id, target_version=target_version, ) parent_offset_sem_vae = 0 expected_latents_shape = ( tokens.shape[0] if not is_diffusion_infill else parent_codes.shape[0] ) if isinstance(parent_codes, np.ndarray): parent_offset_sem_vae = parent_codes.shape[0] - parent_latents.shape[0] # if the parent already have an vae issue -- we want to offset the expected # this is because the end index count from the end of the parent # super complicated -- I hate this # Note: in Jan 2025 we rolled out a fix on vae decoding # It no longer has the 5 sec offset. # But for previous clips, we still have this issue. # parent_from_end_index is counted from the semantic codes end. # THe solution is to subsctract the offset from the parent_from_end_index. if parent_offset_sem_vae > 0: print(f"{item.id} infilling parent offset is {parent_offset_sem_vae}") expected_latents_shape -= parent_offset_sem_vae parent_from_end_index -= parent_offset_sem_vae retrieved_latents_shape = ( parent_from_start_index + infill_with_pad_latents.shape[0] + parent_from_end_index ) if retrieved_latents_shape < expected_latents_shape: parent_from_end_index += expected_latents_shape - retrieved_latents_shape print( f"{item.id} infilling: " f"retrieved_latents_shape {retrieved_latents_shape} < " f"expected_latents_shape {expected_latents_shape}. " f"Pad the end." ) history_latents = parent_latents[:parent_from_start_index] future_latents = ( parent_latents[-parent_from_end_index:] if parent_from_end_index > 0 else np.array([]) ) total_latents = np.concatenate( [history_latents, infill_with_pad_latents, future_latents], axis=0 ) if total_latents is not None: print(f"{item.id}: infill full vae latents shape: {total_latents.shape}") np.savez( os.path.join(td, f"{item.id}_vae.npz"), vae_latents=total_latents, v_vae=target_vae_version, ) s3_client.upload_file( os.path.join(td, f"{item.id}_vae.npz"), "suno-data-uploads", f"studio/uploads/{item.id}_vae.npz", ) except Exception as e: print(f"{item.id}: error loading vae latents: {e}") traceback.print_exc() print( f"{item.id}: Processing time:{round(time.time() - start_time, 2)} seconds. Fetched codes." ) if total_latents is None: audio = VERSION_GPT_MAPPING[f"{target_version}"].codec_decode_stream_to_full_audio( concated_codec_tokens, n_stride_tokens=min(n_default_stride_tokens, concated_codec_tokens.shape[0] - 5), n_overlap_tokens=5, ) print( f"{item.id}: Processing time:{round(time.time() - start_time, 2)} seconds. Decoded coarse latents, {round(audio.duration_s, 2)}s" ) else: assert target_vae_version in vae_version_to_decode_fn audio = vae_version_to_decode_fn[target_vae_version]( total_latents, n_stride_tokens=min(n_default_stride_tokens, total_latents.shape[0] - 5), ) print( f"{item.id}: Processing time:{round(time.time() - start_time, 2)} seconds. Decoded vae latents, {round(audio.duration_s, 2)}s" ) if is_diffusion_infill: _ = self.modal_f_cycle.spawn( f"s3://suno-data-uploads/studio/uploads/{item.id}.mp3", s3_npz_id=item.id, encode_vae_version=None, ) else: new_npz = { f"v{target_version}_raw": tokens, "concat_runner": APP_NAME, "model_version": "4.0.0.0", } full_npz_path = os.path.join(td, "concat.npz") np.savez(full_npz_path, **new_npz) print( f"{item.id}: Processing time:{round(time.time() - start_time, 2)} seconds. Finished writing npz." ) s3_client.upload_file( full_npz_path, "suno-data-uploads", f"studio/uploads/{item.id}.npz" ) print( f"{item.id}: Processing time:{round(time.time() - start_time, 2)} seconds. Finished uploading npz." ) self.worker._write_audio_only( item, audio, gain_adjust=get_vae_version_gain_adjust(target_vae_version), ) print( f"{item.id}: Processing time:{round(time.time() - start_time, 2)} seconds. Finished uploading mp3." ) # TODO: note that we use the previous audio's image self.modal_f_video_generation.spawn(item.json(), f"image_{audio_id}.png") if callback_url: item.callback_url = callback_url print(f"{item.id}: Processing time:{round(time.time() - start_time, 2)} seconds. Finished.") item.notify_progress( { "id": item.id, "type": "concat_infilling", "duration_s": audio.duration_s, } ) @app.cls( cpu=4.0, gpu="A10G", # T4s have slower ffmpegs... secrets=SECRETS, timeout=400, # ppl are crazy... scaledown_window=240, retries=modal.Retries( max_retries=2, backoff_coefficient=2.0, initial_delay=5.0, ), max_containers=CONCURRENCY_LIMITS[DEPLOYMENT_TYPE], min_containers=KEEP_WARM[DEPLOYMENT_TYPE], cloud="aws", region="us-east", ) @modal.concurrent(max_inputs=4) class ConcatStub: def __init__(self): self.worker = ConcatWorker() self.worker.preload() @modal.method() def concat_videos_with_history_info( self, new_id: str, history_info: list[str | dict[str, Any]], callback_url: str | None = None ) -> None: ids = [] continue_at_list = [] include_history_list = [] item = QueueItem(id=new_id, metadata={}) if callback_url: item.callback_url = callback_url for h in history_info: if isinstance(h, str): # either it's a clip with old style history, or a last clip which doesn't specify continue_at ids.append(h) continue_at_list.append(None) include_history_list.append(None) elif isinstance(h, dict): ids.append(h["id"]) continue_at = h.get("continue_at", None) # Non positive continue_at is not allowed if continue_at and continue_at <= 0: continue_at = None continue_at_list.append(continue_at) include_history = h.get("include_history_s", None) if include_history and include_history <= 0: include_history = None include_history_list.append(include_history) if any((clip_id is None or clip_id == "None") for clip_id in ids): item.notify_progress( { "id": new_id, "type": "concat_videos", "ok": 0, "error_type": "generation_failure", "error_message": "Can't retrieve history info for all clips.", } ) return # proceed if the info is complete try: duration_s = self.worker.concat_videos(new_id, ids, continue_at_list, include_history_list) print(f"Finished {new_id}, duration_s, {duration_s}") except Exception as e: print("job failed", new_id, e) traceback.print_exc() item.notify_progress( { "id": new_id, "type": "concat_videos", "ok": 0, "error_type": "generation_failure", "error_message": str(e), } ) return if callback_url: item.notify_progress( { "id": new_id, "type": "concat_videos", "duration_s": duration_s, } ) @app.cls( cpu=4.0, secrets=SECRETS, timeout=200, # ppl are crazy... scaledown_window=240, retries=modal.Retries( max_retries=2, backoff_coefficient=2.0, initial_delay=5.0, ), max_containers=CONCURRENCY_LIMITS[DEPLOYMENT_TYPE], min_containers=KEEP_WARM[DEPLOYMENT_TYPE], cloud="aws", region="us-east", ) @modal.concurrent(max_inputs=8) class CropStub(S3Loader): def __init__(self): self.modal_f_video_generation = modal.Cls.lookup( f"videos-v2-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}", "DummyV0Stub", )().write_video @modal.method() def crop_audio(self, queue_item: str) -> None: """Crop the audio and save it into the new id. prompt_audio is the original audio's id. metadata needs to contain two extra keys: crop_start_time: float crop_end_time: float is_crop_remove: boolean """ print(f"Recieved {queue_item}") item = QueueItem(**json.loads(queue_item)) is_remove_crop_part = item.metadata.get("is_crop_remove", False) # crop audio with tempfile.TemporaryDirectory() as td: # download the npz file first npz_filename = os.path.join(td, f"{item.prompt_audio}.npz") with open(npz_filename, "wb") as tmp_file: fast_retry_s3_download( "suno-data-uploads", f"studio/uploads/{item.prompt_audio}.npz", tmp_file ) npz = np.load(npz_filename) # download the audio file audio_filename = os.path.join(td, f"{item.prompt_audio}.mp3") with open(audio_filename, "wb") as tmp_file: fast_retry_s3_download( "suno-data-uploads", f"studio/uploads/{item.prompt_audio}.mp3", tmp_file ) audio = Audio.from_file(audio_filename, n_channels=2) # get the tokens if "model_version" in npz: full_model_version = npz["model_version"].item() arch, major, _, _ = full_model_version.split(".") target_version = f"{arch}.{major}" tokens = npz[f"v{target_version}_raw"] else: full_model_version = "2.0.0.0" target_version = "2.0" tokens = npz["v1_raw"] codec_freq = 25 audio_duration = tokens.shape[0] / codec_freq start_time_s = int(item.metadata.get("crop_start_time", 0) * codec_freq) / codec_freq end_time_s = ( int(item.metadata.get("crop_end_time", audio_duration) * codec_freq) / codec_freq ) if end_time_s - start_time_s <= 1: item.notify_progress( { "id": item.id, "type": "edit_crop", "ok": 0, "error_type": "edit_failure", "error_message": "Can't process the edit request. The cropped duration is too short.", } ) return if end_time_s - start_time_s > audio_duration + 0.04: item.notify_progress( { "id": item.id, "type": "edit_crop", "ok": 0, "error_type": "edit_failure", "error_message": "Can't process the edit request. Please report this issue.", } ) return # trim the tokens if is_remove_crop_part: new_tokens = np.concatenate( ( tokens[: int(start_time_s * codec_freq)], tokens[int(end_time_s * codec_freq) :], ) ) else: new_tokens = tokens[int(start_time_s * codec_freq) : int(end_time_s * codec_freq)] print(f"Filtered start {start_time_s}, end {end_time_s}, new token shape {new_tokens.shape}") # save the tokens if target_version == "2.0": new_npz = { "v1_raw": new_tokens, } else: new_npz = { f"v{target_version}_raw": new_tokens, "crop_parent": item.prompt_audio, "crop_start_time": start_time_s, "crop_end_time": end_time_s, "model_version": full_model_version, } full_npz_path = os.path.join(td, "crop.npz") np.savez(full_npz_path, **new_npz) s3_client.upload_file(full_npz_path, "suno-data-uploads", f"studio/uploads/{item.id}.npz") # crop the vae latents -- always check if they exist latent_file_exists = False if item.prompt_audio is not None: latent_s3_file = f"studio/uploads/{item.prompt_audio}_vae.npz" try: s3_client.get_object(Bucket="suno-data-uploads", Key=latent_s3_file) latent_file_exists = True except: pass if latent_file_exists and item.prompt_audio is not None: print(f"{item.id}: loading vae latents from {item.prompt_audio}") # we need to load the latents for each clip, and crop... # note in this case this is a lazy load -- no vae conversion is done latents, latents_version = get_latents( os.path.join(td, f"{item.prompt_audio}_vae.npz"), item.prompt_audio, target_vae_version="", ) if latents is not None and latents.shape[0] > 0: if is_remove_crop_part: latents = np.concatenate( ( latents[: int(start_time_s * codec_freq)], latents[int(end_time_s * codec_freq) :], ) ) else: latents = latents[ int(start_time_s * codec_freq) : int(end_time_s * codec_freq), : ] print(f"{item.id}: cropped vae latents shape: {latents.shape}") np.savez( os.path.join(td, f"{item.id}_vae.npz"), vae_latents=latents, crop_start_time=start_time_s, crop_end_time=end_time_s, v_vae=latents_version, parent_clip_id=item.prompt_audio.replace(".mp3", ""), ) s3_client.upload_file( os.path.join(td, f"{item.id}_vae.npz"), "suno-data-uploads", f"studio/uploads/{item.id}_vae.npz", ) # crop the audio if is_remove_crop_part: first_audio = audio.get_segment(from_s=0, to_s=start_time_s) second_audio = audio.get_segment(from_s=end_time_s, to_s=audio.duration_s) audio = first_audio.append(second_audio) else: audio = audio.get_segment(from_s=start_time_s, to_s=end_time_s) self._write_audio_only(item, audio) # copy the images self.copy_and_upload_image(item) # find the chunk of lyrics within crop contained_lyrics = "" try: # download the hoot_json file hoot_json_filename = os.path.join(td, f"{item.prompt_audio}_hoot.json") with open(hoot_json_filename, "wb") as tmp_file: fast_retry_s3_download( "suno-data-uploads", f"studio/uploads/{item.prompt_audio}_hoot.json", tmp_file ) with open(hoot_json_filename, "r") as hoot_file: hoot_json = json.load(hoot_file) # this is pre-sorted lyricsls aligned_lyrics = [word for word in hoot_json if "word" in word.keys()] start_index = None end_index = len(aligned_lyrics) for i, word in enumerate(aligned_lyrics): if word.get("start_s", -1) >= start_time_s and start_index is None: start_index = i if word.get("start_s", -1) >= end_time_s: end_index = i break if start_index is not None: if is_remove_crop_part: contained_lyrics = "".join( [aligned_lyrics[i]["word"] for i in range(start_index)] ) + "".join( [ aligned_lyrics[i]["word"] for i in range(end_index + 1, len(aligned_lyrics)) ] ) else: contained_lyrics = "".join( [aligned_lyrics[i]["word"] for i in range(start_index, end_index + 1)] ) print(f"Contained lyrics {item.id}: {contained_lyrics}") except Exception as e: print(f"Failed to find contained lyrics for {item.id}: {e}") contained_lyrics = "" # if for some reason we can't load the lyrics, but the backend passed in prompt # still try to align it...I guess if not contained_lyrics and item.metadata.get("prompt"): contained_lyrics = item.metadata.get("prompt") # will re-generate video and re-do hoot _ = self.modal_f_video_generation.spawn( json.dumps( { "id": item.id, "prompt_text": contained_lyrics, "metadata": {"tags": " "}, } ), f"image_{item.id}.png", ) duration = end_time_s - start_time_s if is_remove_crop_part: duration = audio_duration - duration item.notify_progress( { "id": item.id, "source_clip_id": item.metadata.get("source_clip_id", item.prompt_audio), "type": "edit_crop", "duration": duration, "audio_s3_id": item.id, "image_s3_id": f"image_{item.id}", "contained_lyrics": contained_lyrics, } ) print(f"Done with {item.id}") return @app.cls( cpu=2.0, secrets=SECRETS, timeout=200, # ppl are crazy... scaledown_window=240, retries=modal.Retries( max_retries=2, backoff_coefficient=2.0, initial_delay=5.0, ), max_containers=CONCURRENCY_LIMITS[DEPLOYMENT_TYPE], min_containers=KEEP_WARM[DEPLOYMENT_TYPE], cloud="aws", region="us-east", ) @modal.concurrent(max_inputs=8) class FadeStub(S3Loader): def __init__(self): self.modal_f_video_generation = modal.Cls.lookup( f"videos-v2-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}", "DummyV0Stub", )().write_video self.modal_f_encode_audio = modal.Cls.lookup( f"cycle-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}", "CycleStub", )().encode_audio @modal.method() def fade_audio(self, queue_item: str) -> None: """Fade out the audio and save it into the new id. prompt_audio is the original audio's id. metadata needs to contain two extra keys: fade_out_time: float Optionally can also have: fade_out_shape: Literal['q', 't', 'l']. Default is 'q' (q = quarter sine, t = linear, l = logarithmic. Refer to pysox Docs) """ print(f"Recieved {queue_item}") item = QueueItem(**json.loads(queue_item)) # fade audio with tempfile.TemporaryDirectory() as td: # download the audio file audio_filename = os.path.join(td, f"{item.prompt_audio}.mp3") with open(audio_filename, "wb") as tmp_file: fast_retry_s3_download( "suno-data-uploads", f"studio/uploads/{item.prompt_audio}.mp3", tmp_file ) audio = Audio.from_file(audio_filename, n_channels=2) audio = fade_out_audio( audio_filename, fade_out_len=audio.duration_s - item.metadata.get("fade_out_time", audio.duration_s), fade_out_shape=item.metadata.get("fade_out_shape", "q"), sample_rate=audio.sample_rate, byte_width=int(audio.bit_depth / 8), n_channels=audio.n_channels, ) self._write_audio_only(item, audio) # encode the audio back to codes # use the same vae version as the original if possible self.modal_f_encode_audio.spawn( audio=audio, s3_npz_id=item.id, encode_vae_version=_get_clip_vae_version(item.prompt_audio), ) # copy the images self.copy_and_upload_image(item) # will re-generate video and re-do hoot _ = self.modal_f_video_generation.spawn( json.dumps( { "id": item.id, "prompt_text": item.metadata.get("prompt"), "metadata": {"tags": " "}, } ), f"image_{item.id}.png", ) if item.callback_url: item.notify_progress( { "id": item.id, "source_clip_id": item.metadata.get("source_clip_id", item.prompt_audio), "type": "edit_fade", "duration": item.metadata.get("duration"), "audio_s3_id": item.id, "image_s3_id": f"image_{item.id}", }, ) print(f"Done with {item.id}") return @app.cls( cpu=2.0, secrets=SECRETS, timeout=200, # ppl are crazy... scaledown_window=240, retries=modal.Retries( max_retries=2, backoff_coefficient=2.0, initial_delay=5.0, ), max_containers=CONCURRENCY_LIMITS[DEPLOYMENT_TYPE], min_containers=0, cloud="aws", region="us-east", ) @modal.concurrent(max_inputs=8) class ChangeSpeedStub(S3Loader): def __init__(self): self.modal_f_video_generation = modal.Cls.lookup( f"videos-v2-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}", "DummyV0Stub", )().write_video self.modal_f_encode_audio = modal.Cls.lookup( f"cycle-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}", "CycleStub", )().encode_audio @modal.method() def change_audio_speed(self, queue_item: str) -> None: """Change the audio speed and save it into the new id. prompt_audio is the original audio's id. metadata needs to contain one extra key: speed_factor: float """ print(f"Recieved {queue_item}") item = QueueItem(**json.loads(queue_item)) final_duration = None # change audio speed with tempfile.TemporaryDirectory() as td: # download the audio file audio_filename = os.path.join(td, f"{item.prompt_audio}.mp3") with open(audio_filename, "wb") as tmp_file: fast_retry_s3_download( "suno-data-uploads", f"studio/uploads/{item.prompt_audio}.mp3", tmp_file ) audio = Audio.from_file(audio_filename, n_channels=2) audio_prefix = audio.get_segment(0, item.metadata.get("change_speed_start_time")) audio_to_transform = audio.get_segment( item.metadata.get("change_speed_start_time"), item.metadata.get("change_speed_end_time"), ) audio_postfix = audio.get_segment( item.metadata.get("change_speed_end_time"), audio.duration_s ) segment_wav_path = os.path.join(td, "segment.wav") audio_to_transform.write_wav(segment_wav_path) audio_to_transform = change_audio_speed( audio_filepath=segment_wav_path, speed_factor=item.metadata.get("speed_factor", 1.0), tempo_only=item.metadata.get("tempo_only", False), sample_rate=audio.sample_rate, byte_width=int(audio.bit_depth / 8), n_channels=audio.n_channels, ) full_audio = audio_prefix.append(audio_to_transform).append(audio_postfix) final_duration = full_audio.duration_s self._write_audio_only(item, full_audio) # encode the audio back to codes # use the same vae version as the original if possible self.modal_f_encode_audio.spawn( audio=full_audio, s3_npz_id=item.id, encode_vae_version=_get_clip_vae_version(item.prompt_audio), ) # copy the images self.copy_and_upload_image(item) # will re-generate video and re-do hoot _ = self.modal_f_video_generation.spawn( json.dumps( { "id": item.id, "prompt_text": item.metadata.get("prompt"), "metadata": {"tags": " "}, } ), f"image_{item.id}.png", ) item.notify_progress( { "id": item.id, "source_clip_id": item.metadata.get("source_clip_id", item.prompt_audio), "type": "edit_speed", "duration": final_duration, "audio_s3_id": item.id, "image_s3_id": f"image_{item.id}", }, ) print(f"Done with {item.id}") return def _test_concat_without_continue_at(app: ConcatStub) -> None: new_id = str(uuid.uuid4()) app.concat_videos_with_history_info.remote( new_id, ["ea69bf30-30c6-497e-8949-49b20822f74c", "08dfecc8-f429-4d1b-b2e3-e23d54cf35e9"], ) print( f"_test_concat_without_continue_at created video: s3://suno-data-uploads/studio/uploads/{new_id}.mp3" ) def _test_concat_with_continue_at(app: ConcatStub) -> None: new_id = str(uuid.uuid4()) # app.concat_videos_with_history_info.remote( # new_id, # [ # {"id": "2d4a9162-3f12-4ddd-9c8e-450a6c4b2086", "continue_at": None}, # {"id": "787b3bb9-2083-4dd2-a995-a133031425eb", "continue_at": None}, # {"id": "ca857a4b-856b-4da1-a457-aa144947b1c1", "continue_at": 31.72}, # {"id": "1163031b-8d4b-41e5-be6c-8c9330616136", "continue_at": 34.48}, # {"id": "204e1a10-9608-4c3c-a975-9d284601b04f", "continue_at": None}, # ], # ) # print( # f"_test_concat_with_continue_at created video: s3://suno-data-uploads/studio/uploads/{new_id}.mp3" # ) # this should trigger the hybrid concat app.concat_videos_with_history_info.remote( new_id, [ {"id": "7deb85a1-9ef3-4c83-87c1-b9638fa5ad1c", "continue_at": 78.0}, {"id": "482d4b62-fd4b-4411-a3d2-47bf94de180c", "continue_at": None}, ], ) print( f"_test_concat_with_continue_at created video: s3://suno-data-uploads/studio/uploads/{new_id}.mp3" ) def _test_crop_audio(app: CropStub) -> None: crop_input = json.dumps( { "id": "9a9718b5-c8d2-4086-92e6-d2595d8d7a90_crop", "prompt_audio": "9a9718b5-c8d2-4086-92e6-d2595d8d7a90", "model_name": "chirp-v3p5", "metadata": { "crop_start_time": 1.2, "crop_end_time": 10.5, "image_s3_id": "image_9a9718b5-c8d2-4086-92e6-d2595d8d7a90", }, } ) app.crop_audio.remote(crop_input) crop_input = json.dumps( { "id": "fe406f51-4023-41ac-958b-d97c6f22e02b", "prompt_audio": "ef47091f-dd47-4c27-a0ca-006260aa6869", "metadata": { "crop_start_time": 0.0, "crop_end_time": 61.4, "image_s3_id": "image_3b617e13-c441-413e-a606-9ced85d12d46", }, } ) app.crop_audio.remote(crop_input) upload_crop_input = json.dumps( { "id": "6c8dab6c-fad9-4d6d-b8b8-4e85a1a08baa", "prompt_audio": "m_255fd292-0139-4b96-b242-018b9e4e5204", "prompt_npz": None, "prompt_text": None, "metadata": { "crop_start_time": 0.6819334049409237, "crop_end_time": 6.122735409953455, "image_s3_id": "image_255fd292-0139-4b96-b242-018b9e4e5204", "source_clip_id": "255fd292-0139-4b96-b242-018b9e4e5204", }, "gen_duration": 12, "model_name": "chirp-v3p5", "title": None, "ids": None, } ) app.crop_audio.remote(upload_crop_input) def _test_infill_audio(app: ConcatInfillingStub) -> None: concat_infill_input = json.dumps( { "id": "413cb2f5-b634-4ab7-b8c9-31ae202dd2da", "prompt_text": "flying down the streets feeling so alive i've got my head ", "prompt_audio": "aed67b71-8857-4891-b497-80b8dc3bbbc7", "metadata": {}, } ) app.concat_infilling_with_queue_item.remote(concat_infill_input) # test parent_from_end_index = 0 case concat_infill_input = json.dumps( { "id": "ae4e79eb-a76d-4a18-9b77-de9ab6139656", "prompt_text": " ", "prompt_audio": "84a67e3f-14a8-4f99-972b-1b3d29d72acd", "metadata": {}, } ) app.concat_infilling_with_queue_item.remote(concat_infill_input) def _test_change_speed(app: ChangeSpeedStub) -> None: change_speed_input = json.dumps( { "id": "58d69cc1-caf0-42fb-b61c-ceae1a963e56", "prompt_audio": "9edce783-5c24-49d5-9332-4a682940b164", "prompt_npz": None, "prompt_text": None, "metadata": { "speed_factor": 1.1, "tempo_only": True, "change_speed_start_time": 0.0, "change_speed_end_time": 15.0, "image_s3_id": "image_9edce783-5c24-49d5-9332-4a682940b164", "source_clip_id": "9edce783-5c24-49d5-9332-4a682940b164", }, "gen_duration": 12, "callback_url": "https://studio-api.prod.suno.com/api/edit/webhook/finish/", "model_name": "chirp-auk-t1", "title": None, "ids": None, } ) app.change_audio_speed.remote(change_speed_input) @app.local_entrypoint() def main(): """Called by modal run, for debugging.""" concat_model = ConcatStub() # (Processing time: 10 seconds) _test_concat_without_continue_at(concat_model) # (Processing time: 85 seconds) _test_concat_with_continue_at(concat_model) crop_model = CropStub() _test_crop_audio(crop_model) infill_model = ConcatInfillingStub() _test_infill_audio(infill_model) change_speed_model = ChangeSpeedStub() _test_change_speed(change_speed_model)