"""Hoot alignment and video generation application on modal.""" import json import re import tempfile import traceback import numpy as np import time import contextlib import os import modal from suno_utils.worker.hoot_stub_mp import HootBackgroundWorker, LoadAudioResult import torch from suno_utils.audio import Audio from suno_utils.gpt import chirp_v2_5 as chirp_v3 from suno_utils.tasks.hoot import ( encode, LOGIT_DURATION_S, ) from suno_utils.tasks.lyrics_alignment.shortest_path_aligner import ( ShortestPathHootV6Config, ) from suno_utils.worker.loader import S3Loader, retry_s3_download from suno_utils.worker.schema import QueueItem from suno_utils.worker.modal_runner_sdxl_worker import IMAGE_PREFIX, LARGE_IMAGE_PREFIX from suno_utils.worker.settings import s3_client from suno_utils.worker.utils import retry_decorator from suno_utils.worker.modal_base import get_modal_base_image from suno_utils.worker.modal_model_volume import ( MODEL_STORE_VOLUME_DIR, model_store_volume, MODEL_STORE_VOLUME_PREFIX, ) ############## CHANGE THESE ############## DEPLOYMENT_TYPE = "dev" # priority, staging, dev, prod, msft #### APP_NAME = f"videos-v2-{DEPLOYMENT_TYPE}" assert APP_NAME.endswith(DEPLOYMENT_TYPE) SECRETS = [ modal.Secret.from_name("studio-aws"), modal.Secret.from_name("redis-test"), modal.Secret.from_name("api-callback-token"), ] base_image = ( get_modal_base_image() .pip_install_private_repos( "github.com/suno-ai/neon.git@0ffdb64#subdirectory=hoot", git_user="mcamac", secrets=[modal.Secret.from_name("victor-modal-github-token")], ) .pip_install("biopython>=1.81") .add_local_python_source("suno_utils", copy=False) ) HOOT_CKPT_PATH = "/checkpoints/hoot_v6/2025-06-23_16-49-02.pt" HOOT_TOKENIZER_PATH = "/checkpoints/hoot_v6/tokenizer.model" S3_BUCKET_NAME = "suno-data-uploads" S3_FOLDER_NAME = "studio/uploads" class HootWorker: def __init__(self): print("Start loading models") chirp_v3.preload_hoot_models( checkpoint_filepath=f"{MODEL_STORE_VOLUME_PREFIX}{HOOT_CKPT_PATH}", tokenizer_filepath=f"{MODEL_STORE_VOLUME_PREFIX}{HOOT_TOKENIZER_PATH}", ) print("Finish loading models") app = modal.App(APP_NAME, image=base_image) N_CPU_HOOT = 32 @app.cls( cpu=N_CPU_HOOT, gpu="A10G", secrets=SECRETS, timeout=200, scaledown_window=300, memory=128 * 1024, min_containers=1 if DEPLOYMENT_TYPE == "dev" else 20, max_containers=400, # cloud="aws", region="us-east", buffer_containers=1 if DEPLOYMENT_TYPE == "dev" else 5, volumes={MODEL_STORE_VOLUME_DIR: model_store_volume}, ) @modal.concurrent(max_inputs=32) class HootAlignmentStub: def __init__(self): from threading import Semaphore # other cores will be used by the background worker torch.set_num_threads(1) os.environ["OMP_NUM_THREADS"] = "1" num_gpus = torch.cuda.device_count() print(f"Found {num_gpus} GPUs.") self.worker = HootWorker() # only allow 2 concurrent encode calls, more risks CUDA OOM self.encode_semaphore = Semaphore(2) t0 = time.time() self.background_worker = HootBackgroundWorker( tokenizer_filepath=f"{MODEL_STORE_VOLUME_PREFIX}{HOOT_TOKENIZER_PATH}", aligner_config=ShortestPathHootV6Config, num_workers=2 * N_CPU_HOOT - 1, ) print(f"HootBackgroundWorker initialized in {time.time() - t0} seconds") @modal.exit() def cleanup_pool(self): self.background_worker.cleanup() def _get_audio_info_and_upload( self, queue_item: str, audio_result: LoadAudioResult, words_timing: list[dict] = [] ) -> None: """Add waveform to the json and upload to s3.""" item = QueueItem(**json.loads(queue_item)) fname = f"{item.id}_hoot.json" s3_fname = f"{S3_FOLDER_NAME}/{fname}" metadata = { "metadata": { "audio_rms": audio_result.waveform, } } words_timing.append(metadata) # Convert words_timing dictionary to a JSON string and encode it to bytes words_timing_str = json.dumps(words_timing, ensure_ascii=False, indent=4) print(f"Hoot {item.id}: before encoding {fname}.") words_timing_bytes = words_timing_str.encode("utf-8") print(f"Hoot {item.id}: before uploading {fname} to s3.") # Upload the bytes directly to S3 s3_client.put_object(Bucket=S3_BUCKET_NAME, Key=s3_fname, Body=words_timing_bytes) print(f"Hoot {item.id}: Uploaded {fname}.") @contextlib.contextmanager def _audio_temp_file(self, s3_id: str): # download to file so we can load in background worker with tempfile.NamedTemporaryFile(suffix=".mp3") as tmp_file: retry_s3_download(S3_BUCKET_NAME, f"{S3_FOLDER_NAME}/{s3_id}.mp3", tmp_file) tmp_file.flush() yield tmp_file.name @modal.method() def get_aligned_lyrics(self, queue_item: str) -> list[dict]: """Basically hit the hoot API, very efficiently.""" start_time = time.time() item = QueueItem(**json.loads(queue_item)) do_callback_aligned_lyrics = item.metadata.get("callback_aligned_lyrics", False) enable_augmentation = item.metadata.get("enable_augmentation", True) words_timing = [] response_json = {"type": "hoot_alignment_finished", "id": item.id} if do_callback_aligned_lyrics: response_json["alignment"] = {} try: print( f"Hoot {item.id}: Starting. full info {queue_item}. took {round(time.time() - start_time, 2)} seconds" ) with self._audio_temp_file(item.id) as audio_path: audio_result = self.background_worker.load_audio(audio_path) lyrics = item.prompt_text if item.prompt_text else "" if not lyrics.strip(): self._get_audio_info_and_upload(queue_item, audio_result, []) item.notify_progress( response_json, ) return [] print( f"Hoot {item.id}: getting alignment. took {round(time.time() - start_time, 2)} seconds" ) with self.encode_semaphore: with torch.cuda.amp.autocast(): logits = encode( audio_result.audio_tensor, return_logits=True, batch_size=4, return_torch=True ) assert isinstance(logits, torch.Tensor) align_interval_start_s = item.metadata.get("align_interval_start_s") align_interval_end_s = item.metadata.get("align_interval_end_s") align_offset = 0.0 if align_interval_start_s is not None and align_interval_end_s is not None: align_first_logit = int(np.floor(align_interval_start_s / LOGIT_DURATION_S)) align_last_logit = int(np.ceil(align_interval_end_s / LOGIT_DURATION_S)) assert ( align_first_logit >= 0 and align_last_logit <= logits.shape[0] and align_first_logit < align_last_logit ) logits = logits[align_first_logit:align_last_logit] print( f"Hoot {item.id}: aligned logits shape: {logits.shape} between {align_first_logit} and {align_last_logit}" ) align_offset = align_first_logit * LOGIT_DURATION_S aligner_start_time = time.time() words_timing, hoot_lyrics, hoot_cer = self.background_worker.align( lyrics, logits, override_enable_jumps=enable_augmentation ) print( f"Hoot {item.id}: shortest path alignment took {round(time.time() - aligner_start_time, 2)} seconds (not cumulative)" ) # these values don't make sense for shortest path alignment # but downstream consumers expect them for w in words_timing: w["success"] = True w["p_align"] = 1.0 w["start_s"] += align_offset w["end_s"] += align_offset words_timing.extend( [ {"hoot_lyrics": hoot_lyrics}, {"hoot_cer": hoot_cer}, {"decoder": f"ShortestPathAligner{'-no-jumps' if not enable_augmentation else ''}"}, ] ) if DEPLOYMENT_TYPE == "dev": print(f"Hoot {item.id}: words_timing: {words_timing}") print( f"Hoot {item.id}: finished alignment. took {round(time.time() - start_time, 2)} seconds" ) if not do_callback_aligned_lyrics: # do not update the json if the call is to just get aligned lyrics # spin this up but not blocking the gpu thread self._get_audio_info_and_upload(queue_item, audio_result, words_timing) except Exception as e: print(f"Hoot {item.id}: failed with: {str(e)}") traceback.print_exc() print(f"Hoot {item.id}: finished. took {round(time.time() - start_time, 2)} seconds") if do_callback_aligned_lyrics: response_json["alignment"] = words_timing # send notification to the client item.notify_progress( response_json, ) return words_timing @app.cls( cpu=4, secrets=SECRETS, timeout=200, scaledown_window=360, retries=modal.Retries( max_retries=2, backoff_coefficient=2.0, initial_delay=5.0, ), # cloud="aws", region="us-east", buffer_containers=1 if DEPLOYMENT_TYPE == "dev" else 20, ) @modal.concurrent(max_inputs=4) class DummyV0Stub: def __init__(self): self.worker = S3Loader() self.modal_f_hoot_alignment = modal.Cls.lookup( f"videos-v2-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}", "HootAlignmentStub", )().get_aligned_lyrics @modal.method() def write_video( self, queue_item: str, image_id: str | None = None, is_square: bool | None = False, # TODO: remove this after redeploy all engines... audio: Audio | None = None, ): print(f"Get video gen request {queue_item}, image id {image_id}") # use wait_seconds 10, since the images / audios should have been generated by now long_retry_s3_download = retry_decorator(3, wait_seconds=10)(s3_client.download_fileobj) item = QueueItem(**json.loads(queue_item)) words_timing = [] try: backlog = self.modal_f_hoot_alignment.get_current_stats().backlog if backlog <= 1000: fn_calls = self.modal_f_hoot_alignment.spawn(item.json()) words_timing = fn_calls.get(timeout=60) words_timing = [word for word in words_timing if "metadata" not in word] except Exception as e: print(f"Failed to get hoot alignment: {str(e)}") if not words_timing: words_timing = [ {"word": s} for s in re.split(r"([\w'\-]+)", item.prompt_text if item.prompt_text else "") if len(s) > 0 ] should_generate_video = item.metadata.get("generate_video", False) if not should_generate_video: # technically might not be required, but just do this to be consistent item.notify_progress( {"type": "video_finished", "video_generated": False, "id": item.id}, ) return elif should_generate_video and image_id: with tempfile.NamedTemporaryFile() as tmp_file: # check if large image exists prefered_image_path = ( f"{S3_FOLDER_NAME}/{image_id.replace(IMAGE_PREFIX, LARGE_IMAGE_PREFIX)}" ) if prefered_image_path.endswith(".png"): prefered_image_jepg_path = prefered_image_path.replace(".png", ".jpeg") try: s3_client.get_object(Bucket=S3_BUCKET_NAME, Key=prefered_image_jepg_path) print(f"Successfully got jpeg image for {item.id}") # replace the path and then download... prefered_image_path = prefered_image_jepg_path except Exception as e: try: print(f"Failed to get jpeg image for {item.id} -- try png: {str(e)}") s3_client.get_object(Bucket=S3_BUCKET_NAME, Key=prefered_image_path) except Exception as e: print(f"Failed to get png image for {item.id} -- try small png: {str(e)}") # large image doesn't exist -- could due to the fact that it is uploaded, not generated prefered_image_path = f"studio/uploads/{image_id}" # check if image exists try: s3_client.get_object(Bucket=S3_BUCKET_NAME, Key=prefered_image_path) except Exception as e: print( f"Failed to get small png image for {item.id} -- use default: {str(e)}" ) # the image doesn't exists -- go to fallback default image FIXED_IMAGE_FILE = "image_large_default_bird.jpeg" prefered_image_path = f"studio/uploads/{FIXED_IMAGE_FILE}" elif prefered_image_path.endswith(".jpeg"): try: s3_client.get_object(Bucket=S3_BUCKET_NAME, Key=prefered_image_path) except Exception as e: print(f"Failed to get jpeg image for {item.id} -- try small jpeg: {str(e)}") # sometimes the large image doesn't exist... prefered_image_path = prefered_image_path.replace( LARGE_IMAGE_PREFIX, IMAGE_PREFIX ) try: s3_client.get_object(Bucket=S3_BUCKET_NAME, Key=prefered_image_path) except Exception as e: print( f"Failed to get small jpeg image for {item.id} -- use default: {str(e)}" ) FIXED_IMAGE_FILE = "image_large_default_bird.jpeg" prefered_image_path = f"{S3_FOLDER_NAME}/{FIXED_IMAGE_FILE}" # this should always be able to download long_retry_s3_download(S3_BUCKET_NAME, prefered_image_path, tmp_file) try: self.worker._write_video( item, show_text=True, visualizer=True, image_path=tmp_file, aligned_text=words_timing, audio=audio, # don't need to go through s3 if audio is around ) item.notify_progress( {"type": "video_finished", "video_generated": True, "id": item.id}, ) except FileNotFoundError: # this means the audio failed to download # could happen if the audio isn't uploaded successfully...or deleted print(f"Failed to download audio for {item.id}") item.notify_progress( { "id": item.id, "type": "error", "error_type": "video_generation_failure", "error_message": "Video generation failed due to missing audio.", } ) @app.local_entrypoint() def main(): import time # Configuration for stress test NUM_STRESS_TASKS = 50000 BATCH_SIZE = 1000 # Process in batches to avoid overwhelming the system print(f"Starting stress test with {NUM_STRESS_TASKS} tasks") # Initialize models hoot_model = HootAlignmentStub() video_model = DummyV0Stub() # Use the same valid test inputs as before - these actually exist in S3 hoot_input_1 = json.dumps( { "id": "6c1dfbff-129b-46f9-b3b4-31b1201cd726", "prompt_audio": None, "prompt_npz": None, "prompt_text": "[Instrumental]", "metadata": { "credit_cost": 5, "feature_flags": 3, "type": "gen", "source": "web", "prompt": "", "tags": "death metal aggressive thunderous", "gpt_prompt": None, "gpt_description_prompt": "metal song", "stream": True, "make_instrumental": True, "priority": 0, "user_id": 224, "is_bot": False, "task": None, "experiment": ["chirp-v3p5-engine-s-8", "chirp-v3p5-engine-s-8"], "experiment_version": "v_156", }, "gen_duration": 12, "callback_url": "https://studio-api.suno.ai/api/generate/finish-clip/", "model_name": "chirp-v3p5-engine-s-8", "title": "Eternal Carnage", "ids": None, } ) hoot_input_2_aug_true = json.dumps( { "id": "1cee79db-a0ee-44e4-9646-d312cd620997", "prompt_text": """ [verse 1] 滚滚长江东逝水， 浪花淘尽英雄。 [verse 2] 是非成败转头空 青山依旧在 几度夕阳红。 [bridge] 白发渔樵江渚上 惯看秋月春风。 [chorus] 一壶浊酒喜相逢。 古今多少事，都付笑谈中。 [Refrain] 一壶浊酒喜相逢。 古今多少事，都付笑谈中。 [end] """, "metadata": { "tags": "中文流行 chinese pop", "enable_augmentation": True, "generate_video": False, "align_interval_start_s": 1.0, "align_interval_end_s": 220.0, }, "title": "临江仙", "callback_url": "https://studio-api.suno.ai/api/generate/finish-clip/", } ) hoot_input_2_aug_false = json.dumps( { "id": "1cee79db-a0ee-44e4-9646-d312cd620997", "prompt_text": """ [verse 1] 滚滚长江东逝水， 浪花淘尽英雄。 [verse 2] 是非成败转头空 青山依旧在 几度夕阳红。 [bridge] 白发渔樵江渚上 惯看秋月春风。 [chorus] 一壶浊酒喜相逢。 古今多少事，都付笑谈中。 [Refrain] 一壶浊酒喜相逢。 古今多少事，都付笑谈中。 [end] """, "metadata": { "tags": "中文流行 chinese pop", "enable_augmentation": False, "generate_video": False, "align_interval_start_s": 1.0, "align_interval_end_s": 220.0, }, "title": "临江仙", "callback_url": "https://studio-api.suno.ai/api/generate/finish-clip/", } ) # Video generation input (same as hoot_input_2_aug_false but with video enabled) video_input = json.dumps( { "id": "1cee79db-a0ee-44e4-9646-d312cd620997", "prompt_text": """ [verse 1] 滚滚长江东逝水， 浪花淘尽英雄。 [verse 2] 是非成败转头空 青山依旧在 几度夕阳红。 [bridge] 白发渔樵江渚上 惯看秋月春风。 [chorus] 一壶浊酒喜相逢。 古今多少事，都付笑谈中。 [Refrain] 一壶浊酒喜相逢。 古今多少事，都付笑谈中。 [end] """, "metadata": { "tags": "中文流行 chinese pop", "enable_augmentation": False, "generate_video": True, "align_interval_start_s": 1.0, "align_interval_end_s": 220.0, }, "title": "临江仙", "callback_url": "https://studio-api.suno.ai/api/generate/finish-clip/", } ) # Test inputs to cycle through test_inputs = [hoot_input_1, hoot_input_2_aug_true, hoot_input_2_aug_false] all_tasks = [] start_time = time.time() print("Spawning tasks in batches...") for batch_num in range(0, NUM_STRESS_TASKS, BATCH_SIZE): batch_tasks = [] batch_end = min(batch_num + BATCH_SIZE, NUM_STRESS_TASKS) print(f"Spawning batch {batch_num//BATCH_SIZE + 1}: tasks {batch_num} to {batch_end-1}") for i in range(batch_num, batch_end): # Cycle through the valid test inputs input_to_use = test_inputs[i % len(test_inputs)] # Spawn hoot alignment task hoot_task = hoot_model.get_aligned_lyrics.spawn(input_to_use) batch_tasks.append(("hoot", hoot_task, f"task_{i}")) # Every 10th task also spawns a video generation task if i % 10 == 0: video_task = video_model.write_video.spawn( video_input, "image_large_default_bird.jpeg", ) batch_tasks.append(("video", video_task, f"video_task_{i}")) all_tasks.extend(batch_tasks) # Small delay between batches to avoid overwhelming the spawn mechanism time.sleep(0.1) spawn_time = time.time() - start_time print(f"All {len(all_tasks)} tasks spawned in {spawn_time:.2f} seconds") print(f"Average spawn rate: {len(all_tasks)/spawn_time:.2f} tasks/second") # Wait for all tasks to complete print("Waiting for all tasks to complete...") completed_tasks = 0 failed_tasks = 0 completion_start = time.time() for task_type, task, task_id in all_tasks: try: result = task.get(timeout=600) # 10 minute timeout per task completed_tasks += 1 if completed_tasks % 1000 == 0: elapsed = time.time() - completion_start print( f"Completed {completed_tasks}/{len(all_tasks)} tasks in {elapsed:.2f}s (rate: {completed_tasks/elapsed:.2f} tasks/s)" ) except Exception as e: failed_tasks += 1 if failed_tasks <= 10: # Only print first 10 failures to avoid spam print(f"Task {task_id} ({task_type}) failed: {str(e)}") elif failed_tasks == 11: print("... (suppressing further failure messages)") total_time = time.time() - start_time completion_time = time.time() - completion_start print("\n" + "=" * 60) print("STRESS TEST RESULTS") print("=" * 60) print(f"Total tasks spawned: {len(all_tasks)}") print(f"Tasks completed successfully: {completed_tasks}") print(f"Tasks failed: {failed_tasks}") print(f"Success rate: {(completed_tasks/len(all_tasks)*100):.2f}%") print(f"Total time (spawn + execution): {total_time:.2f} seconds") print(f"Spawn time: {spawn_time:.2f} seconds") print(f"Execution time: {completion_time:.2f} seconds") print(f"Average task completion rate: {completed_tasks/completion_time:.2f} tasks/second") print("=" * 60) print("Stress test completed!")