"""Audio upload application on modal. Audio similarity detection through 3rd party. """ from pydantic import Field import tempfile import time import logging import os import requests import tarfile import tempfile import subprocess import uuid from datetime import datetime import traceback import json import time from typing import Tuple from pydantic import BaseModel import boto3 import modal from suno_utils.audio import Audio from suno_utils.utils.s3 import _download_s3_file from suno_utils.worker.loader import S3Loader from suno_utils.tasks.match_audio import find_audio_matches from suno_utils.worker.settings import s3_client from suno_utils.worker.utils import retry_decorator from suno_utils.worker.schema import QueueItem from suno_utils.worker.modal_base import get_modal_base_image from suno_utils.worker.modal_model_configs import VAEVersion from datadog import statsd logger = logging.getLogger(__name__) logging.basicConfig() logger.setLevel(logging.INFO) ############## CHANGE THESE ############## DEPLOYMENT_TYPE = "dev" ########################################## DEPLOYMENT_TYPES = {"dev", "prod"} assert DEPLOYMENT_TYPE in DEPLOYMENT_TYPES # orchestrator talks to studio api APP_NAME = f"upload-audio-{DEPLOYMENT_TYPE}" assert APP_NAME.endswith(DEPLOYMENT_TYPE) MOUNT_PATH = "/suno/models" UPLOADS_S3_BUCKET = "suno-data-uploads" retry_s3_download = retry_decorator(3, wait_seconds=20)(s3_client.download_fileobj) USE_GEMINI_BY_DEFAULT = DEPLOYMENT_TYPE == "dev" AUDIBLE_MAGIC_TOOLKIT = "AudibleMagicToolkit_46.9a_Ubuntu" AUDIBLE_MAGIC_CONFIG = "SunoMU_v46.config" # for maximum audio duration -- tbd MAX_VIP_UPLOAD_DURATION = 480 MAX_UPLOAD_DURATION = 120 if DEPLOYMENT_TYPE == "prod" else 480 MAX_FREE_USER_UPLOAD_DURATION = 60 if DEPLOYMENT_TYPE == "prod" else 480 # for audible magic to work MIN_UPLOAD_DURATION = 6 MIN_VIP_UPLOAD_DURATION = 3 MAX_GEMINI_OUTPUT_TOKENS = 1000 class AudibleMagicWorker(S3Loader): """Provide nofication functionality for ChatGptStub.""" def __init__(self): self.audio_magic_bin_path = os.path.join(MOUNT_PATH, AUDIBLE_MAGIC_TOOLKIT, "bin") self.audio_magic_config_path = os.path.join(MOUNT_PATH, AUDIBLE_MAGIC_CONFIG) def is_audio_match_existing_music(self, file_path: str, claimed_artist_names: list[str]) -> bool: """If the audio matches existing music, return true. Status Code Description 2005 No matches were found. 2006 At least one match was found. For reference: https://support.audiblemagic.com/hubfs/Toolkit/v46/docs/AMToolkit%20v46%20Users%20Guide_20231107.pdf """ if not os.path.exists(file_path): raise ValueError(f"File doesn't exists at {file_path}") try: match_output = find_audio_matches( file_path, bin_dir=self.audio_magic_bin_path, config_path=self.audio_magic_config_path, ) if match_output is not None: if match_output.get("statusCode") == 2005: return False elif match_output.get("statusCode") == 2006: print(f"Match found for {file_path}: {match_output}") matched_artist_names = [ match.get("metadata", {}).get("Artist", "").lower() for match in match_output.get("matches", []) ] if any( [ matched_artist_name in set([n.lower() for n in claimed_artist_names]) for matched_artist_name in matched_artist_names ] ): print( f"[{', '.join(matched_artist_names)}] matched a claimed artist name. Overriding artwork match." ) return False return True return False except Exception as e: tags = [ f"env:{DEPLOYMENT_TYPE}", f"exception_type:{type(e).__name__}", ] statsd.increment("audible_magic.errors", tags=tags) print(f"Exception raised for {file_path}, {str(e)}. Will take it as artwork.") return True @staticmethod def download_models() -> None: """Use AWS CLI to download models if they don't exist.""" os.makedirs(MOUNT_PATH, exist_ok=True) target_tgz_path = os.path.join(MOUNT_PATH, f"{AUDIBLE_MAGIC_TOOLKIT}.tgz") print(f"download the tgz file to {target_tgz_path}") _download_s3_file( f"s3://suno-data/tony/AudibleMagic/{AUDIBLE_MAGIC_TOOLKIT}.tgz", target_tgz_path, ) print("download the config file.") _download_s3_file( f"s3://suno-data/tony/AudibleMagic/{AUDIBLE_MAGIC_CONFIG}", os.path.join(MOUNT_PATH, AUDIBLE_MAGIC_CONFIG), ) _download_s3_file( "s3://suno-data/tony/prompt_audio/test_jingle.mp3", os.path.join(MOUNT_PATH, "test_jingle.mp3"), ) with tarfile.open(target_tgz_path) as tar: tar.extractall(MOUNT_PATH) print("Downloaded the tar file successfully.") match_output = find_audio_matches( os.path.join(MOUNT_PATH, "test_jingle.mp3"), os.path.join(MOUNT_PATH, AUDIBLE_MAGIC_TOOLKIT, "bin"), os.path.join(MOUNT_PATH, AUDIBLE_MAGIC_CONFIG), ) print(f"Test match output is {match_output}") print("DONE!!!") aws_secret = modal.Secret.from_name("studio-aws") SECRETS = [ aws_secret, modal.Secret.from_name("datadog-metrics"), modal.Secret.from_name("api-callback-token"), modal.Secret.from_name("google-application-credentials-data"), ] def download_model_wrapper_c() -> None: AudibleMagicWorker.download_models() image = ( get_modal_base_image() .add_local_python_source("suno_utils", copy=True) .run_function(download_model_wrapper_c, secrets=SECRETS) .pip_install("protobuf==3.20.3") # Install other dependencies that might conflict .pip_install("google-api-core<2.0.0") .pip_install("google-cloud-aiplatform<2.0.0") # .pip_install("google") .pip_install("google-generativeai") # Then install vertexai .pip_install("vertexai==1.71.1") ) print("created image:", image) app = modal.App(APP_NAME, image=image) @app.cls( cpu=1.0, secrets=SECRETS, timeout=240, scaledown_window=240, retries=modal.Retries( max_retries=1, backoff_coefficient=2.0, initial_delay=5.0, ), min_containers=1 if DEPLOYMENT_TYPE == "dev" else 10, region="us-east", ) @modal.concurrent(max_inputs=2) class AudioUploadStub: """Audio upload stub and functions.""" def __init__(self): """Set up AudioDetectionStub.""" self.worker = AudibleMagicWorker() self.audio_upload_s3_client = boto3.client( "s3", aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID"), aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY"), ) self.modal_f_image_generator = modal.Cls.from_name( f"sdxl-{DEPLOYMENT_TYPE}", "StableDiffusion", )().generate_image self.modal_f_vocal_detection = modal.Cls.from_name( f"cycle-{DEPLOYMENT_TYPE}", "VocalDetectionStub", )().detect_vocals_from_audio self.modal_f_encode_audio = modal.Cls.from_name( f"cycle-{DEPLOYMENT_TYPE}", "CycleStub", )().encode_audio self.modal_f_whisper = modal.Cls.from_name( f"whisper-{DEPLOYMENT_TYPE}", "WhisperStub", )().transcribe self.modal_f_hoot_alignment = modal.Cls.from_name( f"videos-v2-{DEPLOYMENT_TYPE}", "HootAlignmentStub", )().get_aligned_lyrics self.modal_f_copyright_detector = modal.Cls.from_name( f"orchestrator-{DEPLOYMENT_TYPE}", "CopyrightDetectorStub" )().are_lyrics_copyrighted self.modal_f_audio_describer = AudioUploadDescriberStub().describe_audio @modal.method() def detect_copyright_audio(self, audio_s3_id: str) -> bool: """Detect if the given clip is copyrighted matrial.""" audios_file = f"{audio_s3_id}.mp3" with tempfile.TemporaryDirectory() as td: full_path = os.path.join(td, audios_file) with open(full_path, "wb") as tmp_file: retry_s3_download( "suno-data-uploads", f"studio/uploads/{audios_file}", tmp_file, ) detection_result = self.worker.is_audio_match_existing_music(full_path, []) tags = [ f"env:{DEPLOYMENT_TYPE}", f"copyright_detected:{detection_result}", ] statsd.increment("audible_magic.checks", tags=tags) return detection_result def _download_uploaded_file(self, item, upload_id, clip_id, upload_key): """Download the uploaded file from S3.""" bucket_name = "suno-uploads" s3_path = upload_key temp_input_path = f"/tmp/{upload_key.split('/')[-1]}" print("looking for file on:", bucket_name, s3_path) try: s3_client.get_object(Bucket=bucket_name, Key=s3_path) except Exception as e: print( f"Upload_id: {upload_id} resulted in exception: {e}. File doesn't exist on s3. Did the upload finish?" ) item.notify_progress( { "request_id": upload_id, "id": clip_id, "type": "upload_audio", "ok": 0, "error_type": "upload_failure_decode_audio", "error_message": "Can't fetch the uploaded audio.", } ) return None # Download the file from S3 to a temporary location self.audio_upload_s3_client.download_file(bucket_name, s3_path, temp_input_path) return temp_input_path def _handle_short_uploads(self, upload_id, sound_f, mp3_path): """Add padding to short mobile uploads if needed.""" if sound_f.duration_s < MIN_UPLOAD_DURATION and sound_f.duration_s >= 3: print(f"{upload_id} Upload audio is too short: {sound_f.duration_s}. Pad with silence.") silent = Audio.from_silence( duration_s=6.6 - sound_f.duration_s, sample_rate=sound_f.sample_rate, n_channels=2 ) sound_f = Audio.concatenate([silent, sound_f]) # Overwrite the existing file sound_f.to_hq_mp3(mp3_path) return sound_f def _upload_to_s3(self, file_path, bucket, s3_path): """Upload a file to S3.""" with open(file_path, "rb") as file: self.audio_upload_s3_client.upload_fileobj(file, bucket, s3_path) def _handle_fast_upload(self, item, upload_id, clip_id, s3_id, s3_new_path, title, duration): """Process a fast upload request (bypassing most checks).""" self.modal_f_encode_audio.spawn( audio=f"s3://{UPLOADS_S3_BUCKET}/{s3_new_path}", s3_npz_id=s3_id, encode_vae_version=VAEVersion.V_VAE_25_TUNED_2.value, ) # Default image image_s3_url = "s3://suno-data-uploads/studio/uploads/image_large_defualt_aura_13.jpeg" # Return result item.notify_progress( { "request_id": upload_id, "id": clip_id, "type": "upload_audio", "s3_id": s3_id, "image_s3": image_s3_url, "title": title, "duration": duration, "has_vocal": False, "hoot_lyrics": "", } ) print(f"{upload_id} Done, clip_id {clip_id}, fast upload complete.") return {} def _validate_duration(self, item, upload_id, clip_id, duration, title): """Check if the audio duration is within acceptable limits.""" # Check minimum duration print("validating duration:", item, "duration:", duration) curr_min_duration = MIN_UPLOAD_DURATION if item.metadata.get("is_vip_user", False): curr_min_duration = MIN_VIP_UPLOAD_DURATION if duration < curr_min_duration: print(f"{upload_id} Upload audio is too short: {duration}. Abort.") item.notify_progress( { "request_id": upload_id, "id": clip_id, "type": "upload_audio", "ok": 0, "error_type": "upload_failure_duration_short", "error_message": f"Uploaded audio is too short (currently {duration:.1f} seconds). Minimum duration is {MIN_UPLOAD_DURATION} seconds.", } ) print("too short:", duration) return False print("acceptable length:", duration) # Check for backdoor in title to bypass max duration check limit_max_duration = True if "MikeysLawSchool" in title or "SunoTimesInfinity" in title: limit_max_duration = False title = title.replace("MikeysLawSchool", "").replace("SunoTimesInfinity", "") # Check maximum duration curr_max_duration = ( MAX_FREE_USER_UPLOAD_DURATION if item.metadata.get("is_free_user", True) else MAX_UPLOAD_DURATION ) if item.metadata.get("is_vip_user", False): curr_max_duration = MAX_VIP_UPLOAD_DURATION if duration > curr_max_duration + 1 and limit_max_duration: # Leave some wiggle room print(f"{upload_id} Upload audio is too long: {duration}. Abort.") error_message = ( f"Uploaded audio is too long (currently {duration:.1f} seconds). " f"Maximum duration is {curr_max_duration} seconds. " ) if item.metadata.get("is_free_user", True): error_message += "Subscribe to PRO to unlock longer uploads." item.notify_progress( { "request_id": upload_id, "id": clip_id, "type": "upload_audio", "ok": 0, "error_type": "upload_failure_duration_long", "error_message": error_message, } ) return False return True def _prepare_title_and_prompt(self, title, upload_type): """Clean title and prepare image generation prompt.""" current_time = datetime.now() formatted_current_time = current_time.strftime("%Y-%m-%d_%H:%M:%S") if upload_type == "audio_recording": # Mic recording title = formatted_current_time default_prompt = ( "High contrast album art featuring a playful artistic illustration of 3d soundwaves." ) else: # File upload - strip and clean title = title.split(".")[0].strip() default_prompt = ( f"High contrast album art featuring an abstract artistic illustration of {title}." ) return title, default_prompt def _start_image_generation(self, clip_id, prompt): """Start image generation process.""" return self.modal_f_image_generator.spawn( json.dumps( { "id": clip_id, "prompt_text": prompt, "metadata": {}, "model_name": "user_upload", } ) ) def _start_vocal_detection(self, s3_id): """Start vocal detection process.""" return self.modal_f_vocal_detection.spawn(f"s3://suno-data-uploads/studio/uploads/{s3_id}.mp3") def _trim_recording(self, sound_f): """Trim audio recording by removing margins.""" curr_duration = sound_f.duration_s margin_time = 0.2 return sound_f.get_segment(from_s=margin_time, to_s=curr_duration - margin_time) def _validate_unique_artwork(self, item, upload_id, clip_id, audio_path, claimed_artist_names): """Check if the audio matches existing music.""" # Check for backdoor in title to bypass artwork check # Check for match with existing artwork is_existing_artwork = self.worker.is_audio_match_existing_music(audio_path, claimed_artist_names) tags = [ f"env:{DEPLOYMENT_TYPE}", f"copyright_detected:{is_existing_artwork}", ] statsd.increment("audible_magic.checks", tags=tags) if is_existing_artwork: print(f"{upload_id} Detected matching existing artwork. Abort.") item.notify_progress( { "request_id": upload_id, "id": clip_id, "type": "upload_audio", "ok": 0, "error_type": "upload_failure_artwork", "error_message": "Uploaded audio matches existing work of art.", } ) return False return True def _start_audio_encoding(self, s3_id, s3_path): """Start audio encoding process.""" self.modal_f_encode_audio.spawn( audio=f"s3://{UPLOADS_S3_BUCKET}/{s3_path}", s3_npz_id=s3_id, encode_vae_version=VAEVersion.V_VAE_25_TUNED_2.value, ) def _get_image_result(self, upload_id, image_f_call): """Get image generation result with fallback.""" try: image_s3_url = image_f_call.get(timeout=10).get(timeout=10) except Exception as e: print(f"{upload_id} Failed to get image: {str(e)}") # Fallback to default image image_s3_url = ( "s3://suno-data-uploads/studio/uploads/image_large_defualt_aura_13.jpeg" # [sic] ) print(f"{upload_id} image located", image_s3_url) return image_s3_url def _start_gemini_transcription(self, upload_id, clip_id, s3_id): payload = { "id": clip_id, "metadata": {"s3_id": s3_id}, } return self.modal_f_audio_describer.spawn(json.dumps(payload)) def _transcribe_with_gemini(self, upload_id, clip_id, s3_id): payload = { "id": clip_id, "metadata": {"s3_id": s3_id}, } print("using gemini on:", clip_id, s3_id) transcription, description, display_tags, has_vocals = "", "", "", False try: response = self.modal_f_audio_describer.spawn(json.dumps(payload)).get(timeout=60) except Exception as e: traceback.print_exc() print( f"{upload_id} Failed to get transcription / description / display_tags from Gemini: {(e)}" ) return transcription, description, display_tags, has_vocals print(f"{upload_id} Failed to get transcription / description from Gemini: {(e)}") return transcription, description, display_tags print("gemini response:", response) try: transcription = response["transcription"] description = response["description"] display_tags = response["genre_summary"] except (KeyError, IndexError) as e: print( "failed to get transcription, description and genre_summary from gemini response:", response, ) has_vocals = transcription.lower().strip() == "[instrumental]" return transcription, description, display_tags, has_vocals def _get_and_validate_lyrics( self, item, upload_id, clip_id, s3_id, lyrics_transcription_call, upload_type, is_bypass_in_title, ) -> Tuple[str, bool, str, str, bool]: """Get lyrics transcription and check for copyright issues.""" tic = time.time() print("lyrics call.get") try: response = lyrics_transcription_call.get(timeout=120) except TimeoutError as e: toc = time.time() print(f"lyrics transcription call timed out after {toc - tic}s") raise e toc = time.time() print(f"lyrics_transcription_call.get in {round(toc - tic, 2)}s") transcribed_lyrics, lyrics_are_copyrighted = "", False try: transcribed_lyrics = response["transcription"] inferred_description = response["description"] display_tags = response["genre_summary"] except (KeyError, IndexError) as e: print( "failed to get transcription, description and genre_summary from gemini response:", response, ) has_vocals = transcribed_lyrics.lower().strip() == "[instrumental]" if transcribed_lyrics and len(transcribed_lyrics.strip()) > 0 and not is_bypass_in_title: print(f"{upload_id} Checking transcribed lyrics for copyright...") copyright_tic = time.time() lyrics_are_copyrighted = self._are_transcribed_lyrics_copyrighted( item, upload_id, clip_id, transcribed_lyrics ) copyright_toc = time.time() print( f"{upload_id} lyrics are copyrighted: {lyrics_are_copyrighted}, {copyright_toc - copyright_tic}" ) toc = time.time() print( f"lyrics transcription took: {round(toc - tic, 2)}, using gemini, lyrics are copyrighted:", lyrics_are_copyrighted, ) print("display tags:", display_tags) return ( transcribed_lyrics, lyrics_are_copyrighted, inferred_description or "", display_tags or "", has_vocals, ) def _notify_success(self, item, upload_id, clip_id, status_message): item.notify_progress( { "request_id": upload_id, "id": clip_id, "type": "upload_audio", "ok": 1, "status": status_message, } ) def _are_transcribed_lyrics_copyrighted(self, item, upload_id, clip_id, transcribed_lyrics): try: strict_mode = transcribed_lyrics and len(transcribed_lyrics.split()) <= 10 lyrics_are_copyrighted = self.modal_f_copyright_detector.spawn( transcribed_lyrics, strict_mode ).get(timeout=30) except Exception as e: traceback.print_exc() print(f"{upload_id} Failed to check lyrics copyright: {str(e)}") # Continue processing if copyright check fails if lyrics_are_copyrighted: print(f"{upload_id} Detected copyrighted lyrics in upload. Abort.") item.notify_progress( { "request_id": upload_id, "id": clip_id, "type": "upload_audio", "ok": 0, "error_type": "upload_failure_lyrics_copyright", "error_message": "Uploaded audio contains copyrighted lyrics.", } ) return True else: return False def _detect_bypass_and_reformat_title(self, title: str) -> Tuple[str, bool]: bypasses = ["MurphysLawSchool", "SunoGivesYouWings"] is_bypass_in_title = any(bypass in title for bypass in bypasses) for bypass in bypasses: title = title.replace(bypass, "") print("is bypass in title:", is_bypass_in_title) return title, is_bypass_in_title @modal.method() def process_uploaded_audio(self, queue_item_json: str): """Upload audio to S3, with a special prefix.""" timing_data = {} print("queue_item_json:", queue_item_json) total_start_time = time.time() item = QueueItem(**json.loads(queue_item_json)) # type: ignore missing-argument print("queue item:", item) upload_id = item.id title = item.title or "Uploaded Audio" title, is_bypass_in_title = self._detect_bypass_and_reformat_title(title) # this is the key on s3 that maps to the uploaded audio file upload_key = item.metadata.get("upload_key") upload_type = item.metadata.get("upload_type", "file_upload") use_gemini = item.metadata.get("use_gemini", USE_GEMINI_BY_DEFAULT) print("use gemini:", use_gemini) claimed_artist_names = item.metadata.get("claimed_artist_names", []) # TODO this flag could allow uploads to be abused! # Revert to dev only as soon as we find an alternative for fast generative stem export uploads is_fast_upload = item.metadata.get("is_fast_upload", False) # Generate new identifiers clip_id = str(uuid.uuid4()) s3_id = f"{clip_id}" s3_new_path = f"studio/uploads/{s3_id}.mp3" # Define the S3 path for the uploaded file with timer(timing_data, "download_uploaded_file"): temp_file_path = self._download_uploaded_file(item, upload_id, clip_id, upload_key) if not temp_file_path: # didn't find file on s3 so nothing we can do here. return {"success": False} # Convert and validate the audio file with timer(timing_data, "process_audio_file"): mp3_path, sound_f = _process_audio_file(item, upload_id, clip_id, temp_file_path) if not sound_f: return {"success": False} self._notify_success(item, upload_id, clip_id, "passed_audio_processing") # for uploads < 6s, pad with silence ## start handle_short_uploads sound_f = self._handle_short_uploads(upload_id, sound_f, mp3_path) if upload_type == "audio_recording": with timer(timing_data, "trim_recording"): sound_f = self._trim_recording(sound_f) duration = sound_f.duration_s temp_item = QueueItem(id=clip_id, ids=None, metadata={}) # add a volume normalization step try: if sound_f.loudness > -32: with timer(timing_data, "normalize volume"): sound_f = sound_f.normalize_volume(target_db=-14) except Exception as e: print(f"{upload_id} Error normalizing volume: {e}, duration: {sound_f.duration_s}") with timer(timing_data, "write_audio_only"): self.worker._write_audio_only(temp_item, sound_f) print(f"{upload_id}, s3_new_path: {s3_new_path}. Audio duration is {duration}s.") gemini_transcription_tic = time.time() print("starting gemini transcription") lyrics_transcription_call = self._start_gemini_transcription(upload_id, clip_id, s3_id) if is_fast_upload: # bypass everything and just encode the audio return self._handle_fast_upload( item, upload_id, clip_id, s3_id, s3_new_path, title, duration ) # Validate audio duration with timer(timing_data, "validate duration"): if not self._validate_duration(item, upload_id, clip_id, duration, title): return {"success": False} # Clean up title and prepare for image generation title, default_prompt = self._prepare_title_and_prompt(title, upload_type) # Start parallel processing tasks with timer(timing_data, "starting parallel calls"): image_f_call = self._start_image_generation(clip_id, default_prompt) # tkb # check if the piece matches existing artwork with timer(timing_data, "validating unique artwork"): is_unique_artwork = is_bypass_in_title or self._validate_unique_artwork( item, upload_id, clip_id, mp3_path, claimed_artist_names ) if not is_unique_artwork: return {"success": False} self._notify_success(item, upload_id, clip_id, "passed_artist_moderation") # At this stage the file is valid # The upload will be successful # Trigger the rest of work with timer(timing_data, "start audio encoding"): self._start_audio_encoding(s3_id, s3_new_path) with timer(timing_data, "getting image result"): image_s3_url = self._get_image_result(upload_id, image_f_call) with timer(timing_data, "get and validate lyrics"): ( transcribed_lyrics, lyrics_are_copyrighted, inferred_description, display_tags, has_vocals_in_upload, ) = self._get_and_validate_lyrics( item, upload_id, clip_id, s3_id, lyrics_transcription_call, upload_type, is_bypass_in_title=is_bypass_in_title, ) gemini_transcription_toc = time.time() print("total gemini time:", round(gemini_transcription_toc - gemini_transcription_tic, 2)) # toc = time.time() # timing_data["getting image result"] = toc - tic # print("transcribed lyrics:", transcribed_lyrics) if lyrics_are_copyrighted: print("get and validate lyrics did not pass, returning early") return {"success": False} # send back the information print("notifying progress of:", display_tags) item.notify_progress( { "request_id": upload_id, "id": clip_id, "type": "upload_audio", "s3_id": s3_id, "image_s3": image_s3_url, "title": title, "duration": duration, "has_vocal": has_vocals_in_upload, "hoot_lyrics": transcribed_lyrics, # maybe this isn't the best key "inferred_description": inferred_description, "display_tags": display_tags, } ) print(f"{upload_id} Done, clip_id {clip_id}, TOTAL TIME: {time.time() - total_start_time:.3f}s") # Callback _display_timing_data(timing_data) return {"success": True} from contextlib import contextmanager @contextmanager def timer(timing_data, label): """ Context manager for timing code execution. Args: timing_data: Dictionary to store timing results label: Label for the timing measurement Example: timing_data = {} with timer(timing_data, "getting vocal detection result"): result = some_expensive_operation() """ tic = time.time() print("starting:", label) try: yield finally: toc = time.time() timing_data[label] = toc - tic print("ending:", label, f"{round(toc - tic, 2)}s") def _display_timing_data(timing_data): print("TIMING DATA:") total_time = sum(timing_data.values()) # Find the longest key name for formatting max_key_length = max(len(k) for k in timing_data.keys()) for k, v in timing_data.items(): percentage = v / total_time * 100 # Create bar: one '.' for each 5% (rounded) bar_length = round(percentage / 5) bar = "." * bar_length # Format with consistent spacing print(f"{k:<{max_key_length}}: {round(v, 2):>8}, {round(percentage, 2):>6}% |{bar}") print("Total logged time:", total_time) def _process_audio_file(item, upload_id, clip_id, temp_input_path): """Convert and load audio file, ensuring it's a valid MP3.""" try: # Convert file to MP3 if needed if not temp_input_path.endswith(".mp3"): mp3_path = f"/tmp/{clip_id}.mp3" subprocess.run( [ "ffmpeg", "-loglevel", "error", "-i", temp_input_path, "-vn", "-q:a", "0", "-ar", "48000", "-y", mp3_path, ], check=True, ) else: mp3_path = temp_input_path # Load the audio file if temp_input_path.endswith(".wav"): sound_f = Audio.from_file(temp_input_path, n_channels=2) else: sound_f = Audio.from_file(mp3_path, n_channels=2) print("upon loading, duration:", sound_f.duration_s) return mp3_path, sound_f except Exception as e: print(f"{upload_id} Upload audio: {clip_id} failed with error: {e}") traceback.print_exc() item.notify_progress( { "request_id": upload_id, "id": clip_id, "type": "upload_audio", "ok": 0, "error_type": "upload_failure_decode_audio", "error_message": "Can't parse uploaded audio. Source is corrupted.", } ) return None, None ################################################################################ def _load_audio_as_opus(local_fp): from vertexai.generative_models import Part # type: ignore unresolved-import with tempfile.TemporaryDirectory() as tempdir: if local_fp.endswith(".opus"): out_local_fp = local_fp else: out_local_fp = os.path.join(tempdir, "audio.opus") _ = Audio.from_file(local_fp, sample_rate=48_000, byte_width=2, n_channels=2).to_opus( out_local_fp ) with open(out_local_fp, "rb") as f: audio_bytes = f.read() p = Part.from_data(audio_bytes, mime_type="audio/mpeg") return p AUDIO_UPLOAD_DESCRIBER_SYSTEM_PROMPT = f"""Your task is to describe a musical piece I've attached in a way that would allow an AI music generator to recreate it as accurately as possible. Your response should contain three elements: A Description, a Genre Summary, and a Transcription. * Description Guidelines: - Assume you're instructing a powerful AI music model that has no knowledge of existing artists or songs - Focus on describing the musical elements you can clearly identify (don't guess or include uncertain details) - Include specific details about: * Genre and style * Instruments used * Chord progressions * Vocal qualities and techniques * Song structure * Melody characteristics * Key and tempo * Production elements (mixing, effects, etc.) - The AI already understands lyrical structure formats like "[Verse 1]", so you don't need to explain these - Don't use parentheses or brackets in your descriptions - Never use hedging language like "maybe" or "likely". Only include details you're confident about - Your description should be at the level of a professional music producer: comprehensive and technically precise. - Never start your descriptions with the phrase "This is a". In the first sentence, it's ok to use a noun phrase rather than a complete sentence. * Genre Summary Guidelines: - Summarize the song in terms of the two or three genres or phrases that best describe it. These could be general like "pop" or "rock", or specific like "electroswing" or "silk shirt R&B", as long as they're accurate. - Structure this output as a string with the individual genres separated by commas, for example: "rock, aggressive, baritone male vocals". * Transcription Guidelines - Transcribe the lyrics in full. - Separate the song sections like verses and choruses with newlines - Add section markers in square brackets (like "[Verse]" and ["Chorus"]) at the head of each section - Do not add timestamps - If there are no vocals present, simply give "[Instrumental]" as the transcription. - If there are any vocal adlibs present, transcribe them within parentheses, e.g. "(ooh-yeah)". - Listen especially carefully to vocal adlibs and be sure to keep track of exactly where you are in the song when transcribing them. * General Guidelines - Make sure that your response is valid JSON and conforms to the attached schema. - Make sure that your total response is less than {MAX_GEMINI_OUTPUT_TOKENS} tokens. """ class AudioDescriptionResponse(BaseModel): description: str = Field(description="Professional producer-level description of the musical style") transcription: str = Field(description="Transcription of the lyrics") genre_summary: str = Field( description="Two or three genres or descriptive terms describing the overall feel of the music" ) def _get_temp_path(upload_key): filename = upload_key.split("/")[-1] return f"/tmp/{filename}" def _load_google_creds(environ): creds_json = os.environ["GOOGLE_APPLICATION_CREDENTIALS_DATA"] creds_path = "/tmp/google_creds.json" with open(creds_path, "w") as f: f.write(creds_json) os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = creds_path @app.cls( cpu=1.0, secrets=SECRETS, timeout=240, scaledown_window=240, retries=modal.Retries( max_retries=1, backoff_coefficient=2.0, initial_delay=5.0, ), min_containers=1 if DEPLOYMENT_TYPE == "dev" else 10, region="us-east", ) @modal.concurrent(max_inputs=1000) class AudioUploadDescriberStub: """Audio upload stub and functions.""" def __init__(self): """Set up AudioDetectionStub.""" tic = time.time() _load_google_creds(os.environ) toc = time.time() print(f"Initialized AudioUploadDescriberStub in {toc - tic}s") @modal.method() def describe_audio(self, queue_item_json: str): tic = time.time() from google import genai # type: ignore unresolved-import from google.genai import types # type: ignore unresolved-import print("genai version:", genai.__version__) from suno_utils.audio import Audio toc = time.time() GEMINI_MODEL = "gemini-2.5-flash-preview-04-17" t0 = time.time() tic = time.time() client = genai.Client(vertexai=True, project="ml-gemini-455703", location="us-central1") toc = time.time() def _load_audio_as_prompt(local_fp): print("in _load_audio_as_prompt:", local_fp) with tempfile.TemporaryDirectory() as tempdir: if local_fp.endswith(".opus"): print("found local_fp opus:", local_fp) out_local_fp = local_fp else: tic = time.time() out_local_fp = os.path.join(tempdir, "audio.opus") _ = Audio.from_file( local_fp, sample_rate=48_000, byte_width=2, n_channels=2 ).to_opus(out_local_fp) toc = time.time() print(f"converted local_fp: {local_fp} in {round(toc - tic, 2)}s") with open(out_local_fp, "rb") as f: audio_bytes = f.read() print(f"successfully read: {len(audio_bytes)} bytes") p = types.Part.from_bytes(data=audio_bytes, mime_type="audio/mpeg") return p item = QueueItem(**json.loads(queue_item_json)) # type: ignore missing-argument audio_url = item.metadata.get("audio_url") s3_id = item.metadata.get("s3_id") if audio_url: print("describe_audio working off audio_url") temp_input_path = _get_temp_path(audio_url) tic = time.time() try: response = requests.get(audio_url) with open(temp_input_path, "wb") as f: f.write(response.content) except Exception as e: print(f"Couldn't download file:{e}") return {} opus_file = _load_audio_as_prompt(temp_input_path) elif s3_id: # TKTKTK print("descrbe_audio working off s3_id:", s3_id) temp_input_path = f"/tmp/{s3_id}.opus" s3_path = f"studio/uploads/{s3_id}.opus" tic = time.time() try: waiter = s3_client.get_waiter("object_exists") waiter.wait( Bucket="suno-data-uploads", Key=s3_path, WaiterConfig={ "Delay": 0.2, "MaxAttempts": 50, }, ) with open(temp_input_path, "wb") as f: s3_client.download_fileobj("suno-data-uploads", s3_path, f) toc = time.time() print(f"downloaded {temp_input_path} in {round(toc - tic, 2)}s") opus_file = _load_audio_as_prompt(temp_input_path) except Exception as e: print(f"Failed to download opus file from s3 and convert to Gemini prompt: {e}") return AudioDescriptionResponse( description="", transcription="", genre_summary="" ).dict() t1 = time.time() print(f"calling gemini at {t1}") try: response = client.models.generate_content( model=GEMINI_MODEL, contents=[opus_file, AUDIO_UPLOAD_DESCRIBER_SYSTEM_PROMPT], config=types.GenerateContentConfig( candidateCount=1, audioTimestamp=False, thinkingConfig=types.ThinkingConfig( includeThoughts=False, thinkingBudget=0, ), temperature=0.0, response_mime_type="application/json", response_schema=AudioDescriptionResponse, max_output_tokens=MAX_GEMINI_OUTPUT_TOKENS, frequency_penalty=0.1, ), ) t2 = time.time() print(f"Gemini call completed in: {t2 - t1:.1f}s") except Exception as e: print(f"Got exception from Gemini: {e}") return AudioDescriptionResponse(description="", transcription="", genre_summary="").dict() parsed_response = _parse_response(response) print("TRANSCRIPT" + "-" * 20) print(parsed_response["transcription"]) print("-" * 20) if "description" in parsed_response and "transcription" in parsed_response: print("Gemini returned a valid response") else: print(f"Gemini returned an invalid response on attempt: {parsed_response}") raise RuntimeError() print(f"describe_audio completed in {t2 - t0:.1f} total seconds") return parsed_response def _parse_response(resp: str) -> dict: json_text = resp.candidates[0].content.parts[0].text assert isinstance(json_text, str) print("trying to parse json_text:", json_text) try: data = json.loads(json_text) return data except json.JSONDecodeError as e: print(f"couldn't json decode: {json_text}") if not json_text.endswith('"}'): json_text += '"}"' try: data = json.loads(json_text, strict=False) return data except json.JSONDecodeError as e: print(f"couldn't json decode: {json_text}") return AudioDescriptionResponse(description="", transcription="", genre_summary="").dict() def _test_audio_upload_describer(audio_upload_describer_stub): print("-" * 80) test_data = { "id": "906b108e-0a92-404f-b10e-169fd84b1353", "metadata": { "audio_url": "https://cdn1.suno.ai/736dbf7b-172f-4141-b52a-6a5f2c215b9a.mp3", }, } tic = time.time() result = audio_upload_describer_stub.describe_audio.remote(json.dumps(test_data)) toc = time.time() elapsed_time = toc - tic print(f"got audio upload description in {elapsed_time}s") print("got result:", result) assert "description" in result assert result["description"] # expected_personas = ["amateur", "enthusiast", "producer"] # found_personas = [] # for item in result["description"]: # assert "persona" in item # assert "description" in item # found_personas.append(item["persona"]) # assert set(expected_personas) == set(found_personas) assert "transcription" in result assert result["transcription"] assert "genre_summary" in result assert result["genre_summary"] def _test_process_audio_upload(audio_upload_stub): print("-" * 80) test_data = { "id": "6a39ef27-d555-43d4-88d9-fbe6d5369462", "metadata": {"upload_key": "raw_uploads/6a39ef27-d555-43d4-88d9-fbe6d5369462.mp3"}, } tic = time.time() result = audio_upload_stub.process_uploaded_audio.remote(json.dumps(test_data)) print("result type:", type(result)) result = result if isinstance(result, dict) else result.get() toc = time.time() elapsed_time = toc - tic print(f"got audio upload description in {elapsed_time}s") print("got result:", result) assert result["success"] is True # assert result["transcription"] # assert result["description"] # assert result["display_tags"] # assert result["has_vocal"] def _test_copyrighted_lyrics(audio_upload_stub): print("-" * 80) test_data = { "id": "906b108e-0a92-404f-b10e-169fd84b1353", "metadata": { "upload_key": "raw_uploads/906b108e-0a92-404f-b10e-169fd84b1353.mp3", "use_gemini": True, }, } tic = time.time() # result = audio_upload_stub.process_uploaded_audio.remote(json.dumps(test_data)) result = audio_upload_stub.process_uploaded_audio.remote(json.dumps(test_data)) print("result type:", type(result)) result = result if isinstance(result, dict) else result.get() toc = time.time() elapsed_time = toc - tic print(f"got audio upload description in {elapsed_time}s") print("got result:", result) assert result["success"] is False def _test_copyrighted_audio(audio_upload_stub): print("-" * 80) flashing_lights_data = { "id": "831f6731-2ad6-4406-a899-eb906f685144", "title": "flashing lights", "metadata": { "upload_key": "raw_uploads/831f6731-2ad6-4406-a899-eb906f685144.mp3", "use_gemini": True, }, } tic = time.time() # result = audio_upload_stub.process_uploaded_audio.remote(json.dumps(flashing_lights_data)) result = audio_upload_stub.process_uploaded_audio.remote(json.dumps(flashing_lights_data)) print("result type:", type(result)) result = result if isinstance(result, dict) else result.get() toc = time.time() elapsed_time = toc - tic print(f"got audio upload description in {elapsed_time}s") print("got result:", result) assert result["success"] is False def _test_title_bypass(audio_upload_stub): print("-" * 80) flashing_lights_data = { "id": "831f6731-2ad6-4406-a899-eb906f685144", "title": "flashing lights MurphysLawSchool", "metadata": { "upload_key": "raw_uploads/831f6731-2ad6-4406-a899-eb906f685144.mp3", "use_gemini": True, }, } tic = time.time() # result = audio_upload_stub.process_uploaded_audio.remote(json.dumps(flashing_lights_data)) result = audio_upload_stub.process_uploaded_audio.remote(json.dumps(flashing_lights_data)) print("result type:", type(result)) result = result if isinstance(result, dict) else result.get() toc = time.time() elapsed_time = toc - tic print(f"got audio upload description in {elapsed_time}s") print("got result:", result) assert result["success"] is True # assert result["transcription"] # assert result["description"] # assert result["display_tags"] # assert result["has_vocal"] def _test_connectapella(audio_upload_stub): print("-" * 80) test_data = { "id": "27829588-e14a-4445-823f-a79eeb7b4ee3", "metadata": {"upload_key": "raw_uploads/27829588-e14a-4445-823f-a79eeb7b4ee3.mp3"}, } tic = time.time() result = audio_upload_stub.process_uploaded_audio.remote(json.dumps(test_data)) print("result type:", type(result)) result = result if isinstance(result, dict) else result.get() toc = time.time() elapsed_time = toc - tic print(f"got audio upload description in {elapsed_time}s") print("got result:", result) assert result["success"] is True def _test_nirvana(audio_upload_stub): print("-" * 80) test_data = { "id": "0fe07094-373d-4b0e-9f3b-e271137895a6", "metadata": {"upload_key": "raw_uploads/27829588-e14a-4445-823f-a79eeb7b4ee3.mp3"}, } tic = time.time() result = audio_upload_stub.process_uploaded_audio.remote(json.dumps(test_data)) print("result type:", type(result)) result = result if isinstance(result, dict) else result.get() toc = time.time() elapsed_time = toc - tic print(f"got audio upload description in {elapsed_time}s") print("got result:", result) assert result["success"] is True @app.local_entrypoint() def main(): audio_upload_stub = AudioUploadStub() audio_upload_describer_stub = AudioUploadDescriberStub() _test_process_audio_upload(audio_upload_stub) _test_connectapella(audio_upload_stub) _test_title_bypass(audio_upload_stub) _test_copyrighted_audio(audio_upload_stub) _test_audio_upload_describer(audio_upload_describer_stub) _test_copyrighted_lyrics(audio_upload_stub) _test_nirvana(audio_upload_stub) # this should trigger a detection match result = audio_upload_stub.detect_copyright_audio.remote("6afed9a8-a00b-45a7-8340-394a9b820fd8") print("Audio upload test result:", result) assert result test_data = { "id": "906b108e-0a92-404f-b10e-169fd84b1353", "metadata": { "upload_key": "raw_uploads/906b108e-0a92-404f-b10e-169fd84b1353.mp3", "use_gemini": False, }, } # try hoot _ = audio_upload_stub.process_uploaded_audio.remote(json.dumps(test_data)) # try whisper test_data["metadata"]["upload_type"] = "audio_recording" _ = audio_upload_stub.process_uploaded_audio.remote(json.dumps(test_data)) print("Finish test!")