""" This is the history audio loader. It will do the following: 1. Load the history audio and text 2. If needed, convert the history audio to the correct coarse tokens as the model requires 3. Hoot the history audio and text for alignment 4. Spin up the chirp worker as required by the queue item This worker requires GPU for faster encoding and hoot. """ import json import os import time import traceback from typing import Optional, Tuple import math import modal import numpy as np import torch 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.hoot import encode as hoot_encode from suno_utils.worker.loader import S3Loader from suno_utils.worker.modal_model_configs import get_vae_version, is_vae_model from suno_utils.worker.schema import HistoryPrompt, QueueItem from suno_utils.worker.settings import s3_client from suno_utils.worker.utils import print_gpu_memory_usage from suno_utils.worker.modal_base import get_modal_base_image_with_flash_attention from suno_utils.worker.tracing import serialize_context, distributed_trace, tracer from suno_utils.worker.modal_model_volume import MODEL_STORE_VOLUME_DIR, model_store_volume ############## CHANGE THESE ############## DEPLOYMENT_TYPE = "dev" # dev, prod ########################################## assert DEPLOYMENT_TYPE in {"dev", "prod"}, DEPLOYMENT_TYPE ENCODER_CONCURRENCY_LIMITS = { "dev": 5, "prod": 125, } KEEP_WARM = { "dev": 1, "prod": 4, } # set number of cpus. ENCODER_MAX_INPUT = 4 # encoder can handle much more traffic, but let's be conservative... 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_dict( { "DD_SITE": "datadoghq.com", "DD_ENV": DEPLOYMENT_TYPE, "DD_SERVICE": "history_loader", "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"), ] base_image = get_modal_base_image_with_flash_attention() class CodecWorker(S3Loader): def __init__(self): S3Loader.__init__(self) 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) hoot_ckpt_path = chirp_v3._get_model_if_needed(chirp_v3.HOOT_CKPT_PATH, cache_dir=MOUNT_PATH) hoot_tokenizer_path = chirp_v3._get_model_if_needed( chirp_v3.HOOT_TOKENIZER_PATH, cache_dir=MOUNT_PATH ) # load hoot on cuda as well chirp_v3.preload_hoot_models(hoot_ckpt_path, hoot_tokenizer_path) print("Finish loading models") @staticmethod def download_models(dir_path=MOUNT_PATH): """Download the codec models for dac8, dac12 and hoot.""" print("Start downloading models") os.makedirs(dir_path, exist_ok=True) chirp_v2._get_model_if_needed(chirp_v2.CODEC_CKPT_PATH, cache_dir=dir_path) chirp_v2.preload_codec_models(chirp_v2.CODEC_CKPT_PATH) chirp_v3._get_model_if_needed(chirp_v3.CODEC_CKPT_PATH, cache_dir=dir_path) chirp_v3.preload_codec_models(chirp_v3.CODEC_CKPT_PATH) hoot_ckpt_path = chirp_v3._get_model_if_needed(chirp_v3.HOOT_CKPT_PATH, cache_dir=dir_path) hoot_tokenizer_path = chirp_v3._get_model_if_needed( chirp_v3.HOOT_TOKENIZER_PATH, cache_dir=dir_path ) chirp_v3.preload_hoot_models(hoot_ckpt_path, hoot_tokenizer_path) print("Finish downloading models") def download_model_wrapper_b(): # 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") CodecWorker.download_models() image = base_image.run_function( download_model_wrapper_b, secrets=SECRETS, volumes={MODEL_STORE_VOLUME_DIR: model_store_volume} ).add_local_python_source("_remote_module_non_scriptable", "suno_utils") APP_NAME = f"history_encoder-{DEPLOYMENT_TYPE}" app = modal.App(APP_NAME, image=image) HISTORY_DDOG_SERVICE = "history_loader" @app.cls( gpu="A10G", secrets=SECRETS, timeout=600, scaledown_window=400, memory=15000, retries=modal.Retries( max_retries=2, backoff_coefficient=2.0, initial_delay=5.0, ), max_containers=ENCODER_CONCURRENCY_LIMITS[DEPLOYMENT_TYPE], min_containers=KEEP_WARM[DEPLOYMENT_TYPE], region="us-east", volumes={MODEL_STORE_VOLUME_DIR: model_store_volume}, ) @modal.concurrent(max_inputs=ENCODER_MAX_INPUT) class HistoryEncoderStub: def __init__(self): import torch num_gpus = torch.cuda.device_count() print(f"Found {num_gpus} GPUs.") self.worker = CodecWorker() self.modal_f_orchestrator = modal.Cls.lookup( f"orchestrator-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}", "ConductorStub", )().redirect_and_generate self.modal_f_encode_audio_to_vae = modal.Cls.lookup( f"cycle-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}", "VaeCycleStub", )().encode_audio_to_vae_latents self.modal_f_cycle = modal.Cls.lookup( f"cycle-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}", "CycleStub", )().encode_audio def print_info(self, queue_item_id: str, message: str): print(f"HistoryLoader ({queue_item_id}): {message}") def _get_model_version_and_history_max_duration(self, item: QueueItem) -> Tuple[str, int]: """Get the model full version and the history max duration.""" # note that this name is from queue item and the FE input model name if item.model_name is None: raise ValueError("Model name is None") if "v2" in item.model_name: return "2.0.0.0", 20 # this is kinda stupid # but for upsamples, we only need semantic tokens # and semantic tokens are the same signature as v5 # so we should try to load 5.0 first # if 5.0 exists, it will load semantic only # if 5.0 does not exist, it will load v4.0 / etc's semantic elif "up" in item.model_name and "upload" not in item.model_name or "ahi" in item.model_name: return "5.0.0.0", 120 elif "6b" in item.model_name or "auk" in item.model_name or "bluejay" in item.model_name: # AUK infill is a special case -- we are still using 30b -- we will do the mapping here if item.is_infill: return "4.0.0.0", 120 return "5.0.0.0", 240 elif "stem" in item.model_name: return "5.0.0.0", 20 elif "v3-5" in item.model_name or "v3p5" in item.model_name: return "3.5.0.0", 120 elif "v3" in item.model_name: return "3.0.0.0", 60 elif "v4" in item.model_name: return "4.0.0.0", 120 else: raise ValueError(f"Unknown model {item.model_name}") def _vae_freq(self) -> int: """TODO: make this dynamic based on vae models.""" return 25 def _vae_context_len(self) -> int: """TODO: make this dynamic based on vae models.""" return 30 def _encode_history(self, item: QueueItem) -> HistoryPrompt: # Model version, used for logging and compatibility checks. # Version format: arch.major.minor.ft # arch: for model architecture changes. i.e. 2 for 3b model, 3 for 7b model, etc. # major: major model version, for changes which WILL cause model input/output structure changes. # minor: minor model version, for other changes which will NOT cause model input/output structure changes. # ft: fine-tune version # TODO: this needs to get better...probably through item attributes model_full_version, history_max_duration = self._get_model_version_and_history_max_duration(item) cover_audio_tokens = None artist_audio_tokens = None playlist_audio_tokens = None multi_artist_audio_tokens = None underpainting_audio_tokens = None overpainting_audio_tokens = None self.print_info(item.id, f"clip_ids: {item.ids}") # TODO: if the parent npz does not exist, we need to cycle it probably if item.is_cover_condition: with tracer.trace("load_cover_audio_prompt"): cover_audio_tokens = self.worker._load_special_audio_prompt( item, model_full_version, is_cover=True ) # we trim and keep the start rather than keep the end cover_audio_tokens = cover_audio_tokens[: int(history_max_duration * 25)] self.print_info( item.id, f"is_cover_condition, trimmed audio npz prompt {cover_audio_tokens.shape}" ) if item.is_artist_condition: with tracer.trace("load_artist_audio_prompt"): artist_audio_tokens = self.worker._load_special_audio_prompt( item, model_full_version, is_artist=True ) # we trim and keep the start rather than keep the end artist_audio_tokens = artist_audio_tokens[: int(history_max_duration * 25)] self.print_info( item.id, f"is_artist_condition, trimmed audio npz prompt {artist_audio_tokens.shape}" ) if item.is_underpainting: with tracer.trace("load_underpainting_audio_prompt"): underpainting_audio_tokens = self.worker._load_special_audio_prompt( item, model_full_version, is_underpainting=True ) underpainting_audio_tokens = underpainting_audio_tokens[: int(history_max_duration * 25)] self.print_info( item.id, f"is_underpainting, trimmed audio npz prompt {underpainting_audio_tokens.shape}", ) if item.is_overpainting: with tracer.trace("load_overpainting_audio_prompt"): overpainting_audio_tokens = self.worker._load_special_audio_prompt( item, model_full_version, is_overpainting=True ) overpainting_audio_tokens = overpainting_audio_tokens[: int(history_max_duration * 25)] self.print_info( item.id, f"is_overpainting, trimmed audio npz prompt {overpainting_audio_tokens.shape}", ) if item.is_playlist_condition: with tracer.trace("load_playlist_audio_prompt"): playlist_audio_tokens = self.worker._load_list_of_audio_prompts( item, model_full_version, "playlist_clip_ids" ) # keep the personnas 60s each for now playlist_audio_tokens = [t[: int(60 * 25)] for t in playlist_audio_tokens] # we trim and keep the start rather than keep the end self.print_info( item.id, f"is_playlist_condition, trimmed audio npz prompt {[t.shape for t in playlist_audio_tokens]}", ) assert len(playlist_audio_tokens) == len(item.metadata["playlist_clip_ids"]) if item.is_multi_artist_consistency: with tracer.trace("load_multi_artist_audio_prompt"): multi_artist_audio_tokens = self.worker._load_list_of_audio_prompts( item, model_full_version, "artist_clip_ids" ) # keep the personnas 60s each for now multi_artist_audio_tokens = [t[: int(60 * 25)] for t in multi_artist_audio_tokens] # we trim and keep the start rather than keep the end self.print_info( item.id, f"is_multi_artist_consistency, trimmed audio npz prompt {[t.shape for t in multi_artist_audio_tokens]}", ) assert len(multi_artist_audio_tokens) == len(item.metadata["artist_clip_ids"]) # immediately crop in case they are too long if item.is_artist_cover_condition or item.is_artist_cover_extend: # we need to trade off the artist and cover duration here assert cover_audio_tokens is not None and artist_audio_tokens is not None n_cover_tokens = cover_audio_tokens.shape[0] n_artist_tokens = artist_audio_tokens.shape[0] self.print_info( item.id, f"artist cover condition, pre trim -- n_cover_tokens: {n_cover_tokens}, n_artist_tokens: {n_artist_tokens}", ) max_condition_tokens = history_max_duration * 25 if n_cover_tokens + n_artist_tokens > max_condition_tokens: # keep one mins artist token first n_artist_tokens = min(25 * 60, n_artist_tokens) # keep the rest for cover n_cover_tokens = max_condition_tokens - n_artist_tokens cover_audio_tokens = cover_audio_tokens[:n_cover_tokens] artist_audio_tokens = artist_audio_tokens[:n_artist_tokens] self.print_info( item.id, f"artist cover condition, after trim -- n_cover_tokens: {n_cover_tokens}, n_artist_tokens: {n_artist_tokens}", ) if item.is_artist_extend or item.is_cover_extend or item.is_artist_cover_extend: self.print_info( item.id, "artist extend or cover extend, reducing history audio to half of the max history duration", ) history_max_duration = history_max_duration // 2 # if continue at is not set it will load all the tokens with tracer.trace("load_audio_prompt"): parent_audio_tokens = ( self.worker._load_audio_prompt(item, model_full_version) if item.prompt_audio else None ) if item.prompt_audio and parent_audio_tokens is None: raise ValueError("Can't load audio prompt for: %s. Need to retry." % item.prompt_audio) history_text = item.metadata.get("continued_from_prompt", "") if item.prompt_audio else "" history_latents = None future_latents = None # for most cases, we don't even want to load the history VAE # TODO: redo the history VAE loading logic if ( isinstance(item.model_name, str) # TODO: categorize this better and is_vae_model(item.model_name) and item.prompt_audio is not None # we don't need to load history VAE for upsample -- waste of resources and (not item.is_upsample) ): prompt_audio_s3_id = item.prompt_audio.replace(".mp3", "") target_vae_version = get_vae_version(item.model_name) self.print_info(item.id, f"{item.model_name}, target_vae_version: {target_vae_version}") if item.prompt_audio is not None: latent_s3_file = f"studio/uploads/{prompt_audio_s3_id}_vae.npz" try: with tracer.trace("load_history_vae"): s3_client.get_object(Bucket="suno-data-uploads", Key=latent_s3_file) all_latents = self.worker._load_history_vae(item, target_vae_version) if all_latents is None: self.print_info(item.id, "Failed to load history VAE.") else: self.print_info(item.id, f"Loaded history VAE {all_latents.shape}.") except: s3_new_path = f"studio/uploads/{prompt_audio_s3_id}.mp3" self.print_info(item.id, "encoding prompt audio to VAE latents") with tracer.trace("encode_audio_to_vae"): # TODO: add parent context to this call vae_f_call = self.modal_f_encode_audio_to_vae.spawn( audio_s3_path=f"s3://suno-data-uploads/{s3_new_path}", s3_npz_id=prompt_audio_s3_id, encode_vae_version=target_vae_version, return_vae_latents=True, ) all_latents = vae_f_call.get(timeout=60) if all_latents is not None: print(f"History loader {item.id}: loaded history VAE {all_latents.shape}") if (continue_at := item.metadata.get("continue_at", None)) is not None: history_latents = all_latents[: int(continue_at * self._vae_freq())] elif (infill_start_s := item.metadata.get("infill_start_s", None)) is not None: history_latents = all_latents[: int(infill_start_s * self._vae_freq())] if ( infill_context_start_s := item.metadata.get("infill_context_start_s", None) ) is not None: history_latents = all_latents[ int(infill_context_start_s * self._vae_freq()) : int( infill_start_s * self._vae_freq() ) ] else: # for other tasks -- like stem, we just use the whole history history_latents = all_latents # fetch the future latents as well if (infill_end_s := item.metadata.get("infill_end_s", None)) is not None and ( infill_context_end_s := item.metadata.get("infill_context_end_s", None) ) is not None: future_latents = all_latents[ int(infill_end_s * self._vae_freq()) : int( infill_context_end_s * self._vae_freq() ) ] if not item.is_stem and history_latents is not None: # crop it to max 1 min to avoid passing too much data -- but keep the end history_latents = history_latents[-self._vae_context_len() * self._vae_freq() :, :] self.print_info( item.id, f"""Item.is_stem: {item.is_stem}, finalized loaded history VAE {history_latents.shape}, future VAE {future_latents.shape if future_latents is not None else 0}""", ) if isinstance(parent_audio_tokens, np.ndarray): # NPZ self.print_info(item.id, f"loaded audio npz prompt {parent_audio_tokens.shape}") if parent_audio_tokens.shape[0] == 0: self.print_info(item.id, f"history_audio is empty from input {item.metadata}.") raise ValueError("History audio is empty") estimated_history_audio_duration = parent_audio_tokens.shape[0] / 25 infill_start_s = item.metadata.get("infill_start_s", None) infill_end_s = item.metadata.get("infill_end_s", None) infill_context_start_s = item.metadata.get("infill_context_start_s", None) infill_context_end_s = item.metadata.get("infill_context_end_s", None) if item.is_infill: assert (infill_start_s is not None) or (infill_end_s is not None) # some validations if infill_context_start_s > infill_start_s or infill_context_end_s < infill_end_s: raise ValueError( f"Infilling times are wrong: {infill_context_start_s}, {infill_start_s}, {infill_end_s}, {infill_context_end_s}" ) else: self.print_info( item.id, f"""Infilling times are: Context start: {infill_context_start_s} Infill start: {infill_start_s} Infill end: {infill_end_s} Context end: {infill_context_end_s} """, ) context_duration = (infill_context_end_s or estimated_history_audio_duration) - ( infill_context_start_s or 0 ) if context_duration > 4 * 60: raise ValueError(f"Maxed out infilling context window: {context_duration}") with tracer.trace("load_infilling_audio_prompt"): infilling_audio_prompts = self.worker._load_infilling_audio_prompt( item, model_full_version, context_start_s=infill_context_start_s, start_s=infill_start_s, end_s=infill_end_s, context_end_s=infill_context_end_s, ) if infilling_audio_prompts is None: raise ValueError(f"Can't load the audio prompt for {item.id}") ( pre_prompt_audio, parent_audio_tokens, future_audio, post_prompt_audio, ) = infilling_audio_prompts self.print_info( item.id, f"""trimmed audio npz prmpts: pre_prompt: {pre_prompt_audio.shape} history: {parent_audio_tokens.shape} future: {future_audio.shape} post_future: {post_prompt_audio.shape}""", ) # note that for infilling, we don't change the history text -- passed in from FE correctly return HistoryPrompt( prompt_audio=parent_audio_tokens if parent_audio_tokens.shape[0] > 0 else None, future_audio=future_audio if future_audio.shape[0] > 0 else None, pre_history_arr=pre_prompt_audio if pre_prompt_audio.shape[0] > 0 else None, post_future_arr=post_prompt_audio if post_prompt_audio.shape[0] > 0 else None, prompt_lyrics="", # for infilling history text is empty cover_audio=cover_audio_tokens, artist_audio=artist_audio_tokens, underpainting_audio=underpainting_audio_tokens, overpainting_audio=overpainting_audio_tokens, playlist_audio=playlist_audio_tokens, multi_artist_audio=multi_artist_audio_tokens, history_latents=history_latents, future_latents=future_latents if future_latents is not None else None, ) elif item.is_diff_infill: assert (infill_start_s is not None) and (infill_end_s is not None) # note that history_audio at this stage is the full tokens assert (infill_context_start_s is not None) and (infill_context_end_s is not None) assert history_latents is not None assert future_latents is not None if round(infill_context_diff := infill_context_end_s - infill_context_start_s, 3) != 30: self.print_info( item.id, f"""Failed to trim infill context to 30s: infill_context_start_s: {infill_context_start_s} infill_context_end_s: {infill_context_end_s} infill_start_s: {infill_start_s} infill_end_s: {infill_end_s}""", ) raise ValueError( f"Diffusion context needs to be exactly 30s: {infill_context_start_s}, {infill_context_end_s}, diff: {infill_context_diff}" ) infill_context_start_index = math.ceil(infill_context_start_s * self._vae_freq()) partial_semantic_tokens = parent_audio_tokens[ infill_context_start_index : infill_context_start_index + 30 * 25 ] if partial_semantic_tokens.shape[0] != 750: self.print_info( item.id, f"Partial semantic tokens are not 750: {partial_semantic_tokens.shape}" ) parent_audio_tokens = parent_audio_tokens[: 30 * 25] self.print_info( item.id, f"""diffusion infill, partial_semantic_tokens: {partial_semantic_tokens.shape}, history_latents: {None if history_latents is None else history_latents.shape}, future_latents: {None if future_latents is None else future_latents.shape}""", ) return HistoryPrompt( prompt_audio=partial_semantic_tokens, prompt_lyrics="", # for infilling history text is empty history_latents=history_latents, future_latents=future_latents, ) # THIS IS REALLY BAD but history_audio isn't trimmed # So if it is longer than 60 seconds, it just doesn't work... # we keep the last bit if not (item.is_upsample or item.is_stem): parent_audio_tokens = parent_audio_tokens[-int(history_max_duration * 25) :] self.print_info( item.id, f"trimmed audio npz prmpt {parent_audio_tokens.shape}", ) hoot_start_time = time.time() with tracer.trace("load_history_audio_as_audio"): history_audio_as_audio = self.worker.load_history_audio_as_audio(item) assert history_audio_as_audio is not None continue_at = item.metadata.get("continue_at", None) self.print_info( item.id, f"loaded audio mp3, with duration: {history_audio_as_audio.duration_s}" ) if continue_at is None: self.print_info(item.id, "continue_at is unset, set to audio duration.") continue_at = history_audio_as_audio.duration_s continue_at = max(1, continue_at) # bound continue_at to be greater than 1 sec # hoot input don't need to be longer than the clip # since user generate clips that are longer and concat, we need to trim the input history_audio_segment = history_audio_as_audio.get_segment( from_s=max( 0, continue_at - history_max_duration, ), to_s=continue_at, ) aligned_history_text = item.metadata.get("continued_aligned_prompt", None) if aligned_history_text is None: self.print_info(item.id, "aligned text not provided, running alignment") try: with tracer.trace("hoot_encode_and_align"): hoot_output = chirp_v3.hoot_encode_and_align( history_audio_segment, history_text ).strip() self.print_info( item.id, f""" History_text:\n {history_text}; \n Transcript:\n {hoot_output} \n Runtime: {time.time() - hoot_start_time:.1f}s; \n Audio duration: {history_audio_segment.duration_s:.1f}s""", ) history_text = hoot_output except Exception as e: self.print_info(item.id, f"transcription failed: {e}") # pass in empty history text for now history_text = "" else: history_text = aligned_history_text # make sure we don't use the audio vae for artist or cover conditions # but for extend, we want to use the audio vae still if ( (item.is_artist_condition or item.is_cover_condition) and history_latents is not None and ("extend" not in item.metadata.get("task", "")) ): self.print_info( item.id, f"{item.metadata.get('task')}, setting history_latents to None", ) history_latents = None # should always return sth n_history_tokens = parent_audio_tokens.shape[0] if parent_audio_tokens is not None else 0 n_cover_tokens = cover_audio_tokens.shape[0] if cover_audio_tokens is not None else 0 n_artist_tokens = artist_audio_tokens.shape[0] if artist_audio_tokens is not None else 0 n_total_tokens = n_history_tokens + n_cover_tokens + n_artist_tokens if n_total_tokens > 6 * 60 * 25 and not (item.is_upsample or item.is_stem): # sanity check self.print_info( item.id, f"""HistoryPrompt superlong tokens found: n_history_tokens: {n_history_tokens}, n_cover_tokens: {n_cover_tokens}, n_artist_tokens: {n_artist_tokens}, total: {n_total_tokens}""", ) return HistoryPrompt( prompt_audio=parent_audio_tokens, future_audio=None, prompt_lyrics=history_text, cover_audio=cover_audio_tokens, artist_audio=artist_audio_tokens, underpainting_audio=underpainting_audio_tokens, overpainting_audio=overpainting_audio_tokens, playlist_audio=playlist_audio_tokens, multi_artist_audio=multi_artist_audio_tokens, history_latents=history_latents, ) @modal.method() @distributed_trace("encode_history", HISTORY_DDOG_SERVICE, env_name=DEPLOYMENT_TYPE) def encode_history(self, queue_item: str) -> None: if VERBOSE_MESSAGE: print_gpu_memory_usage(self.__class__.__name__) torch.cuda.reset_max_memory_allocated() print(queue_item) item = QueueItem(**json.loads(queue_item)) # enter some forced retry logic -- after that we notify failure max_retries = 3 for attempt in range(1, max_retries + 1): try: with tracer.trace("encode_history_attempt"): history_prompt = self._encode_history(item) if VERBOSE_MESSAGE: max_mem = torch.cuda.max_memory_allocated() / 1e9 print(f"ChirpV2Stub(EncoderStub): Max memory allocated: {max_mem:.2f} GB") # return the call to the orchetsator with tracer.trace("spawn_orchestrator"): self.modal_f_orchestrator.spawn( item.model_dump_json(), history_prompt, parent_context=serialize_context() ) print(f"Spawned modal job {item.id}.") return # hate to be broad here but both ValueError and TypeErrors are possible # also not obvious why this would happen # but when it happens the worker is in a bad state # we need to time it out to kill it for now except Exception as e: if item.prompt_audio: # check if the s3 npz exists # in rare cases it may not -- we will have to cycle try: # try to find the npz file exists or not s3_client.head_object( Bucket="suno-data-uploads", Key=f"studio/uploads/{item.prompt_audio.replace('.mp3', '.npz')}", ) except: print(f"HistoryLoader {item.id}: encoding prompt audio to codes.") encode_audio_f_call = self.modal_f_cycle.spawn( f"s3://suno-data-uploads/studio/uploads/{item.prompt_audio.replace('.mp3', '')}.mp3", s3_npz_id=item.prompt_audio.replace(".mp3", ""), encode_vae_version=None, ) _ = encode_audio_f_call.get(timeout=30) if attempt != max_retries: print(f"Error in encode history: {e}") time.sleep(1) else: # TODO: Do we need both? what's the difference? print(e) traceback.print_exc() try: item.notify_progress( { "id": item.id, "model": item.model_name, "n_audios": len(item.ids) if item.ids else 1, "ok": 0, "gen_duration": 0, "error_type": "generation_failure", "error_message": "Can't start the continue generation.", } ) except Exception as e2: print(f"Error in notify progress: {e2}") raise e2 @app.cls( gpu="A10G", secrets=SECRETS, timeout=600, scaledown_window=400, memory=15000, retries=modal.Retries( max_retries=2, backoff_coefficient=2.0, initial_delay=5.0, ), max_containers=ENCODER_CONCURRENCY_LIMITS[DEPLOYMENT_TYPE], min_containers=0, region="us-east", volumes={MODEL_STORE_VOLUME_DIR: model_store_volume}, ) @modal.concurrent(max_inputs=ENCODER_MAX_INPUT) class HootEncoderStub: def __init__(self): import torch num_gpus = torch.cuda.device_count() print(f"Found {num_gpus} GPUs.") self.worker = CodecWorker() @modal.method() def hoot_encode(self, s3_url: str, prior_texts: Optional[str] = None): audio = Audio.from_s3(s3_url, n_channels=2) return hoot_encode(audio, prior_texts=prior_texts) @app.local_entrypoint() def main(): model = HistoryEncoderStub() input = json.dumps( { "id": "a276071f-cef6-4bc8-b69b-58a4ccb69e0c", "prompt_audio": "8fdcf425-19c5-4ddb-a1d0-71e9d3dc01eb.mp3", "prompt_npz": None, "prompt_text": "", "metadata": { "type": "gen", "source": "web", "prompt": "", "tags": "experimental hyperstep", "gpt_prompt": None, "gpt_description_prompt": None, "stream": True, "make_instrumental": False, "priority": 0, "continued_from_prompt": "", "history": [ { "id": "8fdcf425-19c5-4ddb-a1d0-71e9d3dc01eb", "continue_at": 33.0, "type": "gen", "source": "web", } ], "audio_prompt_id": "8fdcf425-19c5-4ddb-a1d0-71e9d3dc01eb", "continue_at": 33.0, "options": None, "task": "extend", }, "gen_duration": 12, "callback_url": "https://studio-api.suno.ai/api/generate/finish-clip/", "model_name": "chirp-v3-5", "title": "", "ids": ["bcc8bf5c-ede9-4f8d-968a-0225ac0e25f3", "8addbd4e-5d28-42de-a836-32b908f4f581"], } ) model.encode_history.remote(input) print("finish test encode_history") input = json.dumps( { "id": "f3fb4b05-49ab-45f2-9591-1bcdcef6ece6", "prompt_audio": "99c4280e-c26b-4c8b-87a8-494512a24484.mp3", "prompt_npz": None, "prompt_text": "[verse]\n\n", "metadata": { "tags": "shoegaze, ethereal, art, ego", "gpt_prompt": None, "gpt_description_prompt": None, "stream": True, "make_instrumental": False, "priority": 10, "user_id": 172918, "is_bot": False, "task": "extend", "continued_from_prompt": "[melodic hook]\n\n", "continued_aligned_prompt": None, "history": [ { "id": "99c4280e-c26b-4c8b-87a8-494512a24484", "continue_at": 160.0, "type": "gen", "source": "web", "infill": False, } ], "continue_at": 160.0, "options": None, "edit_session_id": None, "infill": False, "type": "gen", "lang": "English", }, "gen_duration": 12, "callback_url": "https://studio-api.prod.suno.com/api/generate/finish-clip/", "model_name": "chirp-v4", "title": "", "ids": ["e72a78db-a9a2-43a3-8b97-597db7f76c12", "3c0160e0-0924-45a3-9c95-fb2e8ffaaddc"], } ) model.encode_history.remote(input) print("finish test encode_history 2") input = json.dumps( { "id": "790650e3-ab88-4bec-995d-64935371772f", "prompt_audio": "8a200dbe-4f12-4050-b88d-77e16efd760d.mp3", "prompt_npz": None, "prompt_text": "[Outro]\nTest Test Test Test\nEnd end end end\nend\n\n[end]", "metadata": { "credit_cost": 5, "feature_flags": 3, "user_handle": "mikey", "source": "web", "prompt": "[Outro]\nTest Test Test Test\nEnd end end end\nend\n\n[end]", "tags": "dance electronic", "gpt_prompt": None, "gpt_description_prompt": None, "stream": True, "make_instrumental": False, "priority": 0, "user_id": 132, "is_bot": False, "task": "extend", "configurations": {}, "continued_from_prompt": "[Verse]\nShabi\nShabi all night long (ooh-yeah!)\nFeel the rhythm\nFeel so strong\nLights are flashing\nHearts collide\nShabi\nShabi\nLet's take a ride\n\n[Verse 2]\nStep by step\nWe move as one\nDancing till we see the sun\nHands in the air\nTouch the sky\nShabi\nShabi\nWe can fly\n\n[Chorus]\nShabi shabi\nFeel the groove\nShabi shabi\nLet's all move\nShabi shabi\nDon't say no\nShabi shabi\nLet it flow\n\n[Verse 3]\nMusic pumping\nBass so loud\nLost in the crowd\nFeeling proud\nShabi\nShabi\nCan't stand still\nFeel the beat\nIt's such a thrill\n\n[Verse 4]\nTwirl around\nSpin me fast\nMake this moment always last\nShabi\nShabi\nHearts on fire\nTake me higher\nTake me higher\n\n[Chorus]\nShabi shabi\nFeel the groove\nShabi shabi\nLet's all move\nShabi shabi\nDon't say no\nShabi shabi\nLet it flow", "continued_aligned_prompt": None, "history": [ { "id": "8a200dbe-4f12-4050-b88d-77e16efd760d", "continue_at": 180.0, "type": "gen", "source": "web", "infill": False, } ], "continue_at": 180.0, "options": None, "edit_session_id": None, "edited_clip_id": "8a200dbe-4f12-4050-b88d-77e16efd760d", "infill": False, "type": "gen", "lang": "", }, "gen_duration": 12, "callback_url": "https://studio-api.staging.suno.com/api/generate/finish-clip/", "model_name": "chirp-v3p5-h-s-31", "title": "Shabi Shabi", "ids": ["d67755c2-c4a8-4a0a-b0e2-c8cbb8878747", "50c06369-2ae9-4bb0-96ff-ca1b274a3010"], } ) model.encode_history.remote(input) print("finish test encode_history 3") # test infill infill_input = json.loads(input) infill_input["metadata"]["task"] = "infill" infill_input["metadata"]["infill_start_s"] = 5.0 infill_input["metadata"]["infill_end_s"] = 10.0 infill_input["metadata"]["infill_context_start_s"] = 0.0 infill_input["metadata"]["infill_context_end_s"] = 20.0 infill_input["metadata"]["include_future_s"] = 5.0 infill_input["metadata"]["include_history_s"] = 5.0 del infill_input["metadata"]["continue_at"] model.encode_history.remote(json.dumps(infill_input)) print("finish test encode_history infill") # test infill with minimal params infill_input = dict( model_name="chirp-v3-5-tau", id="infilling_01234_modal", prompt_text="", prompt_audio="af3ce26e-5571-4c92-a9d7-e981ccd6827b.mp3", metadata={ "task": "infill", "infill_start_s": 5.0, "infill_end_s": 15.0, "infill_context_start_s": 0.0, "infill_context_end_s": 20.0, "include_future_s": 5.0, "include_history_s": 5.0, }, ) model.encode_history.remote(json.dumps(infill_input)) print("finish test encode_history infill minimal") hoot_model = HootEncoderStub() print( hoot_model.hoot_encode.remote( s3_url="s3://suno-data-uploads/studio/uploads/4a77dea7-19f3-46d2-8b0a-b2b7e9ea9a05.mp3" ) ) time.sleep(10) print("Done")