"""Audio upsample (diffusion) application on modal. Code --> Diffusion --> Audio """ import json import modal import tempfile import time import os import torch import fcntl import select import numpy as np from uuid import uuid4 from typing import Optional from datadog import initialize, statsd import random from suno_utils.worker.loader import S3Loader from suno_utils.worker.schema import QueueItem, HistoryPrompt from suno_utils.diffusion import generation as diffusion_gen from suno_utils.audio import Audio from suno_utils.worker.settings import s3_client from suno_utils.worker.utils import ffmpeg_stream_encode_opus_webm, retry_decorator, ffmpeg_stream_encode from suno_utils.worker.modal_model_configs import MODEL_CONFIG_DICT, VAEVersion from suno_utils.tasks.upsample_engine_old import UpsampleEngine, Request from suno_utils.worker.tracing import distributed_trace, tracer from suno_utils.worker.modal_base import get_modal_base_image_diffusion_with_flash_attention from suno_utils.worker.modal_model_volume import MODEL_STORE_VOLUME_DIR, model_store_volume ############## CHANGE THESE ############## DEPLOYMENT_TYPE = "dev" # dev, prod MODEL_CONFIG = MODEL_CONFIG_DICT["diff_v2"] # Production models: # diff_v1 # diff_v2, diff_v2_test, diff_v2_data # Dev model names: # diff_v1, diff_v1_test, diff_v1_test_2 # diff_v2, diff_v2_test, # diff_seeds_v0 MODEL_VAE_VERSION = MODEL_CONFIG.vae_version MODEL_GAIN_ADJUST = MODEL_CONFIG.gain_adjust MODEL = MODEL_CONFIG.model # Maximum allowed semantic mask ratio (0.8 = 80% masking) MAX_SEMANTIC_MASK_RATIO = 0.95 # do some hacky stuff to adjust gain if MODEL == "diff_v2_data": # diff v2 data collection model is quieter MODEL_GAIN_ADJUST += 1.8 if MODEL == "diff_v2_test": MODEL_GAIN_ADJUST = 0.43 print(f"Model gain adjust for {MODEL} to: {MODEL_GAIN_ADJUST} dB") ########################################## fast_retry_s3_download = retry_decorator(3, wait_seconds=1)(s3_client.download_fileobj) MOUNT_PATH = "/suno/models" DIFFUSION_CKPT_PATH = MODEL_CONFIG.diffusion_ckpt_path # print(f"Upsample model: {DIFFUSION_CKPT_PATH}") aws_secret = modal.Secret.from_name("studio-aws") SECRETS = [ aws_secret, modal.Secret.from_name("api-callback-token"), modal.Secret.from_name("datadog-metrics"), ] # These queue names must correspond to the ones in streaming_api ENV_CONFIGS CHUNK_QUEUE_NAME = f"chunk-queue-{DEPLOYMENT_TYPE}" WEBM_CHUNK_QUEUE_NAME = f"chunk-queue-webm-{DEPLOYMENT_TYPE}" STREAM_KEY_QUEUE_NAME = f"stream-key-queue-{DEPLOYMENT_TYPE}" UPSAMPLE_DDOG_SERVICE = "upsample-worker" mp3_audio_chunk_queue = modal.Queue.from_name(CHUNK_QUEUE_NAME, create_if_missing=True) webm_audio_chunk_queue = modal.Queue.from_name(WEBM_CHUNK_QUEUE_NAME, create_if_missing=True) # Queue used by streaming_api to determine which ephemeral stream_key to stream from. # This corresponds to stream_key_queue_name in streaming_api. appstream_key_queue = modal.Queue.from_name(STREAM_KEY_QUEUE_NAME, create_if_missing=True) # this uses mert... base_image = get_modal_base_image_diffusion_with_flash_attention().pip_install("transformers==4.44.0") class UpsampleWorker(S3Loader, UpsampleEngine): def __init__(self): S3Loader.__init__(self) # configure the upsample's min chunk size here UpsampleEngine.__init__(self, min_chunk_size=25 * 30, vae_version=MODEL_VAE_VERSION) print(f"Model gain adjust: {MODEL_GAIN_ADJUST} dB") start_time = time.time() print("Start loading models") tokenizer_path = diffusion_gen.get_model_if_needed( diffusion_gen.TOKENIZER_FILEPATH, cache_dir=MOUNT_PATH ) semantic_model_path = diffusion_gen.get_model_if_needed( diffusion_gen.SEMANTIC_MODEL_FILEPATH, cache_dir=MOUNT_PATH ) semantic_clusters_path = diffusion_gen.get_model_if_needed( diffusion_gen.SEMANTIC_CLUSTERS_FILEPATH, cache_dir=MOUNT_PATH ) codec_path = diffusion_gen.get_model_if_needed( MODEL_CONFIG.codec_ckpt_path, cache_dir=MOUNT_PATH ) dit_model_path = diffusion_gen.get_model_if_needed(DIFFUSION_CKPT_PATH, cache_dir=MOUNT_PATH) diffusion_gen.preload_models( tokenizer_filepath=tokenizer_path, semantic_model_filepath=semantic_model_path, semantic_clusters_filepath=semantic_clusters_path, codec_filepath=codec_path, dit_model_filepath=dit_model_path, compile=True, # can't compile with flash v2 ) print(f"Finish loading models. Took {round(time.time() - start_time, 2)} seconds") def download_models_to_modal_disk(dir_path=MOUNT_PATH): """Download diffusion models.""" print("Start downloading models.") _ = diffusion_gen.get_model_if_needed(diffusion_gen.TOKENIZER_FILEPATH, cache_dir=dir_path) _ = diffusion_gen.get_model_if_needed(diffusion_gen.SEMANTIC_MODEL_FILEPATH, cache_dir=dir_path) _ = diffusion_gen.get_model_if_needed(diffusion_gen.SEMANTIC_CLUSTERS_FILEPATH, cache_dir=dir_path) _ = diffusion_gen.get_model_if_needed(diffusion_gen.CODEC_FILEPATH, cache_dir=dir_path) _ = diffusion_gen.get_model_if_needed(DIFFUSION_CKPT_PATH, cache_dir=dir_path) print("Finish downloading models.") image = base_image.add_local_python_source("suno_utils", copy=False) APP_NAME = f"upsample-{MODEL_CONFIG.model}-{DEPLOYMENT_TYPE}" app = modal.App(APP_NAME, image=image) @app.cls( gpu=["H100"], # H100 is faster, but if limited by resources, burst into A100 cpu=4, secrets=SECRETS, timeout=240, # the wav encoding can take a while scaledown_window=240, memory=15000, retries=modal.Retries( max_retries=1, backoff_coefficient=2.0, initial_delay=5.0, ), min_containers=1 if DEPLOYMENT_TYPE == "dev" else 2, max_containers=4 if DEPLOYMENT_TYPE == "dev" else 500, buffer_containers=0 if DEPLOYMENT_TYPE == "dev" else 1, volumes={MODEL_STORE_VOLUME_DIR: model_store_volume}, ) @modal.concurrent(max_inputs=4) class UpsampleStub: def __init__(self): import torch print(f"Init with diffusion ckpt: {DIFFUSION_CKPT_PATH}") num_gpus = torch.cuda.device_count() device_name = torch.cuda.get_device_name() print(f"Found {num_gpus} GPUs. Using CUDA device: {device_name}") self.worker = UpsampleWorker() self.modal_f_video_generation = modal.Cls.lookup( f"videos-v2-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}", "DummyV0Stub", )().write_video self.modal_f_cycle = modal.Cls.lookup( f"cycle-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}", "CycleStub", )().encode_audio random_seed = int((time.time() * 1000000) % 100000000) print("Random seed set to:", random_seed) random.seed(random_seed) self.seed_pool = [random.randint(0, 1_000_000) for _ in range(100)] self.worker.start() self.worker.wait_for_warmup() print("Upsample engine warmed up") options = {"statsd_host": "127.0.0.1", "statsd_port": 8125} initialize(**options) @modal.method() @distributed_trace("gpt_generate", UPSAMPLE_DDOG_SERVICE, env_name=DEPLOYMENT_TYPE) def upsample(self, queue_item: str, history: Optional[HistoryPrompt] = None) -> None: """Takes queue item and upsample it acoordingly.""" tracer.current_span().set_tag("deployment_type", DEPLOYMENT_TYPE) item = QueueItem(**json.loads(queue_item)) device_name = torch.cuda.get_device_name() dd_tags = [ f"model:{MODEL}", f"env:{DEPLOYMENT_TYPE}", f"env_name:{APP_NAME}", f"modal_cloud_provider:{os.environ.get('MODAL_CLOUD_PROVIDER', 'unknown')}", f"modal_environment:{os.environ.get('MODAL_ENVIRONMENT', 'unknown')}", f"modal_image_id:{os.environ.get('MODAL_IMAGE_ID', 'unknown')}", f"modal_region:{os.environ.get('MODAL_REGION', 'unknown')}", f"gen_type:{item.metadata.get('type', 'unknown')}", f"gen_task:{item.metadata.get('task', 'unknown')}", # this is not always set for stuff like basic gens f"cuda_device:{device_name}", ] start_time = time.time() # Each unique Chirp request is assigned a unique ephemeral stream_key. # Replace item in the stream_key_queue. # streaming_api streams from the first item in the stream_key_queue. stream_key = "stream-key-" + str(uuid4()) appstream_key_queue.get(block=False, partition=item.id) appstream_key_queue.put(stream_key, partition=item.id, partition_ttl=600) # If specified, use stream_key to read tokens from the token queue. # Otherwise, fall back to the item id. partition = stream_key or item.id print(f"Got {item}") tags = item.metadata.get("tags", "") or "" # if None set to empty string lyrics = item.metadata.get("prompt", "") or "" # if None set to empty string # semantic_skip_factor = item.metadata.get("diff_freedom", 1) or 1 # default to 1 semantic_mask_ratio = item.metadata.get("diff_freedom", 0) or 0 # default to 0 # Ensure input is within 0-MAX_SEMANTIC_MASK_RATIO range semantic_mask_ratio = max(0, min(MAX_SEMANTIC_MASK_RATIO, semantic_mask_ratio)) if MODEL == "diff_v1": text_cfg_coef = 2.0 steps = 10 elif MODEL == "diff_v1_test_2" or MODEL == "diff_v1_test_3" or MODEL == "diff_v1_test_4": text_cfg_coef = 2.0 steps = 10 elif MODEL == "diff_v1_test": text_cfg_coef = 2.0 steps = 10 elif MODEL == "diff_v2" or MODEL == "diff_v2_test" or MODEL == "diff_v2_data": text_cfg_coef = 2.0 steps = 10 elif MODEL == "diff_seeds_v0": text_cfg_coef = 4.0 steps = 16 else: raise ValueError(f"Unknown model: {MODEL}") print( f"{item.id} has diffusion text cfg " f"{text_cfg_coef}, context cfg " f"{steps} steps, semantic mask ratio " f"{semantic_mask_ratio}." ) chosen_seed = random.choice(self.seed_pool) diffusion_inference_config = dict( text_cfg_coef=text_cfg_coef, steps=steps, semantic_mask_ratio=semantic_mask_ratio, seed=chosen_seed, ) # v2 uses a different scale factor and downscales the ctx vector if MODEL_VAE_VERSION == VAEVersion.V_VAE_25_TUNED_2: # use a different scale factor adn downscale the ctx vector diffusion_inference_config["codec_scale_factor"] = 0.4 diffusion_inference_config["scale_ctx_vector"] = True diffusion_inference_config["noise_ctx_level"] = 0.75 diffusion_inference_config["noise_ctx_pad_len"] = 0 # if MODEL == "diff_v2_test": # # do it at more conservative level for v2 data collection # diffusion_inference_config["noise_ctx_level"] = 0.0 if item.is_diff_infill: assert history is not None assert history.history_latents is not None assert history.future_latents is not None assert history.prompt_audio is not None print( f"Upsample {item.id}: " f"history_latents shape: {history.history_latents.shape}, " f"future_latents shape: {history.future_latents.shape}, " f"infill semantics shape: {history.prompt_audio.shape}" ) if history.future_latents.shape[0] > 750: raise ValueError( f"Infill future latents are too long: {history.future_latents.shape[0]}" ) if history.history_latents.shape[0] > 750: raise ValueError( f"Infill history latents are too long: {history.history_latents.shape[0]}" ) diffusion_inference_config["noise_ctx_level"] = 0.0 diffusion_inference_config["noise_ctx_pad_len"] = 0 diffusion_inference_config["infill_prefix_latents"] = torch.from_numpy( history.history_latents ) diffusion_inference_config["infill_suffix_latents"] = torch.from_numpy( history.future_latents ) diffusion_inference_config["drop_semantic_tokens"] = False diffusion_inference_config["text_cfg_coef"] = 3.0 diffusion_inference_config["semantic_mask_ratio"] = 0.65 # if diffusion_inference_config["semantic_skip_factor"] > 1: # diffusion_inference_config["drop_semantic_tokens"] = False # Note that experiment config will overwrite default config if special_config := item.metadata.get("model_config"): # has specific config instruction print(f"Upsample {item.id}: load specific config: {special_config}.") # update the custom config # note this will overwrite existing parameters diffusion_inference_config.update(**special_config) # Note forced infer config will overwrite experiment config if special_config := item.metadata.get("forced_infer_config"): # has specific config instruction print(f"Upsample {item.id}: load forced infer config: {special_config}") # update the custom config # note this will overwrite existing parameters diffusion_inference_config.update( **{k: v for k, v in special_config.items() if v is not None} ) print(f"Upsample {item.id}: diffusion_inference_config: {diffusion_inference_config}") # Filter kwargs that are valid properties of GenerationConfig valid_gconf_props = set(vars(diffusion_gen.DiffusionGenerationConfig()).keys()) invalid_keys = set(diffusion_inference_config.keys()) - valid_gconf_props if invalid_keys: print(f"WARNING: Stem {item.id}: Invalid config keys found: {invalid_keys}") diffusion_inference_config = { k: v for k, v in diffusion_inference_config.items() if k in valid_gconf_props } with tempfile.TemporaryDirectory() as td: if MODEL == "diff_seeds_v0": history_audio = torch.ones(750).long() * 4000 audio_semantic_codes = history_audio.reshape(1, -1).long() else: # download the audio file if not isinstance(item.prompt_audio, str): raise ValueError(f"{item.id} Prompt audio is not a string: {item.prompt_audio}") history_audio = ( history.prompt_audio if history else self.worker._load_audio_prompt(item, "5.0.0.0") ) if history_audio is None: raise ValueError(f"Failed to load audio prompt for {item.id}") audio_semantic_codes = torch.from_numpy(history_audio[:, 0]).reshape(1, -1).long() expected_duration = round(history_audio.shape[0] / 25, 2) print( f"{item.id} got parent audio. Expected raution {expected_duration}s, {(time.time() - start_time):.2f}s elapsed." ) if expected_duration > 16 * 60: # 20 MB modal queue size raise ValueError(f"{item.id} Parent audio is too long: {expected_duration}s") audios: list[Audio] = [] with ( ffmpeg_stream_encode(gain_adjust=MODEL_GAIN_ADJUST) as mp3_proc, ffmpeg_stream_encode_opus_webm(gain_adjust=MODEL_GAIN_ADJUST) as webm_proc, ): # Increase pipe buffer size (e.g., to 1 MB) # Should fix broken pipe error PIPE_BUF_SIZE = 1024 * 1024 # 1 MB fcntl.fcntl(mp3_proc.stdin.fileno(), fcntl.F_SETPIPE_SZ, PIPE_BUF_SIZE) fcntl.fcntl(mp3_proc.stdout.fileno(), fcntl.F_SETPIPE_SZ, PIPE_BUF_SIZE) fcntl.fcntl(webm_proc.stdin.fileno(), fcntl.F_SETPIPE_SZ, PIPE_BUF_SIZE) fcntl.fcntl(webm_proc.stdout.fileno(), fcntl.F_SETPIPE_SZ, PIPE_BUF_SIZE) os.set_blocking(mp3_proc.stdout.fileno(), False) os.set_blocking(webm_proc.stdout.fileno(), False) os.set_blocking(mp3_proc.stdin.fileno(), False) os.set_blocking(webm_proc.stdin.fileno(), False) mp3_poll = select.poll() mp3_poll.register(mp3_proc.stdout, select.POLLIN) webm_poll = select.poll() webm_poll.register(webm_proc.stdout, select.POLLIN) diffusion_generation_config = diffusion_gen.DiffusionGenerationConfig( lyrics=lyrics or "", tags=tags or "", **diffusion_inference_config, ) print( f"Upsample {item.id} adding request: audio_semantic_codes shape: {audio_semantic_codes.shape}" ) job = self.worker.add_request( Request( item.id, diffusion_generation_config, [c for c in audio_semantic_codes[0]], input_tokens_finished=True, ) ) print(f"Upsample {item.id}: added job") mp3_first_log_time = None webm_first_log_time = None sent_notify = False def _poll_reads(): nonlocal mp3_first_log_time, webm_first_log_time, sent_notify mp3_done = False webm_done = False made_progress = False # read audio from ffmpeg if webm_poll.poll(1): webm_read_res = webm_proc.stdout.read(64 * 1024) if webm_read_res: webm_audio_chunk_queue.put( webm_read_res, partition=partition, partition_ttl=600, block=False, timeout=600, ) made_progress = True else: webm_done = True if webm_first_log_time is None: webm_first_log_time = time.time() elapsed_time = webm_first_log_time - start_time print( f"First audio chunk writing: {elapsed_time:.2f}s elapsed since request received.", ) first_audio_write_time_millis = int(time.time() * 1000) elapsed_millis = first_audio_write_time_millis - int(start_time * 1000) if elapsed_millis > 0: print( f"UpsampleStub ({item.id}): Start audio write time from gen request: {elapsed_millis // 1000}s" ) if elapsed_millis > 600_000: print( f"UpsampleStub ({item.id}): Start audio write time from gen request took >10m: {elapsed_millis // 1000}s" ) statsd.distribution( "decoder.audio_write_start_millis.distribution", elapsed_millis, tags=dd_tags + ["format:webm"], ) if mp3_poll.poll(1): mp3_read_res = mp3_proc.stdout.read(64 * 1024) if mp3_read_res: mp3_audio_chunk_queue.put( mp3_read_res, partition=partition, partition_ttl=600, block=False, timeout=600, ) made_progress = True else: mp3_done = True if mp3_first_log_time is None: mp3_first_log_time = time.time() elapsed_time = mp3_first_log_time - start_time print( f"First audio chunk writing: {elapsed_time:.2f}s elapsed since request received.", ) first_audio_write_time_millis = int(time.time() * 1000) elapsed_millis = first_audio_write_time_millis - int(start_time * 1000) if elapsed_millis > 0: print( f"UpsampleStub ({item.id}): Start audio write time from gen request: {elapsed_millis // 1000}s" ) if elapsed_millis > 600_000: print( f"UpsampleStub ({item.id}): Start audio write time from gen request took >10m: {elapsed_millis // 1000}s" ) statsd.distribution( "decoder.audio_write_start_millis.distribution", elapsed_millis, tags=dd_tags + ["format:mp3"], ) if ( mp3_first_log_time is not None or webm_first_log_time is not None ) and not sent_notify: item.notify_progress( { "id": item.id, "type": "streaming", }, ) sent_notify = True return (mp3_done and webm_done), made_progress for audio in job.audio_generator(): audios.append(audio) partially_written_chunk_mp3 = np.ascontiguousarray(audio.array_float.T).tobytes() partially_written_chunk_webm = partially_written_chunk_mp3 while len(partially_written_chunk_mp3) > 0 or len(partially_written_chunk_webm) > 0: made_progress = False # write audio to ffmpeg try: write_len = mp3_proc.stdin.write(partially_written_chunk_mp3) if write_len is not None: partially_written_chunk_mp3 = partially_written_chunk_mp3[write_len:] made_progress |= write_len > 0 except BlockingIOError: pass try: write_len = webm_proc.stdin.write(partially_written_chunk_webm) if write_len is not None: partially_written_chunk_webm = partially_written_chunk_webm[write_len:] made_progress |= write_len > 0 except BlockingIOError: pass _, read_made_progress = _poll_reads() made_progress |= read_made_progress if not made_progress: time.sleep(0.1) mp3_proc.stdin.close() webm_proc.stdin.close() while True: streams_done, made_progress = _poll_reads() if streams_done: break if not made_progress: time.sleep(0.1) mp3_proc.stdout.close() mp3_proc.terminate() webm_proc.stdout.close() webm_proc.terminate() # write empty bytes to modal to signal end of stream mp3_audio_chunk_queue.put( b"", partition=partition, partition_ttl=600, block=False, timeout=600, ) webm_audio_chunk_queue.put( b"", partition=partition, partition_ttl=600, block=False, timeout=600, ) if (gen_request_start_time := item.metadata.get("gen_request_start_time", None)) is not None: first_audio_write_time_millis = int(time.time() * 1000) elapsed_millis = first_audio_write_time_millis - gen_request_start_time if elapsed_millis > 0: print( f"UpsampleStub ({item.id}): Finish audio write time from gen request: {elapsed_millis // 1000}s" ) if elapsed_millis > 600_000: print( f"UpsampleStub ({item.id}): Finish audio write time from gen request took >10m: {elapsed_millis // 1000}s" ) statsd.distribution( "decoder.audio_write_finish_millis.distribution", elapsed_millis, tags=dd_tags, ) upsampled_audio = Audio.concatenate(audios) assert len(job.vae_latents) > 0, f"No vae latents generated for item_id {item.id} {job}" latents = torch.cat(job.vae_latents, dim=0).cpu().numpy() print( f"Upsample {item.id}: vae_latents shape: {latents.shape}. " f"{(time.time() - start_time):.2f}s elapsed." ) self.worker.remove_job(job) self.worker._write_audio_only(item, upsampled_audio, gain_adjust=MODEL_GAIN_ADJUST) # save npz for the vae latents if item.is_diff_infill: assert history is not None assert history.history_latents is not None assert history.future_latents is not None assert history.prompt_audio is not None self.worker._write_vae_latents_npz( item, latents, MODEL_VAE_VERSION.value, APP_NAME, n_sem_tokens=job.processed_tokens, # these are making the format consistent with the GPT infill format # note that for diffusion we care a bit less about padding (cause they shouldn't change) parent_from_start_index=int(item.metadata.get("infill_context_start_s", 0) * 25), parent_from_end_index=int(item.metadata.get("infill_context_end_s", 0) * 25), parent_clip_id=item.prompt_audio, history_latents=history.history_latents if item.is_diff_infill else None, future_latents=history.future_latents if item.is_diff_infill else None, seed=chosen_seed, semantic_tokens=history.prompt_audio if item.is_diff_infill else None, ) else: self.worker._write_vae_latents_npz( item, latents, MODEL_VAE_VERSION.value, APP_NAME, n_sem_tokens=job.processed_tokens, seed=chosen_seed, ) print( f"{item.id} uploaded upsampled audio., Took {round(time.time() - start_time, 2)} seconds" ) item.notify_progress( { "id": item.id, "model": item.model_name, "n_audios": 1, "ok": True, "gen_duration": time.time() - start_time, "ids": [item.id], "durations": [ round(upsampled_audio.duration_s, 3) if upsampled_audio is not None else 0 ], }, ) _ = self.modal_f_video_generation.spawn( item.model_dump_json(), f"image_{item.id}.jpeg", ) _ = self.modal_f_cycle.spawn( f"s3://suno-data-uploads/studio/uploads/{item.id}.mp3", s3_npz_id=item.id, encode_vae_version=None, ) print(f"Done with upsample {item.id}") @app.local_entrypoint() def main(): if MODEL == "diff_seeds_v0": input = { "id": "c7aad429-6af1-4e53-ae93-0b67540afd2a_upsample", "prompt_audio": "", "prompt_text": "", "model_name": "diff_seeds_v0", "callback_url": "https://api-staging.suno.ai/api/generate/finish-clip/", "metadata": { "prompt": "", "tags": "acoustic drum kit, backbeat, 120 bpm", }, } else: # for an upsample test input = { "id": "c7aad429-6af1-4e53-ae93-0b67540afd2a_upsample", "prompt_audio": "829fa774-4f34-4433-a66e-5ca35bdb383a.mp3", "prompt_text": """walking down the streets feeling so alive i've got my head in the clouds got a gleam in my eye every step i take it's like a brand new start no matter where i'm going i'll always find my part life is like a hard wire act we're dancing in the sky no need to worry no need to ask why with a little bit of courage we can chase our dreams no matter what comes our way we'll always be a team we're unstoppable yeah""", "model_name": "upsample", "callback_url": "https://api-staging.suno.ai/api/generate/finish-clip/", "metadata": { "prompt": """walking down the streets feeling so alive i've got my head in the clouds got a gleam in my eye every step i take it's like a brand new start no matter where i'm going i'll always find my part life is like a hard wire act we're dancing in the sky no need to worry no need to ask why with a little bit of courage we can chase our dreams no matter what comes our way we'll always be a team we're unstoppable yeah""", "tags": "epic film orchestra", }, } # if testing code functionality model = UpsampleStub() model.upsample.remote(json.dumps(input)) time.sleep(60) # if testing performance # for i in range(1_000): # new_input = input.copy() # new_input["id"] = f"{i}" # model.upsample.spawn(json.dumps(new_input)) # time.sleep(20000000) # new_input = input.copy() # new_input["id"] = f"{i}" # model.upsample.spawn(json.dumps(new_input)) # time.sleep(20000000) print("Done")