""" This file is responsible for loading the audio and video files from S3. It also contains the logic for loading the tokens and latents from S3. It also contains the version conversion logic for the tokens and latents. Discrete codes are saved in the id.npz files. Continous codes are saved in the id_vae.npz files. - The id_vae.npz files contain the v_vae field, which indicates the version of the vae. - The vae_latents field contains the latents. - When loading a vae, if the loading key doesn't match the v_vae field, we will need to convert the latents to the new vae version. - We save the latents in float16 format. """ import os import re import tempfile import traceback from functools import partial from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Any, List, Optional, Tuple import botocore.exceptions import numpy as np from PIL import Image, ImageOps from suno_utils.audio import Audio from suno_utils.gpt import chirp_v2, chirp_v2_5 from suno_utils.tasks import dac_2c, dac_2c_12cb, mert_25 from suno_utils.worker.schema import QueueItem from suno_utils.worker.settings import s3_client from suno_utils.worker.utils import retry_decorator from suno_utils.worker.video_renderer import generate_video DEFAULT_S3_FOLDER = "studio/uploads/" DEFAULT_S3_BUCKET = "suno-data-uploads" VERSION_GPT_MAPPING = { "2.0": chirp_v2, "3.0": chirp_v2_5, "3.5": chirp_v2_5, "4.0": chirp_v2_5, "5.0": chirp_v2_5, } VERSION_EMBEDDING_RATE_MAPPING = { "2.0": dac_2c.EMBEDDING_RATE, "3.0": dac_2c_12cb.EMBEDDING_RATE, "3.5": dac_2c_12cb.EMBEDDING_RATE, "4.0": dac_2c_12cb.EMBEDDING_RATE, "5.0": mert_25.EMBEDDING_RATE, } # boto3 download_fileobj writes from the current position in the file # so if transfer fails after starting to write, we need to clean up previous attempt def truncate_then_download(Bucket, Key, Fileobj, ExtraArgs=None, Callback=None, Config=None): if hasattr(Fileobj, "truncate") and hasattr(Fileobj, "seek"): Fileobj.truncate(0) Fileobj.seek(0) return s3_client.download_fileobj(Bucket, Key, Fileobj, ExtraArgs, Callback, Config) # retries should be very quick retry_s3_download = retry_decorator(3, wait_seconds=5)(truncate_then_download) fast_retry_s3_download = retry_decorator(2, wait_seconds=1)(truncate_then_download) retry_s3_upload = retry_decorator(3, wait_seconds=5)(s3_client.upload_file) def get_latents( filename: str, audio_id: str, from_s: Optional[float] = None, to_s: Optional[float] = None, target_vae_version: str = "", ) -> Tuple[np.ndarray | None, str | None]: """If any vae doesn't exist or doesn't match. Will fail and return None. If target_vae_version is "", we will just lazy load the vae as primary key is. """ tokens, vae_version = _get_compatible_latents(filename, audio_id, target_vae_version) if tokens is not None: rate = 25 # TODO: we will make this configurable when we need to... l = int(from_s * rate) if from_s is not None else 0 r = int(to_s * rate) if to_s is not None else tokens.shape[0] return tokens[l:r], vae_version return None, None def _get_compatible_latents( filename: str, audio_id: str, target_vae_version: str ) -> Tuple[np.ndarray | None, str | None]: """Get the tokens for the given audio_id and target_version from S3.""" try: with open(filename, "wb") as tmp_file: fast_retry_s3_download(DEFAULT_S3_BUCKET, f"studio/uploads/{audio_id}_vae.npz", tmp_file) npz = np.load(filename) existing_vae_version = str(npz["v_vae"].item()) if target_vae_version == "": # this is an edge case -- by default, we just lazy load the vae as they are target_vae_version = existing_vae_version if existing_vae_version == target_vae_version: return npz.get("vae_latents", None), target_vae_version else: # check if we have a compatible vae version compatible_vae_latents_key = f"vae_latents_{target_vae_version}" if compatible_vae_latents_key in npz: print(f"Found compatible vae version {target_vae_version} for {audio_id}, loading...") return npz[compatible_vae_latents_key], target_vae_version else: print( f"Found incompatible vae version {existing_vae_version} for {audio_id}, need to convert!" ) return None, None except Exception as e: print(f"No valid vae latents found for {audio_id}, error: {e}") traceback.print_exc() return None, None def get_tokens( filename: str, audio_id: str, target_version: str, from_s: Optional[float] = None, to_s: Optional[float] = None, ) -> np.ndarray | None: """Returns the tokens for the given audio_id, target_version, and optionally continue_at. Both semantic and audio tokens are included""" tokens = _get_compatible_tokens(filename, audio_id, target_version) if tokens is not None: rate = VERSION_EMBEDDING_RATE_MAPPING.get(target_version, 25) l = int(from_s * rate) if from_s is not None else 0 r = int(to_s * rate) if to_s is not None else tokens.shape[0] return tokens[l:r] def _get_compatible_tokens(filename: str, audio_id: str, target_version: str) -> np.ndarray | None: """Get the tokens for the given audio_id and target_version from S3.""" try: with open(filename, "wb") as tmp_file: retry_s3_download(DEFAULT_S3_BUCKET, f"studio/uploads/{audio_id}.npz", tmp_file) npz = np.load(filename) if (target_version == "2.0" or target_version == "1.0") and "v1_raw" in npz: print(f"Found compatible {target_version}.npz file for {audio_id}") return npz["v1_raw"] elif f"v{target_version}_raw" in npz: print(f"Found compatible {target_version}.npz file for {audio_id}") return npz[f"v{target_version}_raw"] # 5.0 is using semantic only elif target_version == "5.0": for potential_matching_version in [ "4.0", "3.5", "3.0", "2.0", ]: if f"v{potential_matching_version}_raw" in npz: print( f"Found compatible .npz file {npz[f'v{potential_matching_version}_raw'].shape} for " f"{potential_matching_version} target {target_version}, " f"{audio_id}" ) # get only the semantic tokens return npz[f"v{potential_matching_version}_raw"][:, 0:1] # this is a hack for 3.0 and 3.5 and potentially 4.0 # since they are the same codec... # we can translate them and cross load for now # we actually don't want to cycle them multiple times # particularly for concatenation elif target_version in ["3.0", "3.5", "4.0"]: for potential_matching_version in ["3.0", "3.5", "4.0"]: if f"v{potential_matching_version}_raw" in npz: print( f"Found compatible .npz file for with " f"{potential_matching_version} target {target_version}, " f"{audio_id}" ) return npz[f"v{potential_matching_version}_raw"] # if didn't find --> we still need to convert! print("Found incompatible .npz file for", audio_id, ". Converting...") if "model_version" in npz: old_arch, old_major, _, _ = npz["model_version"].item().split(".") old_version = f"{old_arch}.{old_major}" old_tokens = npz[f"v{old_version}_raw"] if old_version != "5.0": audio = VERSION_GPT_MAPPING[f"{old_version}"].codec_decode(old_tokens[:, 1:]) else: # directly load the audio from s3; note for auk/sem) we need to cycle from audio cause sem only audio = Audio.from_s3( f"s3://suno-data-uploads/studio/uploads/{audio_id}.mp3", n_channels=2 ) else: old_tokens = npz["v1_raw"] print("got old tokens") if old_tokens.shape[1] == 13: audio = VERSION_GPT_MAPPING["3.0"].codec_decode(old_tokens[:, 1:]) else: audio = VERSION_GPT_MAPPING["2.0"].codec_decode(old_tokens[:, 1:]) print("converted old tokens to audio") codec_tokens = VERSION_GPT_MAPPING[f"{target_version}"].codec_encode(audio) semantic_tokens = old_tokens[:, 0:1] if semantic_tokens.shape[0] != codec_tokens.shape[0]: # this can happen...due to edge time in decoding, we just trim semantic min_code_shape = min(semantic_tokens.shape[0], codec_tokens.shape[0]) semantic_tokens = semantic_tokens[:min_code_shape, :] codec_tokens = codec_tokens[:min_code_shape, :] npz = { **npz, f"v{target_version}_raw": np.concatenate([semantic_tokens, codec_tokens], axis=1), } np.savez(filename, **npz) retry_s3_upload(filename, DEFAULT_S3_BUCKET, f"studio/uploads/{audio_id}.npz") return npz[f"v{target_version}_raw"] except Exception as e: print("No valid .npz file found for", audio_id) print(e) traceback.print_exc() class S3Loader: def _load_list_of_audio_prompts( self, item: QueueItem, model_version: str, list_key: str ) -> list[np.ndarray]: """Load the playlist audio prompt based on the QueueItem input.""" playlist_audio_tokens = [] for curr_idx, curr_clip_id in enumerate(item.metadata[list_key]): with tempfile.NamedTemporaryFile(suffix=".npz") as temp_file: arch, major, _, _ = model_version.split(".") target_version = f"{arch}.{major}" tokens = _get_compatible_tokens(temp_file.name, curr_clip_id, target_version) if tokens is None: raise ValueError(f"Couldn't fetch the tokens from {curr_clip_id}") rate = VERSION_EMBEDDING_RATE_MAPPING.get(target_version, 25) # TODO: these flags are not currently set for multi-artist consistency playlist_start_s = item.metadata.get("playlist_clip_starts", None) playlist_end_s = item.metadata.get("playlist_clip_ends", None) from_s = ( playlist_start_s[curr_idx] if playlist_start_s is not None and curr_idx < len(playlist_start_s) else None ) to_s = ( playlist_end_s[curr_idx] if playlist_end_s is not None and curr_idx < len(playlist_end_s) else None ) start_index = int(from_s * rate) if from_s is not None else 0 end_index = int(to_s * rate) if to_s is not None else tokens.shape[0] assert end_index > start_index, f"start_index: {start_index}, end_index: {end_index}" playlist_audio_tokens.append(tokens[start_index:end_index]) return playlist_audio_tokens def _load_special_audio_prompt( self, item: QueueItem, model_version: str, is_cover: bool = False, is_artist: bool = False, is_underpainting: bool = False, is_overpainting: bool = False, ) -> np.ndarray: """Load the artist/cover prompt based on the QueueItem input. Note that this function loads one prompt a time. In case we need to load both, we need to call this twice. """ if is_artist: audio_id = item.metadata.get("artist_clip_id", "") from_s = item.metadata.get("artist_start_s", None) to_s = item.metadata.get("artist_end_s", None) elif is_cover: audio_id = item.metadata.get("cover_clip_id", "") from_s = item.metadata.get("cover_start_s", None) to_s = item.metadata.get("cover_end_s", None) elif is_underpainting: audio_id = item.metadata.get("underpainting_clip_id", "") from_s = item.metadata.get("underpainting_start_s", None) to_s = item.metadata.get("underpainting_end_s", None) elif is_overpainting: audio_id = item.metadata.get("overpainting_clip_id", "") from_s = item.metadata.get("overpainting_start_s", None) to_s = item.metadata.get("overpainting_end_s", None) else: raise ValueError("is_artist or is_cover must be True to load the audio prompt") if "." in audio_id: audio_id = audio_id.split(".")[0] with tempfile.NamedTemporaryFile(suffix=".npz") as temp_file: arch, major, _, _ = model_version.split(".") target_version = f"{arch}.{major}" tokens = _get_compatible_tokens(temp_file.name, audio_id, target_version) if tokens is None: raise ValueError(f"Couldn't fetch the tokens from {audio_id}") rate = VERSION_EMBEDDING_RATE_MAPPING.get(target_version, 25) start_index = int(from_s * rate) if from_s is not None else 0 end_index = int(to_s * rate) if to_s is not None else tokens.shape[0] assert end_index > start_index return tokens[start_index:end_index] def _load_infilling_audio_prompt( self, item: QueueItem, model_version: str, start_s: float | None = None, context_start_s: float | None = None, end_s: float | None = None, context_end_s: float | None = None, ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray] | None: """Return the tuple of 4 audio arrays.""" if not isinstance(item.prompt_audio, str): raise ValueError(f"{item.id} prompt audio is invalid.") if isinstance(item.prompt_audio, str) and "." in item.prompt_audio: audio_id, _ = item.prompt_audio.split(".") else: audio_id = item.prompt_audio with tempfile.NamedTemporaryFile(suffix=".npz") as f: arch, major, _, _ = model_version.split(".") target_version = f"{arch}.{major}" tokens = _get_compatible_tokens(f.name, audio_id, target_version) rate = VERSION_EMBEDDING_RATE_MAPPING.get(target_version, 25) if tokens is not None: context_start_index = int(context_start_s * rate) if context_start_s is not None else 0 start_index = int(start_s * rate) if start_s is not None else 0 end_index = int(end_s * rate) if end_s is not None else tokens.shape[0] context_end_index = ( int(context_end_s * rate) if context_end_s is not None else tokens.shape[0] ) assert context_start_index <= start_index <= end_index <= context_end_index return ( tokens[:context_start_index], tokens[context_start_index:start_index], tokens[end_index:context_end_index], tokens[context_end_index:], ) return None def _load_audio_prompt( self, item: QueueItem, model_version: str, from_s: Optional[float] = None, to_s: Optional[float] = None, ) -> np.ndarray | None: # TODO: if we really want to be backwards compatible with v1 # we will need to load mert and be able to encode audio mp3s # To be done another day continue_at = item.metadata.get("continue_at", None) if continue_at is not None: from_s = 0 to_s = continue_at if not isinstance(item.prompt_audio, str): raise ValueError(f"{item.id} prompt audio is invalid.") if "." in item.prompt_audio: audio_id, _ = item.prompt_audio.split(".") else: audio_id = item.prompt_audio with tempfile.NamedTemporaryFile(suffix=".npz") as f: arch, major, _, _ = model_version.split(".") tokens = get_tokens( f.name, audio_id, target_version=f"{arch}.{major}", from_s=from_s, to_s=to_s, ) if tokens is not None: return tokens else: print(f"Load hisotry prompt failed! {item.model_dump_json()}") def load_history_audio_as_audio(self, item: QueueItem) -> Audio | None: """Load the history audio as an Audio object.""" if not isinstance(item.prompt_audio, str): raise ValueError(f"{item.id} prompt audio is invalid.") if "." in item.prompt_audio: audio_id, _ = item.prompt_audio.split(".") else: audio_id, _ = item.prompt_audio, "mp3" with tempfile.NamedTemporaryFile(suffix=".webm") as f: try: s3_client.download_fileobj(DEFAULT_S3_BUCKET, f"studio/uploads/{audio_id}.webm", f) prompt_audio = Audio.from_file(f.name) return prompt_audio except Exception as e: print(f"No .webm file found for {audio_id}. Fallback to mp3. Error: {e}") with tempfile.NamedTemporaryFile(suffix=".mp3") as f: try: s3_client.download_fileobj(DEFAULT_S3_BUCKET, f"studio/uploads/{audio_id}.mp3", f) prompt_audio = Audio.from_file(f.name) return prompt_audio except botocore.exceptions.ClientError: print("No .mp3 file found for", audio_id) return None def _load_history_prompt(self, item: QueueItem): with tempfile.NamedTemporaryFile(suffix=".npz") as f: retry_s3_download(DEFAULT_S3_BUCKET, f"studio/uploads/{item.prompt_npz}.npz", f) history_prompt = np.load(f.name) return history_prompt def _load_history_vae(self, item: QueueItem, target_vae_version: str) -> np.ndarray | None: if not isinstance(item.prompt_audio, str): raise ValueError(f"{item.id} prompt audio is invalid {item.prompt_audio}.") if isinstance(item.prompt_audio, str) and "." in item.prompt_audio: audio_id, _ = item.prompt_audio.split(".") else: audio_id = item.prompt_audio print(f"Loading history VAE for {audio_id}, {item.prompt_audio}") with tempfile.NamedTemporaryFile(suffix=".npz") as f: fast_retry_s3_download(DEFAULT_S3_BUCKET, f"studio/uploads/{audio_id}_vae.npz", f) history_prompt = np.load(f.name) existing_vae_version = str(history_prompt["v_vae"].item()) if existing_vae_version == target_vae_version: return history_prompt["vae_latents"] else: if f"vae_latents_{target_vae_version}" in history_prompt: print( f"Found compatible vae version {target_vae_version} for {audio_id}, don't convert." ) return history_prompt[f"vae_latents_{target_vae_version}"] else: raise ValueError( f"Found incompatible vae version {existing_vae_version} for {audio_id}, need to convert..." ) def _decode_history_prompt(self, id: str): with tempfile.NamedTemporaryFile(suffix=".npz") as f: retry_s3_download(DEFAULT_S3_BUCKET, f"studio/uploads/{id}.npz", f) history_prompt = np.load(f.name) return history_prompt def _get_video_title(self, item: QueueItem): MAX_TITLE_LENGTH = 25 title = ( item.title or re.sub(r"\[.*?\]", "", item.prompt_text or "").strip().split("\n")[0] or "Untitled" ) return f"{title[:MAX_TITLE_LENGTH]}..." if len(title) > MAX_TITLE_LENGTH else title def _write_video( self, item: QueueItem, show_text=True, visualizer=False, image_path=None, audio=None, aligned_text=[], ): waveform_file = None if audio is None: get_audio_with_retry = retry_decorator(3, wait_seconds=10)(Audio.from_s3) audio = get_audio_with_retry( f"s3://suno-data-uploads/studio/uploads/{item.id}.mp3", n_channels=2, ) text = item.prompt_text if show_text and item.prompt_text else "" if audio is None: raise ValueError(f"{item.id} audio is invalid.") if image_path is None: raise ValueError(f"{item.id} image is invalid.") with tempfile.NamedTemporaryFile(suffix=".mp4") as waveform_tmp_file: if visualizer: tags = None if item.metadata.get("promotion") == "vday": tags = "vdaysong.com" else: tags = item.metadata.get("tags", None) use_hq_settings = item.metadata.get("is_high_quality", False) # Use new video pipeline for promotional videos promotion = item.metadata.get("promotion", None) waveform_file = generate_video( audio, image_path, aligned_text, waveform_tmp_file.name, song_title=self._get_video_title(item), song_user_handle=item.metadata.get("user_handle"), tags=tags, resolution_factor=2 if use_hq_settings else 1, fps=10 if use_hq_settings else 5, video_mode=promotion, ) retry_s3_upload( waveform_file, DEFAULT_S3_BUCKET, f"studio/uploads/{item.id}.mp4", ExtraArgs={ "ContentType": "video/mp4", }, ) async def _write_audio_only_async( self, item: QueueItem, audio: Audio, s3_bucket: str = DEFAULT_S3_BUCKET, s3_folder: str = DEFAULT_S3_FOLDER, gain_adjust: float = 0, write_wav: bool = False, ): import asyncio applied_gain = 0 if gain_adjust != 0: audio, applied_gain = await asyncio.get_running_loop().run_in_executor( None, audio.apply_gain, gain_adjust ) gain_adjust_meta = [] if applied_gain != 0: gain_adjust_meta = ["-metadata", f"sga={str(float(applied_gain))}"] async def _upload_result( task: asyncio.Task, local_path: str, bucket: str, remote_path: str, extra_args: dict[str, str], add_ttl: bool = False, ): await task def _upload_with_extra_args(): with open(local_path, "rb") as f: if add_ttl: s3_client.put_object( Bucket=bucket, Key=remote_path, Body=f, Tagging="wav-ttl=one-day&IntelligentTiering=true", **extra_args, ) else: s3_client.put_object( Bucket=bucket, Key=remote_path, Body=f, **extra_args, ) await asyncio.get_running_loop().run_in_executor( None, _upload_with_extra_args, ) with tempfile.TemporaryDirectory() as temp_dir: opus_path = os.path.join(temp_dir, "opus.opus") mp3_path = os.path.join(temp_dir, "mp3.mp3") webm_path = os.path.join(temp_dir, "webm.webm") wav_path = os.path.join(temp_dir, "wav.wav") t_opus = asyncio.create_task(audio.write_opus_async(opus_path)) async def write_opus_and_webm(): await t_opus proc = await asyncio.create_subprocess_exec( "ffmpeg", "-hide_banner", "-loglevel", "error", "-y", "-i", opus_path, "-c:a", "copy", *gain_adjust_meta, "-f", "webm", webm_path, stdin=asyncio.subprocess.DEVNULL, stdout=None, stderr=None, ) await proc.wait() if proc.returncode != 0: raise RuntimeError(f"ffmpeg recontainer exited with code {proc.returncode}") t_webm = asyncio.create_task(write_opus_and_webm()) t_mp3 = asyncio.create_task(audio.write_hq_mp3_async(mp3_path)) if write_wav: t_wav = asyncio.create_task(audio.write_wav_async(wav_path)) tasks = [ (t_opus, opus_path, "opus", "audio/ogg"), (t_mp3, mp3_path, "mp3", "audio/mp3"), (t_webm, webm_path, "webm", "audio/webm; codecs=opus"), ] if write_wav: tasks.append((t_wav, wav_path, "wav", "audio/wav")) upload_tasks = [ _upload_result( task, local_path, s3_bucket, os.path.join(s3_folder, f"{item.id}.{ext}"), {"ContentType": content_type}, add_ttl=ext == "wav", ) for task, local_path, ext, content_type in tasks ] await asyncio.gather(*upload_tasks) def _write_audio_only( self, item: QueueItem, audio: Audio, s3_bucket: str = DEFAULT_S3_BUCKET, s3_folder: str = DEFAULT_S3_FOLDER, gain_adjust: float = 0, ): def process_audio_format( audio, item_id, s3_bucket, s3_folder, format_info, gain_adjust: float = 0 ): """Process a single audio format and upload to S3""" format_name, file_ext, content_type, write_method = format_info print(f"Writing {format_name} audio for {item_id}") with tempfile.NamedTemporaryFile(suffix=f".{file_ext}") as temp_file: # Call the appropriate write method on the audio object getattr(audio, write_method)(temp_file.name, gain_adjust=gain_adjust) path = os.path.join(s3_folder, f"{item_id}.{file_ext}") print(f"Writing {format_name} to {s3_bucket} {path}") retry_s3_upload( temp_file.name, s3_bucket, path, ExtraArgs={"ContentType": content_type}, ) return f"Completed {format_name} for {item_id}" # Define the formats with their properties: (name, file extension, content type, audio write method) formats = [ ("mp3", "mp3", "audio/mp3", "write_hq_mp3"), ("opus", "opus", "audio/ogg", "write_opus"), ("webm", "webm", "audio/webm; codecs=opus", "write_webm_opus"), ] # Create a partial function with the common arguments process_fn = partial( process_audio_format, audio, item.id, s3_bucket, s3_folder, gain_adjust=gain_adjust ) # Run all format conversions and uploads in parallel with ThreadPoolExecutor(max_workers=3) as executor: # Submit all tasks to the executor future_to_format = { executor.submit(process_fn, format_info): format_info[0] for format_info in formats } # Process results as they complete for future in as_completed(future_to_format): format_name = future_to_format[future] try: result = future.result() print(f"Result: {result}") except Exception as e: print(f"Error processing {format_name} format: {e}") def _write_audio( self, item: QueueItem, audio: Audio, prompt_arrs: list[Any], show_text=True, make_gif=False, # for backwards compatibility gif_prompt="dancing dog", # for backwards compatibility visualizer=False, s3_bucket: str = DEFAULT_S3_BUCKET, s3_folder: str = DEFAULT_S3_FOLDER, ): self._write_audio_only(item, audio) # audio = audio.convert(48000, 2, 2) self._write_video( item, show_text=show_text, visualizer=visualizer, audio=audio, ) if prompt_arrs: with tempfile.NamedTemporaryFile(suffix=".npz") as temp_file: full_generation = { "semantic_prompt": prompt_arrs[0], "coarse_prompt": prompt_arrs[1], "fine_prompt": prompt_arrs[2], } np.savez(temp_file.name, **full_generation) retry_s3_upload(temp_file.name, s3_bucket, os.path.join(s3_folder, f"{item.id}.npz")) def _write_npz( self, item: QueueItem, audio_arrs: np.ndarray, inference_runner: str, model_version: str, generated_arr: np.ndarray | None = None, history_arr: np.ndarray | None = None, future_arr: np.ndarray | None = None, pre_history_arr: np.ndarray | None = None, post_future_arr: np.ndarray | None = None, cover_arr: np.ndarray | None = None, artist_arr: np.ndarray | None = None, playlist_arr: List[np.ndarray] | None = None, multi_artist_arr: List[np.ndarray] | None = None, overpainting_arr: np.ndarray | None = None, underpainting_arr: np.ndarray | None = None, history_lyrics: str | None = None, fuller_array: np.ndarray | None = None, n_skip_semantic: int = 1, s3_bucket: str = DEFAULT_S3_BUCKET, s3_folder: str = DEFAULT_S3_FOLDER, ): """ Write the npz file to S3. audio_arrs: the audio's npz arrays. Will match the codec decoded output. inference_runner: str -- the modal worker name model_version: str -- used for codec decoding generated_arr: generated codes. history_arr: history codes. future_arr: future codes -- infilling. pre_history_arr: pre history codes -- infilling. post_future_arr: post future codes -- infilling. cover_arr: cover codes -- cover. artist_arr: artist codes -- artist. history_lyrics: history lyrics -- for prompting extend. For different tasks: - cover: cover_arr contains the cover codes for the generation. - need to concat with audio_arrs to get input for DPO. - lyrics should be consistent with metadata's lyrics. - artist: artist_arr contains the artist codes for the generation. - need to concat with audio_arrs to get input for DPO. - lyrics should be consistent with metadata's lyrics. - infilling: full_arr contains the infilling codes for the generation. - do NOT need to concat with audio_arrs. - lyrics should be consistent with metadata's lyrics. - extend: history_arr contains the history codes for the generation. - full array is history + generated -- can be used directly. - history lyrics is hoooted -- so we will save it here. """ full_arr = None # full arr will be the full audio, end to end history_start_index = None generated_start_index = None future_start_index = None post_future_start_index = None # we need to keep the generated dimesions consistent cause auk is sem only # but the history loader actually loads the full code + coarse... n_gen_dim = audio_arrs.shape[1] if history_arr is not None: history_arr = history_arr[:, :n_gen_dim] full_arr = generated_arr if pre_history_arr is not None: pre_history_arr = pre_history_arr[:, :n_gen_dim] history_start_index = pre_history_arr.shape[0] generated_start_index = pre_history_arr.shape[0] + history_arr.shape[0] full_arr = np.concatenate([pre_history_arr, history_arr, full_arr], axis=0) else: history_start_index = 0 generated_start_index = history_arr.shape[0] full_arr = np.concatenate([history_arr, full_arr], axis=0) if future_arr is not None: future_arr = future_arr[:, :n_gen_dim] if full_arr is None: full_arr = generated_arr assert full_arr is not None future_start_index = full_arr.shape[0] full_arr = np.concatenate([full_arr, future_arr], axis=0) if post_future_arr is not None: post_future_arr = post_future_arr[:, :n_gen_dim] post_future_start_index = full_arr.shape[0] full_arr = np.concatenate([full_arr, post_future_arr], axis=0) with tempfile.NamedTemporaryFile(suffix=".npz") as temp_file: arch, major, _, _ = model_version.split(".") full_generation = { f"v{arch}.{major}_raw": audio_arrs, "inference_runner": inference_runner, "model_version": model_version, } if full_arr is not None: full_arr = full_arr.astype(np.int32) full_generation["full_arr"] = full_arr # save the context, infilling start indices. # sufficient info for slicing out the full arr if history_start_index is not None: full_generation["history_start_index"] = history_start_index if generated_start_index is not None: full_generation["generated_start_index"] = generated_start_index if future_start_index is not None: full_generation["future_start_index"] = future_start_index if post_future_start_index is not None: full_generation["post_future_start_index"] = post_future_start_index if history_arr is not None: history_arr = history_arr.astype(np.int32) full_generation["history_arr"] = history_arr history_text = str(history_lyrics) if history_lyrics else "" # wrap it as np array to save full_generation["history_text"] = np.array(history_text) if cover_arr is not None: cover_arr = cover_arr.astype(np.int32) full_generation["cover_arr"] = cover_arr if artist_arr is not None: artist_arr = artist_arr.astype(np.int32) full_generation["artist_arr"] = artist_arr if playlist_arr is not None: # Convert each sub-array to int32 and store as object array to handle inhomogeneous shapes # Concatenate the list of arrays into one, and keep track of the original lengths playlist_arr = [ np.asarray(sub_playlist_arr, dtype=np.int32) for sub_playlist_arr in playlist_arr ] playlist_lengths = [arr.shape[0] for arr in playlist_arr] playlist_concat = np.concatenate(playlist_arr, axis=0) full_generation["playlist_arr"] = playlist_concat full_generation["playlist_arr_len"] = np.array(playlist_lengths, dtype=np.int32) if multi_artist_arr is not None: # Convert each sub-array to int32 and store as object array to handle inhomogeneous shapes multi_artist_arr = [ np.asarray(sub_multi_artist_arr, dtype=np.int32) for sub_multi_artist_arr in multi_artist_arr ] multi_artist_lengths = [arr.shape[0] for arr in multi_artist_arr] multi_artist_concat = np.concatenate(multi_artist_arr, axis=0) full_generation["multi_artist_arr"] = multi_artist_concat full_generation["multi_artist_arr_len"] = np.array(multi_artist_lengths, dtype=np.int32) if overpainting_arr is not None: overpainting_arr = overpainting_arr.astype(np.int32) full_generation["overpainting_arr"] = overpainting_arr if underpainting_arr is not None: underpainting_arr = underpainting_arr.astype(np.int32) full_generation["underpainting_arr"] = underpainting_arr if fuller_array is not None: fuller_array = fuller_array.astype(np.int32) full_generation["fuller_arr"] = fuller_array full_generation["n_skip_semantic"] = n_skip_semantic np.savez(temp_file.name, **full_generation) retry_s3_upload(temp_file.name, s3_bucket, os.path.join(s3_folder, f"{item.id}.npz")) def _write_vae_latents_npz( self, item: QueueItem, vae_latents: np.ndarray, vae_version: str, inference_runner: str, parent_from_start_index: int | None = None, parent_from_end_index: int | None = None, parent_clip_id: str | None = None, history_latents: np.ndarray | None = None, history_text: str | None = None, n_sem_tokens: int | None = None, seed: int | None = None, future_latents: np.ndarray | None = None, semantic_tokens: np.ndarray | None = None, ): """Save the vae latents to S3. Note the file name is hardcoded to include "_vae". They will be saved NOT compressed. - compression takes about 3x to even 10x longer. - not significant savings in storage. item: QueueItem vae_latents: np.ndarray -- the vae latents to save vae_version: str -- the vae version inference_runner: str -- the model inference worker name parent_from_start_index: int -- the start index of the parent clip in the latent space array This index excludes the pad 5 sections. parent_from_end_index: int -- the end index of the parent clip in the latent space array This index excludes the pad 5 sections. ************* THIS IS VERY FUCKED UP right now for v4 VAE latents are not expected to agree in length with the semantic codes. Mostly due to the end CROP_CHUNK_END In normal gens, VAE latents will be 5 tokens shorter than the semantic codes. This is known but hard to fix TODO: need VICTOR to fix this at some point. ************* parent_clip_id: str -- the parent clip id -- for look up retrieval. history_latents: np.ndarray | None -- the history latents to save history_text: str | None -- the history text to save n_sem_tokens: int | None -- the number of semantic tokens processed seed: int | None -- the seed used for the generation future_latents: np.ndarray | None -- the future latents to save semantic_tokens: np.ndarray | None -- the semantic tokens prompted as input """ with tempfile.NamedTemporaryFile(suffix=".npz") as temp_file: full_generation = { "v_vae": vae_version, "vae_latents": vae_latents.astype(np.float16), "inference_runner": inference_runner, } if parent_from_start_index is not None: full_generation["parent_from_start_index"] = parent_from_start_index if parent_from_end_index is not None: full_generation["parent_from_end_index"] = parent_from_end_index if parent_clip_id is not None: full_generation["parent_clip_id"] = parent_clip_id.replace(".mp3", "") if history_latents is not None: full_generation["history_arr"] = history_latents.astype(np.float16) if future_latents is not None: full_generation["future_arr"] = future_latents.astype(np.float16) if history_text is not None: # wrap it as np array to save full_generation["history_text"] = np.array(history_text) if n_sem_tokens is not None: full_generation["n_sem_tokens"] = n_sem_tokens if semantic_tokens is not None: full_generation["semantic_tokens"] = semantic_tokens.astype(np.int32) if seed is not None: full_generation["seed"] = seed np.savez(temp_file.name, **full_generation) retry_s3_upload(temp_file.name, DEFAULT_S3_BUCKET, f"studio/uploads/{item.id}_vae.npz") def copy_and_upload_image(self, item: QueueItem): """This function is used to duplicate the image and upload to s3 for the new id.""" # download the images image_s3_id = item.metadata.get( "image_s3_id", "image_" + item.prompt_audio if item.prompt_audio else "image_large_default_bird", ) if image_s3_id is None: image_s3_id = "image_large_default_bird" # add a protection if image_s3_id and "." in image_s3_id: image_s3_id, _ = image_s3_id.split(".") # we want to fetch the large image, if they are generated, not user uploads if image_s3_id.startswith("image_") and not image_s3_id.startswith("image_large_"): image_s3_id = image_s3_id.replace("image_", "image_large_") image_s3_path = f"studio/uploads/{image_s3_id}.jpeg" # check if the image exists. try: s3_client.get_object(Bucket=DEFAULT_S3_BUCKET, Key=image_s3_path) except Exception as e: # fall back to small jpeg first image_s3_id = image_s3_id.replace("image_large_", "image_") image_s3_path = f"studio/uploads/{image_s3_id}.jpeg" try: s3_client.get_object(Bucket=DEFAULT_S3_BUCKET, Key=image_s3_path) except Exception as e: # fall back to png print(f"Can't get object from {image_s3_path}, falling back to png. Error: {e}") image_s3_path = image_s3_path.replace(".jpeg", ".png") try: s3_client.get_object(Bucket=DEFAULT_S3_BUCKET, Key=image_s3_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" image_s3_path = f"studio/uploads/{FIXED_IMAGE_FILE}" with tempfile.TemporaryDirectory() as td: image_filename = os.path.join(td, "image_small.jpeg") image_large_filename = os.path.join(td, "image_large.jpeg") with open(image_large_filename, "wb") as tmp_file: retry_s3_download(DEFAULT_S3_BUCKET, image_s3_path, tmp_file) # load the pil image, and do the conversion pil_image = Image.open(image_large_filename) # JPEG doesn't support transpanrency... pil_image = pil_image.convert("RGB") if image_s3_path.endswith(".png"): # overwrite the file with jpeg pil_image.save(image_large_filename, quality=75) # Using image.thumbnail to conserve aspect ratio thumbnail_image = ImageOps.exif_transpose(pil_image.copy()) thumbnail_width, thumbnail_height = 360, 360 width, height = thumbnail_image.size # if not square, get thumbnail max dimensions where shorter side is 360px if width != height: thumbnail_width, thumbnail_height = self._get_thumbnail_dimensions(width, height) thumbnail_image.thumbnail((thumbnail_width, thumbnail_height), Image.Resampling.LANCZOS) thumbnail_image.save(image_filename, quality=50) # copy the images s3_client.upload_file( image_filename, DEFAULT_S3_BUCKET, f"studio/uploads/image_{item.id}.jpeg" ) s3_client.upload_file( image_large_filename, DEFAULT_S3_BUCKET, f"studio/uploads/image_large_{item.id}.jpeg" ) # add a notification to the item item.notify_progress( { "id": item.id, "type": "image", "image_id": f"image_{item.id}", }, ) def _get_thumbnail_dimensions( self, width, height, max_shorter_dimension=360, max_longer_dimension=1080 ): if width < height: new_width = max_shorter_dimension new_height = max_longer_dimension else: new_height = max_shorter_dimension new_width = max_longer_dimension return new_width, new_height