import time import os import tempfile import ffmpeg import numpy as np from typing import Any import uuid import json import traceback import modal from suno_utils.tasks.gpt.generation import * from suno_utils.gpt import chirp_v2 from suno_utils.gpt import chirp_v2_5 from suno_utils.audio import Audio import subprocess as sp from suno_utils.worker.loader import VERSION_GPT_MAPPING, S3Loader, get_tokens from suno_utils.worker.utils import download_models_to_dir from suno_utils.worker.modal_base import get_modal_base_image, MODAL_MOUNTS from suno_utils.worker.schema import QueueItem from suno_utils.worker.settings import s3_client from suno_utils.worker.utils import retry_decorator ############## CHANGE THESE ############## DEPLOYMENT_TYPE = "dev" # dev, prod ########################################## ENCODER_MAX_INPUT = 4 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", ) MOUNT_PATH = "/suno/models" retry_s3_download = retry_decorator(3, wait_seconds=20)(s3_client.download_fileobj) 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, bg_image=None): super().__init__(bg_image=bg_image) self.bg_image = bg_image self.modal_f_video_generation = modal.Function.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']}", ) 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()) 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 ): print("request new_id", new_id) print(ids, continue_at_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) 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) 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: 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": STUB_NAME, "model_version": full_model_version, } full_npz_path = os.path.join(td, "concat.npz") np.savez(full_npz_path, **new_npz) concated_codec_tokens = np.concatenate(tokens_list, axis=0)[:, 1:] estimated_duration = concated_codec_tokens.shape[0] / 25 print("estimated audio duration", estimated_duration) if estimated_duration < 1: raise ValueError("The estimated duration is less than 1 second.") 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 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, ) 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("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"Total audio length is {audio.duration_s}.") audio = audio.normalize_volume(-12) # this is only used for video generation full_wav_path = os.path.join(td, "concat.wav") audio.write_wav(full_wav_path) full_mp3_path = os.path.join(td, "concat.mp3") audio.write_hq_mp3(full_mp3_path) # download all the videos and get the trimmed videos filenames = [f"{id}.mp4" for id in ids] full_paths = [] for fname in filenames: full_path = os.path.join(td, fname) full_paths.append(full_path) with open(full_path, "wb") as tmp_file: try: s3_client.download_fileobj( "suno-data-uploads", f"studio/uploads/{fname}", tmp_file, ) except: print("Failed to download", fname, "will try to regenerate.") clip_id = fname.replace(".mp4", "") # Note that this is a dummy call. fn_calls = self.modal_f_video_generation.spawn( json.dumps( { "id": clip_id, "prompt_text": "", "metadata": {"tags": " "}, } ), f"image_{clip_id}.png", ) try: _ = fn_calls.get(timeout=60) print("finished regenerating", fname) # then retry the download retry_s3_download( "suno-data-uploads", f"studio/uploads/{fname}", tmp_file, ) except Exception as e: print("Failed to regenerate video", fname, e) # get the merged target width, height, and frame per second dimension_and_fps = [get_dimensions_and_fps(i) for i in full_paths] target_w = max([x[0] for x in dimension_and_fps]) target_h = max([x[1] for x in dimension_and_fps]) target_fps = min([x[2] for x in dimension_and_fps]) rescaled_inputs = [] for p, continue_at in zip(full_paths, continue_at_list): input = ffmpeg.input(p).filter("scale", target_w, target_h).filter("fps", target_fps) if continue_at: input = input.trim(end=continue_at) rescaled_inputs.append(input) full_path = os.path.join(td, "concat.mp4") joined = ffmpeg.concat(*rescaled_inputs, v=1, a=0) ffmpeg.output(joined, full_path).run(quiet=True) # then concact the audio full_path_2 = os.path.join(td, "concat_2.mp4") sp.run( [ "ffmpeg", "-i", full_path, "-i", full_wav_path, "-c:v", "copy", "-c:a", "aac", "-b:a", "192k", "-map", "0:v", "-map", "1:a", "-shortest", full_path_2, ] ) # ffmpeg.concat(ffmpeg.input(full_path), ffmpeg.input(full_audio_path), v=1, a=1).output( # full_path_2, # ac=2, # c_audio="aac", # audio_bitrate="192k", # ).run(quiet=True) # finish and upload all the concacted info s3_client.upload_file( full_path_2, "suno-data-uploads", f"studio/uploads/{new_id}.mp4", ExtraArgs={ "ContentType": "video/mp4", }, ) s3_client.upload_file( full_mp3_path, "suno-data-uploads", f"studio/uploads/{new_id}.mp3", ExtraArgs={ "ContentType": "audio/mp3", }, ) s3_client.upload_file(full_npz_path, "suno-data-uploads", f"studio/uploads/{new_id}.npz") return audio.duration_s def download_model_wrapper_5(): print("Downloading models...") ConcatWorker.download_models() image = get_modal_base_image().run_function( download_model_wrapper_5, secret=modal.Secret.from_name("studio-aws") ) STUB_NAME = f"concat-v3-{DEPLOYMENT_TYPE}" stub = modal.Stub(STUB_NAME, image=image) @stub.cls( cpu=2.0, gpu=modal.gpu.A10G(count=1), secrets=[ modal.Secret.from_name("studio-aws"), modal.Secret.from_dict({"SUNO_ASSETS_PATH": "/suno/models/assets"}), ], timeout=400, # ppl are crazy... container_idle_timeout=240, mounts=MODAL_MOUNTS, retries=modal.Retries( max_retries=2, backoff_coefficient=2.0, initial_delay=5.0, ), concurrency_limit=CONCURRENCY_LIMITS[DEPLOYMENT_TYPE], allow_concurrent_inputs=ENCODER_MAX_INPUT, ) class ConcatStub: def __enter__(self): self.worker = ConcatWorker(bg_image="/suno/models/assets/wave-bg-3.png") self.worker.preload() # don't need to keep this one around... @modal.method(keep_warm=1) def concat_infilling_with_queue_item(self, queue_item: str, callback_url: str | None = None): """For inflling we do something special.""" if DEPLOYMENT_TYPE == "dev": print(f"Recieved {queue_item}") item = QueueItem(**json.loads(queue_item)) # 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 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") if tokens is None: raise ValueError(f"Can't load item {item.id}, audio {audio_id}") concated_codec_tokens = tokens[:, 1:] estimated_duration = concated_codec_tokens.shape[0] / 25 print("estimated audio duration", estimated_duration) if estimated_duration < 1: raise ValueError("The estimated duration is less than 1 second.") 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 target_version = "4.0" # this is fixed for now 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, ) new_npz = { f"v{target_version}_raw": tokens, "concat_runner": STUB_NAME, "model_version": "4.0.0.0", } full_npz_path = os.path.join(td, "concat.npz") np.savez(full_npz_path, **new_npz) full_mp3_path = os.path.join(td, "concat.mp3") audio.write_hq_mp3(full_mp3_path) s3_client.upload_file( full_mp3_path, "suno-data-uploads", f"studio/uploads/{item.id}.mp3", ExtraArgs={ "ContentType": "audio/mp3", }, ) s3_client.upload_file(full_npz_path, "suno-data-uploads", f"studio/uploads/{item.id}.npz") # 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: self.worker.notify_finish( QueueItem(id=item.id, metadata={}, callback_url=callback_url), { "id": item.id, "type": "concat_infilling", "duration_s": estimated_duration, }, ) @modal.method(keep_warm=KEEP_WARM[DEPLOYMENT_TYPE]) 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 = [] 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) 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) if any((clip_id is None or clip_id == "None") for clip_id in ids): if callback_url: self.worker.notify_finish( QueueItem(id=new_id, metadata={}, callback_url=callback_url), { "id": new_id, "type": "concat_videos", "ok": 0, "error_type": "generation_failure", "error_message": "Can't retrieve history info for all clips.", }, queue_name="results:q", ) return else: print("Can't concat with bad ids and no callback url", ids) raise ValueError("Can't retrieve history info for all clips.") # proceed if the info is complete try: duration_s = self.worker.concat_videos( new_id, ids, continue_at_list, ) print(f"Finished {new_id}, duration_s, {duration_s}") except Exception as e: print("job failed", new_id, e) traceback.print_exc() if callback_url: self.worker.notify_finish( QueueItem(id=new_id, metadata={}, callback_url=callback_url), { "id": new_id, "type": "concat_videos", "ok": 0, "error_type": "generation_failure", "error_message": str(e), }, queue_name="results:q", ) return if callback_url: self.worker.notify_finish( QueueItem(id=new_id, metadata={}, callback_url=callback_url), { "id": new_id, "type": "concat_videos", "duration_s": duration_s, }, ) def _test_concat_without_continue_at(stub: ConcatStub) -> None: new_id = str(uuid.uuid4()) stub.concat_videos.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}.mp4" ) def _test_concat_with_continue_at(stub: ConcatStub) -> None: new_id = str(uuid.uuid4()) stub.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}.mp4" ) @stub.local_entrypoint() def main(): """Called by modal run, for debugging.""" model = ConcatStub() _test_concat_without_continue_at(model) _test_concat_with_continue_at(model)