"""Audio generation application on modal.""" import fcntl import itertools import json import os import queue import random import select import threading import time import traceback from enum import Enum from typing import Optional from uuid import uuid4 import modal import numpy as np import torch from datadog import initialize, statsd from suno_utils.audio import Audio from suno_utils.utils.numbers import map_range from suno_utils.diffusion import generation as diffusion_gen from suno_utils.gpt.engine import GenerationConfig from suno_utils.gpt.generation_prompt import ALL_AUDIO_PROMPTS from suno_utils.gpt.generation_engine import make_request, semantic_codes, unshift_arrays_v2 from suno_utils.gpt.rpc_zmq import start_service_processes from suno_utils.gpt.rpc_zmq_upsample import start_service_processes as start_upsample_service_processes from suno_utils.tasks.codec_engine import CodecEngine from suno_utils.tasks.codec_engine import Request as CodecRequest from suno_utils.tasks.hoot import clean_text from suno_utils.tasks.upsample_engine import DiffusionGenerationConfig, Request from suno_utils.worker.event_queue import INBOUND_EVENT_QUEUE_NAME, Event, EventType from suno_utils.worker.loader import S3Loader from suno_utils.worker.modal_model_configs import MODEL_CONFIG_DICT, VAEVersion from suno_utils.worker.schema import HistoryPrompt, QueueItem from suno_utils.worker.tracing import distributed_trace, serialize_context, tracer from suno_utils.worker.utils import ( print_gpu_memory_usage, recursive_ls_dir, ffmpeg_stream_encode, ffmpeg_stream_encode_opus_webm, ) from suno_utils.utils.text import clean_text_tags from suno_utils.worker.modal_base import ( get_modal_base_image_with_flash_attention, get_modal_base_image_diffusion_with_flash_attention, ) from suno_utils.worker.modal_model_volume import MODEL_STORE_VOLUME_DIR, model_store_volume from suno_utils.db.clip_generation_config_handler import clip_generation_config_handler ############## CHANGE THESE ############## MODEL_CONFIG = MODEL_CONFIG_DICT["6b_sem_t1"] # Production models: # 13b_special_32, 13b_special_32_fast # 30b_t6, 30b_t6_infill # 6b_sem_t1 # Dev models: # 13b_special_32, 13b_special_test_diff, 13b_special_32_fast # 30b_t6, 30b_t6_infill # 6b_sem, 6b_sem_test, 6b_sem_test_2, 6b_sem_task, 6b_sem_eval # 6b_sem_t1, 6b_sem_t2 # 6b_sem_bluejay, 6b_sem_bluejay_test, 6b_sem_bluejay_test_2 DEPLOYMENT_TYPE = "dev" # dev, prod, msft ########################################## DURATION = MODEL_CONFIG.duration MODEL = MODEL_CONFIG.model # v2, v2_ft13 for prod for now COMPILE = MODEL_CONFIG.compile GPU = MODEL_CONFIG.gpu # A100, A10G MODEL_VERSION = MODEL_CONFIG.model_version MODEL_VAE_VERSION = MODEL_CONFIG.vae_version MODEL_GAIN_ADJUST = MODEL_CONFIG.gain_adjust # do some hacky stuff to adjust gain when testing diffusion if ( MODEL == "6b_sem_bluejay_test" or MODEL == "6b_sem_bluejay_test_2" or MODEL == "6b_sem_bluejay" or MODEL == "6b_sem_t2" ): MODEL_GAIN_ADJUST += 1.8 print(f"Model gain adjust for {MODEL} to: {MODEL_GAIN_ADJUST} dB") # only use when testing diffusion # if MODEL == "6b_sem_test": # MODEL_GAIN_ADJUST = 0.43 # print(f"Model gain adjust for {MODEL} to: {MODEL_GAIN_ADJUST} dB") # if MODEL == "6b_sem_test_2": # MODEL_GAIN_ADJUST = 3.08 # print(f"Model gain adjust for {MODEL} to: {MODEL_GAIN_ADJUST} dB") IS_3B = "v2" in MODEL IS_7B = "7b" in MODEL IS_13B = "13b" in MODEL IS_30B = "30b" in MODEL if IS_3B: from suno_utils.gpt import chirp_v2 else: from suno_utils.gpt import chirp_v2_5 as chirp_v2 GPT_CKPT_PATH = MODEL_CONFIG.gpt_ckpt_path DIFFUSION_CKPT_PATH = MODEL_CONFIG.diffusion_ckpt_path # set number of cpus. N_CPU = 8 UPSAMPLE_MAX_INPUT = MODEL_CONFIG.upsample_concurrency DECODER_MAX_INPUT = MODEL_CONFIG.decoder_concurrency TOKEN_TIMEOUT_DURATION = 60 * 12 # increase this a bit than max timeout 10 EVENT_TIMEOUT_DURATION = 60 WRITE_TO_REDIS = DEPLOYMENT_TYPE != "msft" VERBOSE_MESSAGE = DEPLOYMENT_TYPE == "dev" MOUNT_PATH = "/suno/models" TOKENS_PER_CHUNK = 25 * 15 # controls the min chunk size of streaming if "30b" in MODEL or "sem" in MODEL or "6b" in MODEL: # for auk, boost the quality for now... TOKENS_PER_CHUNK = 25 * 30 if "13b_special_32_fast" in MODEL: TOKENS_PER_CHUNK = 25 * 5 # 5 seconds -- for fast first response if MODEL == "13b_special_test": TOKENS_PER_CHUNK = 25 * 5 if MODEL == "13b_special_test_2": TOKENS_PER_CHUNK = 25 * 10 # for auk prod, faster first token if MODEL.startswith("6b_sem"): TOKENS_PER_CHUNK = 25 * 10 WORKER_NAME = f"chirpv4_engine_{MODEL}_{DEPLOYMENT_TYPE}_{DURATION}s" if GPU == "A100" or GPU == "A100-80GB": WORKER_NAME += "_A100" elif GPU == "H100": WORKER_NAME += "_H100" elif GPU == "H200": WORKER_NAME += "_H200" APP_NAME = f"engine-{WORKER_NAME}" GPT_DDOG_SERVICE = "gpt-worker" UPSAMPLE_DDOG_SERVICE = "upsample-worker" DECODER_DDOG_SERVICE = "decoder-worker" SEMANTIC_HZ = 25 EMBEDDING_RATE = 25 # 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}" TOKEN_QUEUE_NAME = f"token-queue-{DEPLOYMENT_TYPE}" CODEC_INPUT_QUEUE_NAME = f"codec-input-queue-{DEPLOYMENT_TYPE}" EVENTS_QUEUE_NAME = f"events-queue-{DEPLOYMENT_TYPE}" ERROR_QUEUE_NAME = f"error-queue-{DEPLOYMENT_TYPE}" NUM_GPT_WORKER_LIMITS = { "priority": 80, "staging": 350, "dev": 5, "prod": 450, "msft": 6, } NUM_CODEC_WORKER_LIMITS = { "priority": 16, "staging": 70, "dev": 20, "prod": 450, "msft": 6, } KEEP_WARM_DECODER = MODEL_CONFIG.keep_warm_decoder KEEP_WARM_UPSAMPLE = MODEL_CONFIG.keep_warm_upsample KEEP_WARM_GPT = MODEL_CONFIG.keep_warm_gpt KEEP_WARM_BUFFER = MODEL_CONFIG.keep_warm_buffer aws_secret = modal.Secret.from_name("studio-aws") SECRETS = [ aws_secret, modal.Secret.from_name("openai-secret"), modal.Secret.from_dict( { "DD_SITE": "datadoghq.com", "DD_ENV": DEPLOYMENT_TYPE, "DD_SERVICE": WORKER_NAME, "DD_LOGS_ENABLED": "false", "DD_TRACE_ENABLED": "true" if DEPLOYMENT_TYPE == "dev" else "false", # memory leak }, ), modal.Secret.from_name("datadog-metrics"), modal.Secret.from_name("api-callback-token"), ] base_image = get_modal_base_image_with_flash_attention() class TokenSignalCode(Enum): STREAM_COMPLETE = 0 GPT_ERROR = 1 DECODER_INPUT_TIMED_OUT = 2 UPSAMPLE_INPUT_TIMED_OUT = 3 class ChirpWorker(S3Loader): def __init__(self, max_sequences=8, max_length_s=80): S3Loader.__init__(self) num_gpus = torch.cuda.device_count() cuda_device = torch.cuda.current_device() print(f"Found {num_gpus} GPUs. Using GPU {cuda_device}.") gpt_ckpt_path = chirp_v2._get_model_if_needed(GPT_CKPT_PATH, cache_dir=MOUNT_PATH) tokenizer_path = chirp_v2._get_model_if_needed(chirp_v2.TOKENIZER_PATH, cache_dir=MOUNT_PATH) min_batch_size = min(32, max_sequences) if MODEL_CONFIG.min_batch_size is not None: min_batch_size = MODEL_CONFIG.min_batch_size self.engine_rpc_client = start_service_processes( gpt_ckpt_path, tokenizer_path=tokenizer_path, max_sequences=max_sequences, max_length_s=max_length_s, world_size=MODEL_CONFIG.world_size, compile=COMPILE, batch_increment=MODEL_CONFIG.batch_increment, min_batch_size=min_batch_size, max_batch_size=MODEL_CONFIG.max_batch_size, ) self.tokenizer = self.engine_rpc_client.get_tokenizer() self.model_cfg = self.engine_rpc_client.get_model_cfg() @staticmethod def download_models(dir_path=MOUNT_PATH): """Use AWS CLI to download models if they don't exist.""" chirp_v2.preload_models( gpt_ckpt_path=GPT_CKPT_PATH, load_whisper=False, load_semantic=False, fetch_only=True, cache_dir=dir_path, ) recursive_ls_dir(dir_path) def code_generator(self, job: str): i = 0 while True: # this is very performance sensitive. should investigate a faster method of ipc # codes = np.zeros((100, 14)) codes = self.engine_rpc_client.get_generated_codes(job, start=i) if codes is None: return i += len(codes) if isinstance(codes, np.ndarray): yield from codes # would like to make this shorter, but slower if i < TOKENS_PER_CHUNK: time.sleep(0.3) elif i < TOKENS_PER_CHUNK * 2: time.sleep(1) else: time.sleep(4) class UpsampleWorker(S3Loader): def __init__(self): S3Loader.__init__(self) start_time = time.time() print("Start loading models") tokenizer_path = diffusion_gen.get_model_if_needed( diffusion_gen.TOKENIZER_FILEPATH, cache_dir=MOUNT_PATH, ) dit_model_path = diffusion_gen.get_model_if_needed(DIFFUSION_CKPT_PATH, cache_dir=MOUNT_PATH) self.default_chunk_size = TOKENS_PER_CHUNK self.engine_rpc_client = start_upsample_service_processes( min_chunk_size=self.default_chunk_size, compile=True, dit_model_filepath=dit_model_path, tokenizer_filepath=tokenizer_path, ) print(f"Finish loading models. Took {round(time.time() - start_time, 2)} seconds") def download_model_wrapper_codec_g(): # this print is necessary to have modal rerun this when MODEL changes # Modal tracks referenced global variables # Change the name of the function to force a rerun print(f"Downloading Codec model {MODEL} {DIFFUSION_CKPT_PATH}") print("Start downloading models") _ = diffusion_gen.get_model_if_needed(diffusion_gen.TOKENIZER_FILEPATH, cache_dir=MOUNT_PATH) _ = diffusion_gen.get_model_if_needed(diffusion_gen.SEMANTIC_MODEL_FILEPATH, cache_dir=MOUNT_PATH) _ = diffusion_gen.get_model_if_needed(diffusion_gen.SEMANTIC_CLUSTERS_FILEPATH, cache_dir=MOUNT_PATH) _ = diffusion_gen.get_model_if_needed(MODEL_CONFIG.codec_ckpt_path, cache_dir=MOUNT_PATH) _ = diffusion_gen.get_model_if_needed(DIFFUSION_CKPT_PATH, cache_dir=MOUNT_PATH) print("Finish downloading models") def download_model_wrapper_gpt_a(): # this print is necessary to have modal rerun this when MODEL changes # Modal tracks referenced global variables # Change the name of the function to force a rerun print(f"Downloading GPT model {MODEL} {GPT_CKPT_PATH}") ChirpWorker.download_models() image = base_image.add_local_python_source("suno_utils", copy=False) diff_image = get_modal_base_image_diffusion_with_flash_attention().add_local_python_source( "suno_utils", copy=False ) # if we need model dependent dependencies, we can add them here # if "sem" in MODEL: # image = image.pip_install("numpy==1.26.4") app = modal.App(APP_NAME, image=image) # Queue used by Chirp/Decoder to stream tokens apptoken_queue = modal.Queue.from_name(TOKEN_QUEUE_NAME, create_if_missing=True) codec_input_queue = modal.Queue.from_name(CODEC_INPUT_QUEUE_NAME, create_if_missing=True) # Queue used by Decoder to send mp3 chunks under ephemeral stream_key. # This corresponds to chunk_queue_name in streaming_api, used to stream chunks. mp3_appchunk_queue = modal.Queue.from_name(CHUNK_QUEUE_NAME, create_if_missing=True) webm_appchunk_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) # Queue used by streaming_api to stream events like audio ready, events_queue = modal.Queue.from_name(EVENTS_QUEUE_NAME, create_if_missing=True) appinbound_event_queue = modal.Queue.from_name(INBOUND_EVENT_QUEUE_NAME, create_if_missing=True) error_queue = modal.Queue.from_name(ERROR_QUEUE_NAME, create_if_missing=True) if GPU == "A100": gpu = f"A100{':' + str(MODEL_CONFIG.world_size) if MODEL_CONFIG.world_size > 1 else ''}" elif GPU == "H100": gpu = f"H100{':' + str(MODEL_CONFIG.world_size) if MODEL_CONFIG.world_size > 1 else ''}" elif GPU == "H200": gpu = f"H200{':' + str(MODEL_CONFIG.world_size) if MODEL_CONFIG.world_size > 1 else ''}" elif GPU == "A100-80GB": gpu = f"A100-80GB{':' + str(MODEL_CONFIG.world_size) if MODEL_CONFIG.world_size > 1 else ''}" else: gpu = f"A10G{':' + str(MODEL_CONFIG.world_size) if MODEL_CONFIG.world_size > 1 else ''}" CHIRP_TIMEOUT = 660 QUEUE_ADD_TIMEOUT = 600 @app.cls( gpu=gpu, cloud=None, # "oci" if GPU == "H100" else None, secrets=SECRETS, timeout=CHIRP_TIMEOUT, # We never want to hit this timeout, it kills the worker and all jobs on it scaledown_window=400, memory=15000, retries=modal.Retries( max_retries=1, backoff_coefficient=2.0, initial_delay=5.0, ), min_containers=KEEP_WARM_GPT[DEPLOYMENT_TYPE], max_containers=NUM_GPT_WORKER_LIMITS[DEPLOYMENT_TYPE], buffer_containers=KEEP_WARM_BUFFER[DEPLOYMENT_TYPE], volumes={MODEL_STORE_VOLUME_DIR: model_store_volume}, ) @modal.concurrent(max_inputs=MODEL_CONFIG.model_concurrency) class ChirpV2Stub: # TODO: this should be renamed to sth else @modal.enter() def on_modal_enter(self): import torch torch.set_num_threads(N_CPU) num_gpus = torch.cuda.device_count() print(f"Init with gpt ckpt: {GPT_CKPT_PATH}") print(f"Found {num_gpus} GPUs.") print(f"Model gain adjust: {MODEL_GAIN_ADJUST} dB") options = {"statsd_host": "127.0.0.1", "statsd_port": 8125} initialize(**options) chirp_v2.preload_models( gpt_ckpt_path=GPT_CKPT_PATH, load_gpt=False, load_semantic=False, load_whisper=False, cache_dir=MOUNT_PATH, ) self.worker = ChirpWorker( max_sequences=MODEL_CONFIG.max_sequences, max_length_s=MODEL_CONFIG.duration, ) def _get_gconf_kwargs(self, item: QueueItem): text = item.prompt_text or "" text_tags = item.metadata.get("tags", None) if isinstance(text_tags, str): text_tags = clean_text_tags(text_tags) custom_config = {} if MODEL == "v2": cfg_coef = 1.25 # condition on non-english, higher cfg text # note this uses fastext, coupled with usage of whisper now detected_lang = chirp_v2._get_text_lang(text) if (detected_lang is not None and not detected_lang == "en") and text: cfg_coef = 1.7 print(f"detected non-english lyrics, increase cfg, lang is {detected_lang}") # condition on genre, rap elif text_tags and ("rap" in text_tags or "hip-hop" in text_tags): cfg_coef = 1.4 custom_config = dict( cfg_coef=cfg_coef, cfg_coef_tags=1.9, temp_semantic=0.9, temp_coarse=0.85, min_text_offset=128, top_k_semantic=1000, top_k_coarse=1000, top_p_semantic=None, top_p_coarse=None, eos_pad_duration_s=0.24, ) elif MODEL == "v2_ft13" or MODEL == "v2_dev": custom_config = dict( cfg_coef=1.25, cfg_coef_tags=1.9, temp_semantic=0.9, temp_coarse=0.85, min_text_offset=128, top_k_semantic=1000, top_k_coarse=1000, top_p_semantic=None, top_p_coarse=None, use_whisper=False, eos_pad_duration_s=0.24, ) elif MODEL == "7b_dpo": custom_config = dict( cfg_coef=1.1, # collect the data for now cfg_coef_neg_tags=-2, cfg_coef_tags_max_steps=None, # collect the data for now min_eos_p=0.1, use_whisper=False, ) elif MODEL == "7b": custom_config = dict( cfg_coef=1.25, # increase the cfg a bit for now cfg_coef_neg_tags=-2, use_whisper=False, ) elif MODEL == "7b_ipo" or MODEL == "7b_special": custom_config = dict( cfg_coef=1.0, # no text cfg cfg_coef_tags=2.0, cfg_coef_neg_tags=-2, min_eos_p=0.1, # this is okay use_whisper=False, ) elif MODEL == "7b_fast": custom_config = dict( cfg_coef=1.0, # for dpo-ed without cfg cfg_coef_tags=2.0, cfg_coef_neg_tags=-1.0, cfg_coef_tags_max_steps=250, # for dpo-ed without cfg min_text_offset=0, min_eos_p=0.1, use_whisper=False, n_repeat_tags=3, n_repeat_neg_tags=3, ) elif MODEL == "13b_ft_1": # this is effectively a bot model custom_config = dict( cfg_coef=1.0, # turn off the text cfg -- use it in exp cfg_coef_tags=2.0, cfg_coef_neg_tags=-1.0, min_eos_p=0.1, # this is okay use_whisper=False, min_text_offset=0, n_repeat_tags=1, n_repeat_neg_tags=1, ) elif ( MODEL == "13b_special_8" or MODEL == "13b_special_test" or MODEL == "13b_special_test_2" or MODEL == "13b_special_test_3" or MODEL == "13b_special_31" or MODEL == "13b_special_32" or MODEL == "13b_short" or MODEL == "13b_tech" or MODEL == "13b_special_test_diff" or MODEL == "13b_special_test_tech_1" or MODEL == "13b_special_33" or MODEL == "13b_special_32_fast" ): custom_config = dict( cfg_coef=1.0, # no text cfg cfg_coef_tags=2.0, cfg_coef_neg_tags=-1.0, min_eos_p=0.1, # this is okay use_whisper=False, min_text_offset=0, max_tag_len=128, n_repeat_tags=1, min_p_semantic=0.005, min_p_coarse=0.005, ) # if MODEL == "13b_special_test": # custom_config["temp_semantic"] = 0.8 elif MODEL == "13b_upload_4": custom_config = dict( cfg_coef=1.0, # no text cfg cfg_coef_tags=2.5, cfg_coef_neg_tags=-1.0, min_eos_p=0.1, # this is okay use_whisper=False, min_text_offset=0, max_tag_len=128, n_repeat_tags=2, n_repeat_neg_tags=2, ) elif ( MODEL == "30b" or MODEL == "30b_t2" or MODEL == "30b_t3" or MODEL == "30b_t5" or MODEL == "30b_t6" or MODEL == "30b_t6_test" or MODEL == "30b_t6_tech" or MODEL == "30b_t6_eval" or MODEL == "30b_classical" or MODEL == "30b_dance" or MODEL == "30b_test_diff" or MODEL == "30b_special_test_tech_1" or MODEL == "30b_t6_infill" ): custom_config = dict( cfg_coef=1.3 if (MODEL == "30b" and DEPLOYMENT_TYPE == "dev") else 1.0, # for dpo-ed without cfg cfg_coef_tags=2.0, cfg_coef_neg_tags=-1.0, cfg_coef_tags_max_steps=250, # for dpo-ed without cfg min_text_offset=0, min_eos_p=0.1, use_whisper=False, max_tag_len=128, n_repeat_tags=3, n_repeat_neg_tags=3, min_p_semantic=0.005, min_p_coarse=0.005, ) # if MODEL == "30b_t6_test": # custom_config["min_p_semantic"] = 0.1 # custom_config["min_p_coarse"] = 0.1 # increase artist / cover text cfg and max steps if item.is_artist_condition or item.is_cover_condition: custom_config["cfg_coef"] = 1.2 custom_config["cfg_coef_tags_max_steps"] = 25 * 30 custom_config["cfg_coef_max_steps"] = 25 * 30 elif MODEL == "6b_sem": custom_config = dict( cfg_coef=1.1, cfg_coef_tags=2.5 if not item.is_infill else 0.0, cfg_coef_max_steps=None, cfg_coef_tags_max_steps=None, n_repeat_tags=1, n_repeat_neg_tags=1, cfg_coef_neg_tags=0.0, temp_semantic=0.92, top_k_semantic=1500, top_p_semantic=None, min_p_semantic=0.005, min_text_offset=0, max_tag_len=512, ) elif ( MODEL == "3b_sem" or MODEL == "3b_sem_orig" or MODEL == "3b_sem_test" or MODEL == "6b_sem_test" or MODEL == "6b_sem_test_2" or MODEL == "6b_sem_tech" or MODEL == "6b_sem_t1" or MODEL == "6b_sem_eval" or MODEL == "6b_sem_task" or MODEL == "6b_sem_bluejay" or MODEL == "6b_sem_bluejay_test" or MODEL == "6b_sem_bluejay_test_2" or MODEL == "6b_sem_t2" ): # DPO-ed model, don't need crunches custom_config = dict( cfg_coef=1.0, # no text cfg cfg_coef_tags=2.0, cfg_coef_neg_tags=-1.0, min_eos_p=0.1, # this is okay use_whisper=False, min_text_offset=0, max_tag_len=512, n_repeat_tags=1, temp_semantic=0.90, top_k_semantic=1500, top_p_semantic=None, min_p_semantic=0.005, eos_pad_duration_s=0, cfg_coef_tags_max_steps=25 * 60 * 2, ) # if MODEL == "6b_sem_test": # custom_config["cfg_coef_tags_max_steps"] = 25 * 10 if item.is_artist_condition or item.is_cover_condition: custom_config["cfg_coef"] = 1.2 custom_config["cfg_coef_tags_max_steps"] = 25 * 30 custom_config["cfg_coef_max_steps"] = 25 * 30 elif item.is_playlist_condition: # Add extra audio cfg for playlist conditions boosted_prompts = ["tag", "lyrics"] + ALL_AUDIO_PROMPTS null_prompts = ["tag", "lyrics"] + ["future", "history"] # altered null stream helps for tasks here custom_config["cfg_coef_tags_max_steps"] = 25 * 120 custom_config["cfg_coef_max_steps"] = 25 * 120 custom_config["cfg_coef"] = 1.2 custom_config["custom_null_fields"] = ALL_AUDIO_PROMPTS custom_config["cfg_streams"] = [ { "stream_type": "custom", "prompts": boosted_prompts, "null_prompts": null_prompts, "weight": 2.0, "max_steps": 25 * 120, }, ] else: custom_config["n_repeat_tags"] = ( 3 if (isinstance(text_tags, str) and len(text_tags) <= 50) else 1 ) print( f"ChirpStub {item.id}: n_repeat_tags update to: {custom_config['n_repeat_tags']} for tag: {text_tags}" ) # slider supports if MODEL.startswith("6b_sem"): # get values from artist sliders # TODO: to be deprecated temp = item.metadata.get("temperature", None) constraint = item.metadata.get("distribution_size", None) tag_weight = item.metadata.get("style_weight", None) audio_weight = item.metadata.get("audio_weight", None) mixed_weirdness = item.metadata.get("weirdness_constraint", None) control_sliders = item.metadata.get("control_sliders", None) if (special_config := item.metadata.get("model_config", None)) is not None: # mask the slider controls when enabled if "masked" in special_config: control_sliders = None print(f"ChirpStub {item.id}: control_sliders: {control_sliders} is masked.") if control_sliders is not None: temp = control_sliders.get("temperature", None) constraint = control_sliders.get("distribution_size", None) tag_weight = control_sliders.get("style_weight", None) audio_weight = control_sliders.get("audio_weight", None) mixed_weirdness = control_sliders.get("weirdness_constraint", None) # first deal with the simple cases... if temp is not None and mixed_weirdness is None: if temp < 0.5: temp_scaled = map_range(temp, in_min=0.0, in_max=0.5, out_min=0.5, out_max=0.9) elif temp < 0.75: temp_scaled = map_range(temp, in_min=0.5, in_max=0.75, out_min=0.9, out_max=1.2) elif temp < 0.9: temp_scaled = map_range(temp, in_min=0.75, in_max=0.9, out_min=1.2, out_max=1.5) else: temp_scaled = map_range(temp, in_min=0.9, in_max=1.0, out_min=1.5, out_max=2.5) custom_config["temp_semantic"] = temp_scaled print(f"Using custom temperature {temp_scaled} for {item.id}") if tag_weight is not None: tag_weight_scaled = map_range(tag_weight, out_min=0.2, out_max=3.8) custom_config["cfg_coef_tags"] = tag_weight_scaled print(f"Using custom tag weight {tag_weight_scaled} for {item.id}") if constraint is not None and mixed_weirdness is None: # squishes distribution with a combination of top_k and min_p min_p_scaled = map_range(constraint, out_min=0.005, out_max=0.1) top_k_scaled = map_range(constraint, out_min=100, out_max=1500, inverse=True) custom_config["min_p_semantic"] = min_p_scaled custom_config["top_k_semantic"] = top_k_scaled print( f"Using custom min_p_semantic={min_p_scaled} and top_k_semantic={top_k_scaled} for {item.id}" ) if mixed_weirdness is not None: # this overwrites temp and constraint sliders min_p_scaled = None top_k_scaled = None temp_scaled = None temp = map_range(mixed_weirdness, in_min=0.0, in_max=1.0, out_min=0.0, out_max=1.0) constraint = map_range( mixed_weirdness, in_min=0.0, in_max=1.0, out_min=0.0, out_max=1.0, inverse=True ) if mixed_weirdness < 0.25: temp_scaled = map_range(temp, in_min=0.0, in_max=0.25, out_min=0.5, out_max=0.7) min_p_scaled = map_range( constraint, in_min=0.75, in_max=1.0, out_min=0.005, out_max=0.1 ) top_k_scaled = map_range( constraint, in_min=0.75, in_max=1.0, out_min=100, out_max=1500, inverse=True ) elif mixed_weirdness < 0.5: temp_scaled = map_range(temp, in_min=0.25, in_max=0.5, out_min=0.7, out_max=0.9) elif mixed_weirdness < 0.75: temp_scaled = map_range(temp, in_min=0.5, in_max=0.75, out_min=0.9, out_max=1.2) elif mixed_weirdness < 0.9: temp_scaled = map_range(temp, in_min=0.75, in_max=0.9, out_min=1.2, out_max=2.0) min_p_scaled = map_range( constraint, in_min=0.1, in_max=0.25, out_min=0.005, out_max=0.1 ) top_k_scaled = map_range( constraint, in_min=0.1, in_max=0.25, out_min=250, out_max=1500, inverse=True ) else: temp_scaled = map_range(temp, in_min=0.9, in_max=1.0, out_min=1.5, out_max=2.5) if min_p_scaled is not None: custom_config["min_p_semantic"] = min_p_scaled print(f"Using custom min_p_semantic={min_p_scaled} for {item.id}") if top_k_scaled is not None: custom_config["top_k_semantic"] = top_k_scaled print(f"Using custom top_k_semantic={top_k_scaled} for {item.id}") if temp_scaled is not None: custom_config["temp_semantic"] = temp_scaled print(f"Using custom temperature {temp_scaled} for {item.id}") if audio_weight is not None: # add an additional CFG stream for audio adherence null_prompts = None boosted_prompts = ["tag", "lyrics"] + ALL_AUDIO_PROMPTS if item.is_artist_condition: null_prompts = ["tag", "lyrics"] + ["cover", "future", "history"] audio_weight_scaled = map_range(audio_weight, out_min=0.0, out_max=2.0) elif item.metadata.get("task") == "extend" or item.is_upload_extend: audio_weight_scaled = map_range(audio_weight, out_min=0.0, out_max=2.5) null_prompts = ["tag", "lyrics"] + ["artist", "cover", "future"] elif item.is_cover_condition: audio_weight_scaled = map_range(audio_weight, out_min=0.0, out_max=2.0) null_prompts = ["tag", "lyrics"] + ["artist", "future", "history"] # TODO: consider separating this out for stacked tasks elif item.is_cover_extend: audio_weight_scaled = map_range(audio_weight, out_min=0.0, out_max=2.0) null_prompts = ["tag", "lyrics"] + ["artist", "future"] elif item.is_artist_cover_condition: audio_weight_scaled = map_range(audio_weight, out_min=0.0, out_max=2.0) null_prompts = ["tag", "lyrics"] + ["future", "history"] elif item.is_artist_cover_extend: audio_weight_scaled = map_range(audio_weight, out_min=0.0, out_max=2.0) null_prompts = ["tag", "lyrics"] + ["future"] elif item.is_multi_artist_consistency: audio_weight_scaled = map_range(audio_weight, out_min=0.0, out_max=4.0) null_prompts = ["tag", "lyrics"] + ["future", "history"] elif item.is_playlist_condition: audio_weight_scaled = map_range(audio_weight, out_min=0.0, out_max=4.0) null_prompts = ["tag", "lyrics"] + ["future", "history"] elif item.is_overpainting: audio_weight_scaled = map_range(audio_weight, out_min=0.0, out_max=2.0) null_prompts = ["tag", "lyrics"] + ["future", "history"] elif item.is_underpainting: audio_weight_scaled = map_range(audio_weight, out_min=0.0, out_max=2.0) null_prompts = ["tag", "lyrics"] + ["future", "history"] if null_prompts is not None: # we're running a valid task # altered null stream helps for tasks here print(f"Using custom audio weight {audio_weight_scaled} for {item.id}") custom_config["cfg_coef_tags_max_steps"] = 25 * 120 custom_config["cfg_coef_max_steps"] = 25 * 120 custom_config["cfg_coef"] = 1.2 custom_config["custom_null_fields"] = ALL_AUDIO_PROMPTS custom_config["cfg_streams"] = [ { "stream_type": "custom", "prompts": boosted_prompts, "null_prompts": null_prompts, "weight": audio_weight_scaled, "max_steps": 25 * 120, }, ] else: raise ValueError(f"Unknown model {MODEL}") # for infill and extend, we want to use a lower temperature, to be more consistent if item.is_infill or item.metadata.get("task") == "extend" or item.is_upload_extend: custom_config["temp_semantic"] = 0.8 input_negative_tags = item.metadata.get("negative_tags") if isinstance(input_negative_tags, str) and len(input_negative_tags) > 0: # if negative tags are provided, we want to repeat them more times custom_config["n_repeat_neg_tags"] = 3 # Note that experiment config will overwrite default config if special_config := item.metadata.get("model_config"): # has specific config instruction print(f"ChirpStub: load specific config: {special_config}, {item.id}") # update the custom config # note this will overwrite existing parameters custom_config.update(**special_config) # Note forced infer config will overwrite experiment config if (special_config := item.metadata.get("forced_infer_config")) and DEPLOYMENT_TYPE == "dev": # has specific config instruction print(f"item.metadata: {item.metadata}") print(f"ChirpStub: load forced infer config: {special_config}, {item.id}") # update the custom config # note this will overwrite existing parameters custom_config.update(**{k: v for k, v in special_config.items() if v is not None}) # Filter kwargs that are valid properties of GenerationConfig valid_gconf_props = set(vars(GenerationConfig()).keys()) filtered_config = {k: v for k, v in custom_config.items() if k in valid_gconf_props} # Print removed keys for debugging removed_keys = set(custom_config.keys()) - set(filtered_config.keys()) if removed_keys: print(f"ChirpStub: Removed invalid GenerationConfig properties: {removed_keys}") custom_config = filtered_config return custom_config @tracer.wrap(service=GPT_DDOG_SERVICE) def _make_gconf(self, item: QueueItem, history: HistoryPrompt) -> GenerationConfig: # build a config text = item.prompt_text or "" text_tags = item.metadata.get("tags", None) if text_tags == "random": text_tags = None text_neg_tags = None text_start_control_tags = None text_end_control_tags = None is_mumble = item.metadata.get("is_mumble", False) if is_mumble: text = "[mumble mode]" # cleaned_text will remove all meta tags like: [chorus] cleaned_text = clean_text(text).strip() is_instrumental = len(cleaned_text) == 0 and not is_mumble text_neg_tags = "repetitive, loop," if IS_13B: text_neg_tags += " noisy, distorted," if is_instrumental: # there is no asr-able text # very likely this should be instrumental # consistent with the FE update text_tags = "instrumental, " + (text_tags or "") + ", instrumental" if MODEL == "13b_upload_4": # we want to trick the audio upload worker # to generate a bit longer audio text = "[instrumental] \n [instrumental] \n ... \n [instrumental] \n [instrumental]" print(f"Request {item.id} hack text input to : {text}") else: text = "[instrumental, instrumental, instrumental]" + text text_neg_tags += " vocal, sing, voices, chant, hum, speech" if item.is_one_box_generation or history.prompt_audio is None: # can be None, or empty "" # for one box, we want to try to end the songs print("ChirpStub: Pad start/end to one-box generation or history is None.") # from v3 for 7b we will use the `{}` annotation # note that for instrumental, we don't want vocals, and we don't pad control {end} text_start_control_tags = "{start:0} " if IS_30B else "{start} " if not is_instrumental and (MODEL == "13b_short" or MODEL == "7b_fast"): text_start_control_tags = "{vocals:start} " text_end_control_tags = "{end:vocals:end} " if not is_instrumental and IS_7B: text_start_control_tags = "{start;vocals:start} " if not is_instrumental and IS_30B: text_start_control_tags = ( "{start:0;vocals:start} " if random.random() > 0.5 else "{start:0;vocals:intro} " ) if history.prompt_audio is None: text_end_control_tags = "{end} " # we will try to end one box using end tag if item.is_one_box_generation: text = text + (" \n [end]" if not is_instrumental else "") # for semantic only models, we don't want to pass in any control tags yet if MODEL.startswith("6b_sem") or MODEL.startswith("3b_sem"): text_start_control_tags = None text_end_control_tags = None # special infilling control tag overwrite if IS_30B and item.is_infill: infill_start_s = item.metadata.get("infill_start_s", 0) infill_end_s = item.metadata.get("infill_end_s", None) # these are intro/outro durations control_tags = [] if (infill_duration := item.metadata.get("infill_dur_s")) and item.metadata.get( "task" ) == "infill": replaced_duration = max(round(infill_duration), 1) if infill_start_s == 0: control_tags.append("start:0") history_context_duration = ( history.prompt_audio.shape[0] // 25 if isinstance(history.prompt_audio, np.ndarray) else 0 ) future_context_duration = ( history.future_audio.shape[0] // 25 if isinstance(history.future_audio, np.ndarray) else 0 ) total_duration = int( history_context_duration + replaced_duration + future_context_duration ) control_tags.append("duration:" + str(total_duration)) elif item.metadata.get("task") == "infill_intro": replaced_duration = max(int(infill_end_s - infill_start_s), 1) intro_duration = item.metadata.get("infill_dur_s", 10) future_context_duration = ( history.future_audio.shape[0] // 25 if isinstance(history.future_audio, np.ndarray) else 20 ) total_duration = int(future_context_duration + replaced_duration + intro_duration) control_tags.append("start:0") control_tags.append("duration:" + str(total_duration)) elif item.metadata.get("task") == "infill_outro": replaced_duration = max(int(infill_end_s - infill_start_s), 1) outro_duration = item.metadata.get("infill_dur_s", 10) history_context_duration = ( history.prompt_audio.shape[0] // 25 if isinstance(history.prompt_audio, np.ndarray) else 20 ) total_duration = int(history_context_duration + replaced_duration + outro_duration) control_tags.append("duration:" + str(total_duration)) if control_tags: text_start_control_tags = "{" + ";".join(control_tags) + "} " print(f"ChirpStub: item {item.id} infilling control tags: {text_start_control_tags}.") elif item.is_infill and self.worker.model_cfg.use_delta_infill: # auk uses delta infill # infill_lyrics = item.metadata.get("infill_lyrics", "") if (infill_duration := item.metadata.get("infill_dur_s")) and item.metadata.get( "task" ) == "infill": replaced_duration = max(round(infill_duration), 1) history_context_duration = ( history.prompt_audio.shape[0] // 25 if isinstance(history.prompt_audio, np.ndarray) else 0 ) future_context_duration = ( history.future_audio.shape[0] // 25 if isinstance(history.future_audio, np.ndarray) else 0 ) total_duration = int( history_context_duration + replaced_duration + future_context_duration ) text_start_control_tags = "{" + "duration:" + str(total_duration) + "} " # text = infill_lyrics # load the other inference parameters here inference_parameters = self._get_gconf_kwargs(item) # overwrite the negative tags from prompt input_negative_tags = item.metadata.get("negative_tags") if isinstance(input_negative_tags, str) and len(input_negative_tags) > 0: text_negative_tags = input_negative_tags # if negative tags are provided, we want to repeat them more times inference_parameters["n_repeat_neg_tags"] = 3 else: text_negative_tags = text_neg_tags if input_negative_tags: print(f"ChirpStub: load specific negative_tags: {input_negative_tags}, {item.id}") vocal_gender = item.metadata.get("vocal_gender", None) if vocal_gender is not None: if vocal_gender == "m": text_tags = "male, " + (text_tags or "") text_neg_tags = "female, " + (text_neg_tags or "") elif vocal_gender == "f": text_tags = "female, " + (text_tags or "") text_neg_tags = "male, " + (text_neg_tags or "") else: print(f"ChirpStub: item {item.id} invalid vocal gender: {vocal_gender}. Won't augment.") if DEPLOYMENT_TYPE == "dev": print( f"ChirpStub: item {item.id} vocal gender: {vocal_gender}, text_tags: {text_tags}, text_neg_tags: {text_neg_tags}" ) # set the generation max duration if history.prompt_audio is not None: if MODEL == "13b_upload_4" or item.is_upload_extend: # will dynamically change based on input duration max_gen_duration_s = MODEL_CONFIG.duration - history.prompt_audio.shape[0] // 25 print( f"Upload extend clip: {item.id} customized max generation duration is {max_gen_duration_s}." ) else: max_gen_duration_s = MODEL_CONFIG.duration - MODEL_CONFIG.max_history_duration_s else: max_gen_duration_s = MODEL_CONFIG.duration if item.is_infill: max_gen_duration_s = 120 # for bot generation, we want to limit the generation duration, but keep it random :) if item.is_bot_generation and MODEL == "13b_ft_1": # this is a beta distribution, mean is roughply 10 / (10 + 10) = 0.5 ~ 120 sec max_gen_duration_s = 10 + int(random.betavariate(10, 10) * 230) print(f"Bot generation: {item.id}, max duration: {max_gen_duration_s}.") history_text = history.prompt_lyrics # for personalization test only # in training the augmentation happen before the text # hence we should also do this when extending # if not item.is_one_box_generation: # augment_personlized_text = "\n[++++++++++]\n" # else: # augment_personlized_text = "\n[----------]\n" # print( # f"ChirpStub {item.id}: One-box generation {item.is_one_box_generation}, adding augment_personlized_text: {augment_personlized_text}" # ) # if MODEL == "13b_special_test_3": # if history_text: # history_text = augment_personlized_text + history_text # else: # text = augment_personlized_text + text # Trim history audio cause if we AB test auk (under v4) with v4, history loader will load v4 code with coarse. if "6b" in MODEL: n_gen_dim = 1 # make sure we trim all the history audio so that it works... if history.prompt_audio is not None and history.prompt_audio.shape[-1] > n_gen_dim: print( f"ChirpStub {item.id}: Trimming history.prompt_audio from {history.prompt_audio.shape[-1]} to {n_gen_dim}" ) history.prompt_audio = history.prompt_audio[:, :n_gen_dim] if history.future_audio is not None: history.future_audio = history.future_audio[:, :n_gen_dim] if history.cover_audio is not None: history.cover_audio = history.cover_audio[:, :n_gen_dim] if history.artist_audio is not None: history.artist_audio = history.artist_audio[:, :n_gen_dim] if history.playlist_audio is not None: history.playlist_audio = [t[:, :n_gen_dim] for t in history.playlist_audio] if history.multi_artist_audio is not None: history.multi_artist_audio = [t[:, :n_gen_dim] for t in history.multi_artist_audio] if history.underpainting_audio is not None: history.underpainting_audio = history.underpainting_audio[:, :n_gen_dim] if history.overpainting_audio is not None: history.overpainting_audio = history.overpainting_audio[:, :n_gen_dim] # trim the audio condition for playlist condition if item.is_playlist_condition: # no extra augmentation for now. text = item.prompt_text or "" text_tags = item.metadata.get("tags", None) text_negative_tags = item.metadata.get("negative_tags", None) if item.is_overpainting and history.overpainting_audio is not None: expected_duration = int(history.overpainting_audio.shape[0] // 25) text_start_control_tags = "{" + "duration:" + str(expected_duration) + "} " print(f"ChirpStub: item {item.id} overpainting control tags: {text_start_control_tags}.") if item.is_underpainting and history.underpainting_audio is not None: expected_duration = int(history.underpainting_audio.shape[0] // 25) text_start_control_tags = "{" + "duration:" + str(expected_duration) + "} " print(f"ChirpStub: item {item.id} underpainting control tags: {text_start_control_tags}.") # hack control tags if isinstance(item.title, str) and item.title.startswith("HACK{") and item.title.endswith("}"): text_start_control_tags = item.title.replace("HACK", "") print(f"ChirpStub: item {item.id} hack control tags: {text_start_control_tags}.") # make the configuration cfg = chirp_v2.GenerationConfig( text=text, text_tags=text_tags, text_neg_tags=text_negative_tags, text_start_control_tags=text_start_control_tags, text_end_control_tags=text_end_control_tags, n_batch=1, max_gen_duration_s=max_gen_duration_s, rep_penality=0, history_arr=history.prompt_audio, history_text=history_text, future_arr=history.future_audio, cover_arr=history.cover_audio, artist_arr=history.artist_audio, playlist_arr=history.playlist_audio, multi_artist_arr=history.multi_artist_audio, overpaint_arr=history.overpainting_audio, underpaint_arr=history.underpainting_audio, history_do_asr=False, stream=True, allow_eos=True, **inference_parameters, # pass in the other configuration parameters ) if VERBOSE_MESSAGE: print(f"ChirpStub: Model {item.model_name}: gconf: {cfg}") return cfg @modal.method() @distributed_trace("gpt_generate", GPT_DDOG_SERVICE, env_name=DEPLOYMENT_TYPE) def generate(self, queue_item: str, history: Optional[HistoryPrompt] = None): tracer.current_span().set_tag("model", MODEL) tracer.current_span().set_tag("duration", DURATION) tracer.current_span().set_tag("deployment_type", DEPLOYMENT_TYPE) start_time = time.time() item = QueueItem(**json.loads(queue_item)) if item.metadata.get("gen_request_start_time") is None: item.metadata["gen_request_start_time"] = int(time.time() * 1000) item_id = item.id ids = [item_id] if history is None: history = HistoryPrompt() try: if DEPLOYMENT_TYPE == "dev": # these are mostly for debugging but not practical for prod os.environ["CUDA_LAUNCH_BLOCKING"] = "1" os.environ["TORCH_CUDA_ALLOC_SYNC"] = "1" if VERBOSE_MESSAGE: print_gpu_memory_usage(self.__class__.__name__) torch.cuda.reset_max_memory_allocated() request_prep_time = time.time() print( f"ChirpStub: Adding job to queue: {item_id}, " f"modalID {modal.current_input_id()}, {request_prep_time - start_time:.2f}s elapsed." ) device_name = torch.cuda.get_device_name() dd_tags = [ f"model:{MODEL}", f"env:{DEPLOYMENT_TYPE}", f"env_name:{WORKER_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}", f"custom_lyrics:{not bool(item.metadata.get('gpt_description_prompt'))}", # f"modal_task_id:{os.environ.get('MODAL_TASK_ID', 'unknown')}", # this increases costs ] # this watches for events from clients # commented out for now since seems to cause performance dips when times out # inbound_event_queue = InboundEventQueue( # appinbound_event_queue, # item_id, # timeout=TOKEN_TIMEOUT_DURATION, # ) # inbound_event_queue.start() gconf = self._make_gconf(item, history) max_history_duration_s = MODEL_CONFIG.max_history_duration_s # for infilling, we don't want to limit the history duration # they can be as long as possible... if item.is_infill: max_history_duration_s = None expected_history_dim = ( self.worker.model_cfg.semantic_n_codebooks + self.worker.model_cfg.coarse_n_codebooks ) if gconf.history_arr is not None and gconf.history_arr.shape[-1] != expected_history_dim: print( f"ChirpStub: History misshapen for {item_id}: expected (_, {expected_history_dim}), got {gconf.history_arr.shape}" ) gconf = chirp_v2.prep_gconf( gconf, max_history_duration_s=max_history_duration_s, cfg=self.worker.model_cfg, ) # 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=TOKEN_TIMEOUT_DURATION) print(f"ChirpStub: Put stream key {stream_key} for {item_id}") request = make_request(item_id, gconf, self.worker.model_cfg, self.worker.tokenizer) with tracer.trace("add_request_to_worker"): start_add_time = time.time() print( f"Adding request to worker: {item_id}, {start_add_time - start_time:.2f}s elapsed." ) request = make_request( item_id, gconf, self.worker.model_cfg, self.worker.tokenizer, ) clip_generation_config_handler.put_gpt_config(item_id, gconf) result = self.worker.engine_rpc_client.add_request(request) if result == "Job already exists": raise ValueError(f"ChirpStub: Job already exists for {item_id}") while self.worker.engine_rpc_client.get_job_state(request.id) is None: time.sleep(0.2) end_add_time = time.time() statsd.distribution( "engine.add_request_duration_seconds", end_add_time - start_add_time, tags=dd_tags ) statsd.distribution( "engine.free_caches", self.worker.engine_rpc_client.get_free_caches(), tags=dd_tags ) statsd.distribution( "engine.used_caches", self.worker.engine_rpc_client.get_used_caches(), tags=dd_tags ) statsd.gauge( "engine.evictions_last_minute", self.worker.engine_rpc_client.get_evictions_last_minute(), tags=dd_tags, ) num_jobs = self.worker.engine_rpc_client.get_num_jobs() num_streams = self.worker.engine_rpc_client.get_num_streams() avg_streams_per_job = num_streams / num_jobs if num_jobs > 0 else 0 print( f"Number of jobs: {num_jobs}, Number of streams: {num_streams}, Avg streams per job: {avg_streams_per_job}" ) statsd.distribution("engine.num_jobs", num_jobs, tags=dd_tags) statsd.distribution("engine.num_streams", num_streams, tags=dd_tags) statsd.distribution("engine.streams_per_job", avg_streams_per_job, tags=dd_tags) utilization = torch.cuda.utilization() print(f"CUDA Utilization: {utilization}") statsd.distribution("cuda.utilization", utilization, tags=dd_tags) print( f"{item_id}: Time till added request to worker: {end_add_time - start_time:.2f} seconds" ) statsd.distribution("engine.time_before_add", start_add_time - start_time, tags=dd_tags) statsd.distribution("engine.time_after_add", end_add_time - start_time, tags=dd_tags) code_stream = self.worker.code_generator(item_id) aligned_code_stream = semantic_codes( code_stream, self.worker.model_cfg, n_skip_semantic=gconf.n_skip_semantic ) # if we want to include history in the output include_history_s = item.metadata.get("include_history_s", 0) include_future_s = item.metadata.get("include_future_s", 0) if include_history_s > 0 and gconf.history_arr is not None: history_codes = gconf.history_arr[-int(25 * include_history_s) :, 0].astype(np.int32) print( f"ChirpStub: Including {include_history_s} seconds of history. {history_codes.shape} tokens." ) aligned_code_stream = itertools.chain( torch.from_numpy(history_codes), aligned_code_stream, ) if include_future_s > 0 and gconf.future_arr is not None: future_codes = gconf.future_arr[: int(25 * include_future_s), 0].astype(np.int32) print( f"ChirpStub: Including {include_future_s} seconds of future. {future_codes.shape} tokens." ) aligned_code_stream = itertools.chain( aligned_code_stream, torch.from_numpy(future_codes), ) codes = [] stream_size = 0 start_decoder = False with tracer.trace("generate_tokens") as span: span_str = serialize_context(span.context) for code in aligned_code_stream: # if not inbound_event_queue.local_queue.empty(): # event = inbound_event_queue.local_queue.get() # print(f"ChirpStub: Received event {event}") # if event.event_type is EventType.CANCEL: # print(f"ChirpStub: Canceling generation for {item_id}") # # Note this doesnt remove the job from the engine immediately # # it will be removed when rest of this function finishes # break codes.append(code.numpy().astype(np.int16)) if len(codes) >= TOKENS_PER_CHUNK: with tracer.trace("put_tokens"): tracer.current_span().set_tag("stream_size", stream_size) tracer.current_span().set_tag("n_codes", len(codes)) apptoken_queue.put( codes, partition=stream_key, partition_ttl=TOKEN_TIMEOUT_DURATION, block=False, timeout=QUEUE_ADD_TIMEOUT, ) stream_size += len(codes) codes = [] if not start_decoder: start_decoder = True # start the decoder with sth in the dict UpsampleStub().generate.spawn( item.model_dump_json(), history=history, stream_key=stream_key, parent_context=span_str, ) print( f"ChirpStub: Start the decoder: {item_id}, {time.time() - start_time:.2f}s elapsed." ) statsd.distribution( "engine.time_to_start_decoder", time.time() - start_time, tags=dd_tags, ) if time.time() - start_time > CHIRP_TIMEOUT - 120: # manually catch timeouts to prevent modal from killing the worker print( f"ChirpStub: Timeout for {item_id}, {time.time() - start_time:.2f}s elapsed." ) raise TimeoutError( f"ChirpStub: Timeout for {item_id}, {time.time() - start_time:.2f}s elapsed." ) print(f"{item_id}: Time till sending last chunk: {time.time() - start_time:.2f} seconds") # if the gen is really short, we need to send the last chunk, and still start the job with tracer.trace("send_last_chunk", resource="token_queue_operation"): apptoken_queue.put( codes + [TokenSignalCode.STREAM_COMPLETE], partition=stream_key, partition_ttl=TOKEN_TIMEOUT_DURATION, block=False, timeout=QUEUE_ADD_TIMEOUT, ) stream_size += len(codes) print(f"ChirpStub: Item id: {item_id}, total stream size: {stream_size}") print(f"{item_id}: Time till sent last chunk: {time.time() - start_time:.2f} seconds") statsd.distribution( "engine.time_after_send_last_chunk", time.time() - start_time, tags=dd_tags ) if not start_decoder: start_decoder = True UpsampleStub().generate.spawn( item.model_dump_json(), history=history, stream_key=stream_key, ) start_write_npz_time = time.time() with tracer.trace("write_npz"): while True: # we need to wait for the job to complete before writing the npz # we get all semantic codes before completion so we need to wait job_state = self.worker.engine_rpc_client.get_job_state(request.id) if job_state["completed"]: break time.sleep(0.5) print(f"{item_id}: Job state: {job_state}") codes = self.worker.engine_rpc_client.get_generated_codes(request.id) raw_arrays = unshift_arrays_v2( codes, self.worker.model_cfg, job_state["eos_step"], semantic_only=self.worker.model_cfg.is_semantic_only, ) raw_arrays = raw_arrays.cpu().numpy() if len(list(semantic_codes(codes, self.worker.model_cfg))) != raw_arrays.shape[0]: print( f"{item_id}: Semantic codes length {len(list(semantic_codes(codes, self.worker.model_cfg)))} does not match raw arrays length {raw_arrays.shape[0]}" ) # raise AssertionError( # f"{item_id}: Semantic codes length {len(list(semantic_codes(codes, self.worker.model_cfg)))} does not match raw arrays length {raw_arrays.shape[0]}" # ) self.worker._write_npz( item, raw_arrays, APP_NAME, MODEL_VERSION, generated_arr=raw_arrays, history_arr=gconf.history_arr, future_arr=gconf.future_arr, pre_history_arr=history.pre_history_arr, post_future_arr=history.post_future_arr, cover_arr=history.cover_audio, artist_arr=history.artist_audio, playlist_arr=history.playlist_audio, multi_artist_arr=history.multi_artist_audio, overpainting_arr=history.overpainting_audio, underpainting_arr=history.underpainting_audio, history_lyrics=history.prompt_lyrics, fuller_array=job_state["full_array"], n_skip_semantic=gconf.n_skip_semantic, ) statsd.distribution("engine.job_itl", job_state["itl"], tags=dd_tags) statsd.distribution("engine.job_ttl", job_state["ttl"], tags=dd_tags) statsd.distribution("engine.job_ttft", job_state["ttft"], tags=dd_tags) print(f"{item_id}: Time writing npz: {time.time() - start_write_npz_time:.2f} seconds") print(f"{item_id}: Time till removing job: {time.time() - start_time:.2f} seconds") statsd.distribution("engine.time_before_remove_job", time.time() - start_time, tags=dd_tags) self.worker.engine_rpc_client.remove_job(request.id) if VERBOSE_MESSAGE: max_mem = torch.cuda.max_memory_allocated() / 1e9 print(f"ChirpStub: Max memory allocated: {max_mem:.2f} GB") except Exception as e: error_message = str(e) # or e.args[0] if args attribute is used if "Job already exists" in error_message: print(f"ChirpStub: Job {item_id} already exists, skipping.") return # put the error in the error queue error_queue.put(error_message, partition=item_id, partition_ttl=TOKEN_TIMEOUT_DURATION) finish_time = time.time() apptoken_queue.put( [TokenSignalCode.GPT_ERROR], partition=item_id, partition_ttl=TOKEN_TIMEOUT_DURATION, block=False, ) item.notify_progress( { "id": item_id, "model": item.model_name, "n_audios": len(ids), "ok": 0, "gen_duration": finish_time - start_time, "error_type": "generation_failure", "error_message": "Failed generation.", } ) # TODO: such workers should be stopped ASAP. find a proper modal way to do this if "CUDA error: an illegal memory access was encountered" in error_message: print("ChirpStub: Caught a CUDA memory error, this worker is busted. Sleep to death.") modal.experimental.stop_fetching_inputs() if "stream has been closed" in error_message: print("ChirpStub: Caught a stream closed error, this worker is busted. Sleep to death.") modal.experimental.stop_fetching_inputs() self.worker.engine_rpc_client.remove_job(request.id) print(f"ChirpStub: item_id: {item_id} failed with error: {e}") traceback.print_exc() raise e @modal.exit() def cleanup_processes(self): print("ChirpStub: Cleaning up processes") self.worker.engine_rpc_client.terminate_processes() def fill_local_token_queue_upsample_upsample(local_token_queue: queue.Queue, partition: str): # Ensure the queue has not expired. if apptoken_queue.len(partition=partition) == 0: # Fail the decoding job for this partition. print(f"UpsampleStub: No tokens in queue {partition}, putting UPSAMPLE_INPUT_TIMED_OUT") local_token_queue.put([TokenSignalCode.UPSAMPLE_INPUT_TIMED_OUT]) return # Stream tokens into local queue. for token_batch in apptoken_queue.iterate(partition=partition, item_poll_timeout=120): local_token_queue.put(token_batch) if ( isinstance(token_batch[-1], TokenSignalCode) and token_batch[-1] == TokenSignalCode.STREAM_COMPLETE ): print(f"UpsampleStub: Received STREAM_COMPLETE token for {partition}") return # print( # f"Put {len(token_batch)} tokens into queue, {local_token_queue.qsize()} tokens in queue {partition}." # ) # If we exited the loop without STREAM_COMPLETE, Chirp lagged out and stopped generating tokens. print(f"UpsampleStub: No STREAM_COMPLETE token, putting UPSAMPLE_INPUT_TIMED_OUT for {partition}") local_token_queue.put([TokenSignalCode.UPSAMPLE_INPUT_TIMED_OUT]) # TODO: when the GPT worker is spinning up -- they should trigger a multi-warm up call to the decoder # So the decoder scales up before the tokens are generated @app.cls( image=diff_image, cpu=12, gpu="H100", secrets=SECRETS, timeout=600, # TODO: with current issue this is too long scaledown_window=400, memory=32000, retries=modal.Retries( max_retries=1, backoff_coefficient=2.0, initial_delay=5.0, ), min_containers=KEEP_WARM_UPSAMPLE[DEPLOYMENT_TYPE], max_containers=NUM_CODEC_WORKER_LIMITS[DEPLOYMENT_TYPE], buffer_containers=KEEP_WARM_BUFFER[DEPLOYMENT_TYPE], volumes={MODEL_STORE_VOLUME_DIR: model_store_volume}, # cloud="oci", # to pin to h100s ) @modal.concurrent(max_inputs=UPSAMPLE_MAX_INPUT) class UpsampleStub: @modal.enter() def on_modal_enter(self): import torch print(f"Init with diffusion ckpt: {DIFFUSION_CKPT_PATH}") torch.set_num_threads(N_CPU) num_gpus = torch.cuda.device_count() device_name = torch.cuda.get_device_name() print(f"Found {num_gpus} GPUs. Using CUDA device: {device_name}") options = {"statsd_host": "127.0.0.1", "statsd_port": 8125} initialize(**options) self.worker = UpsampleWorker() 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)] def print_info(self, queue_item_id: str, message: str): print(f"UpsampleStub ({queue_item_id}): {message}") @modal.method() @distributed_trace("upsample_generate", UPSAMPLE_DDOG_SERVICE, env_name=DEPLOYMENT_TYPE) def generate(self, item_json: str, history: HistoryPrompt, stream_key: Optional[str] = None): start_time = time.time() if VERBOSE_MESSAGE: print_gpu_memory_usage(self.__class__.__name__) torch.cuda.reset_max_memory_allocated() item = QueueItem(**json.loads(item_json)) item_id = item.id device_name = torch.cuda.get_device_name() dd_tags = [ f"model:{MODEL}", f"env:{DEPLOYMENT_TYPE}", f"env_name:{WORKER_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}", ] gen_request_start_time = item.metadata.get("gen_request_start_time", None) if gen_request_start_time is not None: statsd.distribution( "diffusion.gen_request_start_time", int(time.time() * 1000) - gen_request_start_time, tags=dd_tags, ) print( f"{item_id}: Adding job to queue, " f"modalID {modal.current_input_id()}, {time.time() - start_time:.2f}s elapsed." ) token_queue = queue.Queue() # 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 loader_thread = threading.Thread( target=fill_local_token_queue_upsample_upsample, args=(token_queue, partition), daemon=True, ) loader_thread.start() error = None attempts = 2 self.print_info(item_id, "Starting decoder job") text_cfg_coef = item.metadata.get("text_cfg_coef", 2.0) steps = item.metadata.get("steps", 10) if MODEL == "30b_test_diff" or MODEL == "13b_special_test_diff": text_cfg_coef = 2.0 steps = 10 chosen_seed = random.choice(self.seed_pool) diffusion_inference_config = dict( text_cfg_coef=text_cfg_coef, steps=steps, seed=chosen_seed, ) 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 # default noise level is 0.5 -- we should still tune and understand this better diffusion_inference_config["noise_ctx_level"] = 0.75 diffusion_inference_config["noise_ctx_pad_len"] = 0 # Note that experiment config will overwrite default config if special_config := item.metadata.get("model_config"): # has specific config instruction print(f"UpsampleStub {item.id}: load specific config: {special_config}.") # update the custom config # note this will overwrite existing parameters diffusion_inference_config.update(**special_config) # Filter kwargs that are valid properties of GenerationConfig valid_gconf_props = set(vars(DiffusionGenerationConfig()).keys()) filtered_config = {k: v for k, v in diffusion_inference_config.items() if k in valid_gconf_props} # Print removed keys for debugging removed_keys = set(diffusion_inference_config.keys()) - set(filtered_config.keys()) if removed_keys: print(f"UpsampleStub {item_id}: Removed invalid GenerationConfig properties: {removed_keys}") print(f"UpsampleStub {item.id}: diffusion_inference_config: {filtered_config}") try: received_tokens = [] while attempts > 0: error = None attempts -= 1 audio_idx = 0 # TODO: Tony add some ab testing variations of diffusion config here, maybe? current_history_latents = None # t, l if history.history_latents is not None: current_history_latents = torch.from_numpy(history.history_latents) assert current_history_latents.dtype == torch.float16 print(f"UpsampleStub {item_id}: loaded history VAE {current_history_latents.shape}") if current_history_latents.shape[0] == 0: print(f"Warning {item_id}: No history latents loaded?") current_history_latents = None if item.is_infill: include_history_s = item.metadata.get("include_history_s", 0) include_future_s = item.metadata.get("include_future_s", 0) if include_history_s > 0 and current_history_latents is not None: # truncate the history latents current_total_history_n_tokens = current_history_latents.shape[0] # TODO: 25 should be the vae freq current_total_history_n_tokens -= int(25 * include_history_s) if current_total_history_n_tokens > 0: current_history_latents = current_history_latents[ :current_total_history_n_tokens ] print( f"UpsampleStub {item_id}: infilling, truncated history latents to {current_history_latents.shape}" ) current_future_latents = None if ( (history.future_latents is not None) # Only v2 diffusion with infill can accept future latents # and 30b_t6_infill is the only model that uses it so far and (MODEL_VAE_VERSION == VAEVersion.V_VAE_25_TUNED_2) and (MODEL == "30b_t6_infill" or MODEL.startswith("6b_sem")) ): current_future_latents = torch.from_numpy(history.future_latents) assert current_future_latents.dtype == torch.float16 print(f"UpsampleStub {item_id}: loaded future VAE {current_future_latents.shape}") if current_future_latents.shape[0] == 0: print(f"Warning {item_id}: No future latents loaded?") current_future_latents = None if item.is_infill: include_future_s = item.metadata.get("include_future_s", 0) if include_future_s > 0 and current_future_latents is not None: # truncate the future latents current_total_future_n_tokens = current_future_latents.shape[0] current_total_future_n_tokens -= int(25 * include_future_s) if current_total_future_n_tokens > 0: current_future_latents = current_future_latents[ -current_total_future_n_tokens: ] print( f"UpsampleStub {item_id}: infilling, truncated future latents to {current_future_latents.shape}" ) # if this is empty...we need to clear it if current_history_latents is not None and current_history_latents.shape[0] == 0: print(f"UpsampleStub {item_id}: history latents are empty, setting to None") current_history_latents = None generation_config = DiffusionGenerationConfig( lyrics=item.prompt_text or "", tags=item.metadata["tags"] or "", generation_history_latents=current_history_latents, history_lyrics=history.prompt_lyrics, infill_suffix_latents=current_future_latents, **filtered_config, # pass in experimental configurations ) clip_generation_config_handler.put_diffusion_config(item_id, generation_config) # not currently necessary to input tokens list since its empty, but could be filled in the future with tracer.trace("add_request"): job = self.worker.engine_rpc_client.add_request( Request(item_id, generation_config, tokens=received_tokens) ) time_decoding_completed = None # wait for decoding to complete. while time_decoding_completed is None: # check for new decoded audio new_codes = self.worker.engine_rpc_client.get_generated_codes(job, start=audio_idx) if new_codes is None: time_decoding_completed = time.time() elif isinstance(new_codes, np.ndarray): if audio_idx == 0: statsd.distribution( "diffusion.first_generated_audio_time", time.time() - start_time, tags=dd_tags, ) span_str = serialize_context() DecoderStub().generate.spawn( item.model_dump_json(), stream_key=stream_key, parent_context=span_str ) print( f"UpsampleStub {item_id}: Start the decoder: {item_id}, {time.time() - start_time:.2f}s elapsed." ) statsd.distribution( "diffusion.time_to_start_decoder", time.time() - start_time, tags=dd_tags, ) audio_idx += len(new_codes) # stream the vae latent to codec worker print( f"UpsampleStub {item_id}: streaming vae latent {audio_idx} to codec worker" ) # Stream in chunks of up to 125 tokens latents = new_codes chunks = [] chunk_size = 125 for i in range(0, len(latents), chunk_size): chunk = latents[i : i + chunk_size] # Take up to 125 tokens chunks.append(chunk) with tracer.trace("put_many_chunks"): codec_input_queue.put_many( chunks, partition=partition, partition_ttl=TOKEN_TIMEOUT_DURATION, block=True, timeout=QUEUE_ADD_TIMEOUT, ) if error is not None: break # pull new input tokens from thread queue while True: try: token_batch = token_queue.get_nowait() except queue.Empty: break are_all_codes = all( not isinstance(token, TokenSignalCode) for token in token_batch ) if are_all_codes: # if all are codes we convert them to torch in a batch for speed with tracer.trace("add_tokens_to_job_batch"): tokens_batch = np.stack(token_batch) self.worker.engine_rpc_client.add_token(job, tokens_batch) received_tokens.extend(tokens_batch) if ( len(received_tokens) >= self.worker.default_chunk_size > len(received_tokens) - len(tokens_batch) ): with tracer.trace("received_enough_tokens"): print( f"UpsampleStub {item_id}: received enough ({len(received_tokens)}) tokens, {time.time() - start_time:.2f}s elapsed." ) statsd.distribution( "diffusion.received_enough_tokens", time.time() - start_time, tags=dd_tags, ) else: tokens_to_add = [] set_input_tokens_finished = False job_state = self.worker.engine_rpc_client.get_job_state(job) for token in token_batch: if not isinstance(token, TokenSignalCode): tokens_to_add.append(token) received_tokens.append(token) else: if token == TokenSignalCode.GPT_ERROR: error = "gpt_error" attempts = 0 elif ( token == TokenSignalCode.DECODER_INPUT_TIMED_OUT and not job_state.get("input_tokens_finished") ): error = "decoder_input_timed_out" attempts = 0 elif ( token == TokenSignalCode.UPSAMPLE_INPUT_TIMED_OUT and not job_state.get("input_tokens_finished") ): error = "upsample_input_timed_out" attempts = 0 elif token == TokenSignalCode.STREAM_COMPLETE: self.print_info(item_id, "Received EOS token.") else: raise ValueError("This shouldn't happen.") set_input_tokens_finished = True self.worker.engine_rpc_client.add_token(job, np.array(tokens_to_add)) if set_input_tokens_finished: self.worker.engine_rpc_client.set_input_tokens_finished(job, True) time.sleep(1) if error is None: break elif attempts == 0: self.print_info( item_id, f"Failed to process after few attempts, error: {error}", ) else: self.print_info( item_id, f"Retrying due to {error}, {attempts} attempts left.", ) time.sleep(5) codec_input_queue.put( [TokenSignalCode.STREAM_COMPLETE], partition=stream_key, partition_ttl=TOKEN_TIMEOUT_DURATION, block=False, timeout=QUEUE_ADD_TIMEOUT, ) self.print_info(item_id, "Upsample job complete") if VERBOSE_MESSAGE: max_mem = torch.cuda.max_memory_allocated() / 1e9 self.print_info(item_id, f"Max memory allocated: {max_mem:.2f} GB") loop_end_time = time.time() engine_state = self.worker.engine_rpc_client.get_engine_state() if error is None: statsd.increment("diffusion.complete", 1, tags=dd_tags) statsd.distribution( "diffusion.duration_seconds", loop_end_time - start_time, tags=dd_tags ) statsd.distribution("diffusion.ready_jobs", engine_state["ready_jobs"], tags=dd_tags) statsd.distribution("diffusion.num_jobs", engine_state["num_jobs"], tags=dd_tags) statsd.distribution( "diffusion.min_buffer_size", engine_state["min_buffer_size_s"], tags=dd_tags ) cuda_util = torch.cuda.utilization() print(f"UpsampleStub {item_id}: CUDA utilization: {cuda_util}%") statsd.distribution("diffusion.cuda_utilization", cuda_util, tags=dd_tags) else: raise Exception(f"Failed to generate audio for item_id {item_id}") # save npz for the vae latents latents = self.worker.engine_rpc_client.get_generated_codes(job) print(f"UpsampleStub {item_id}: latents shape: {latents.shape}") assert latents.shape[0] > 0, f"No vae latents generated for item_id {item_id} {job}" print(f"Saving {latents.shape} latents.") parent_from_start_index = None parent_from_end_index = None parent_clip_id = None if item.is_infill: history_arr = history.prompt_audio future_arr = history.future_audio pre_history_arr = history.pre_history_arr post_future_arr = history.post_future_arr parent_from_start_index = ( history_arr.shape[0] if isinstance(history_arr, np.ndarray) else 0 ) + (pre_history_arr.shape[0] if isinstance(pre_history_arr, np.ndarray) else 0) parent_from_end_index = ( future_arr.shape[0] if isinstance(future_arr, np.ndarray) else 0 ) + (post_future_arr.shape[0] if isinstance(post_future_arr, np.ndarray) else 0) parent_clip_id = item.prompt_audio # subsctract the two padding tokens from the start and end include_history_s = item.metadata.get("include_history_s", 0) include_future_s = item.metadata.get("include_future_s", 0) if include_history_s > 0 and history.prompt_audio is not None: history_codes = history.prompt_audio[ -int(SEMANTIC_HZ * include_history_s) :, 0 ].astype(np.int32) parent_from_start_index -= history_codes.shape[0] if include_future_s > 0 and history.future_audio is not None: future_codes = history.future_audio[: int(SEMANTIC_HZ * include_future_s), 0].astype( np.int32 ) parent_from_end_index -= future_codes.shape[0] print( f"{item.id} infilling: parent_from_start_index: {parent_from_start_index}, parent_from_end_index: {parent_from_end_index}" ) job_state = self.worker.engine_rpc_client.get_job_state(job) self.worker._write_vae_latents_npz( item, latents, MODEL_VAE_VERSION.value, APP_NAME, parent_from_start_index=parent_from_start_index, parent_from_end_index=parent_from_end_index, parent_clip_id=parent_clip_id, history_latents=history.history_latents, history_text=history.prompt_lyrics, n_sem_tokens=job_state["processed_tokens"], seed=chosen_seed, ) self.worker.engine_rpc_client.remove_job(job) return item_id except Exception as e: error_message = str(e) self.print_info(item_id, f"Exception occurs: {e}") error_queue.put(error_message, partition=item_id, partition_ttl=TOKEN_TIMEOUT_DURATION) traceback.print_exc() raise e @modal.exit() def cleanup_processes(self): print("UpsampleStub: Cleaning up processes") self.worker.engine_rpc_client.terminate_processes() def fill_local_token_queue_codec(local_token_queue: queue.Queue, partition: str): # Ensure the queue has not expired. # if codec_input_queue.len(partition=partition) == 0: # # Fail the decoding job for this partition. # local_token_queue.put([TokenSignalCode.DECODER_INPUT_TIMED_OUT]) # return # Stream tokens into local queue. for token_batch in codec_input_queue.iterate(partition=partition, item_poll_timeout=60 * 4): local_token_queue.put(token_batch) if ( isinstance(token_batch[-1], TokenSignalCode) and token_batch[-1] == TokenSignalCode.STREAM_COMPLETE ): print(f"DecoderStub: STREAM_COMPLETE signal received for partition {partition}.") return # print(f"Put {len(token_batch)} tokens into queue, {tokens.qsize()} tokens in queue {partition}.") # If we exited the loop without STREAM_COMPLETE, Chirp lagged out and stopped generating tokens. local_token_queue.put([TokenSignalCode.DECODER_INPUT_TIMED_OUT]) class CodecWorker(S3Loader, CodecEngine): def __init__(self): S3Loader.__init__(self) # saves memory with dynamic batch sizes. must be set before loading the model os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" # start the codec engine print(f"CodecWorker: loading model {MODEL_VAE_VERSION} from {MODEL_CONFIG.codec_ckpt_path}") if MODEL_VAE_VERSION == VAEVersion.V_VAE_25_PEAQ_1: from suno_utils.tasks.dac_vae_100hz_peaq import load_model elif MODEL_VAE_VERSION == VAEVersion.V_VAE_25_TUNED_2: from suno_utils.tasks.dac_vae_fixed_25hz import load_model else: raise ValueError(f"Unknown model {MODEL_VAE_VERSION}") local_codec_ckpt_path = chirp_v2._get_model_if_needed( MODEL_CONFIG.codec_ckpt_path, cache_dir=MOUNT_PATH, ) codec_model = load_model(local_codec_ckpt_path, device="cuda") print("CodecWorker: loaded model.") config_do_soft_clip = True if MODEL_VAE_VERSION == "v_vae_25_peaq_2": config_do_soft_clip = False CodecEngine.__init__( self, model=codec_model, is_vae=True, compile=True, do_soft_clip=config_do_soft_clip ) print("CodecWorker: initialized engine.") # TODO: when the GPT worker is spinning up -- they should trigger a multi-warm up call to the decoder # So the decoder scales up before the tokens are generated @app.cls( cpu=12, gpu="A10G", secrets=SECRETS, timeout=600, # TODO: with current issue this is too long scaledown_window=400, memory=32000, retries=modal.Retries( max_retries=1, backoff_coefficient=2.0, initial_delay=5.0, ), min_containers=KEEP_WARM_DECODER[DEPLOYMENT_TYPE], max_containers=NUM_CODEC_WORKER_LIMITS[DEPLOYMENT_TYPE], buffer_containers=1 if DEPLOYMENT_TYPE == "prod" else 0, volumes={MODEL_STORE_VOLUME_DIR: model_store_volume}, ) @modal.concurrent(max_inputs=DECODER_MAX_INPUT) class DecoderStub: @modal.enter() def on_modal_enter(self): import torch print(f"Init with codec ckpt: {MODEL_CONFIG.codec_ckpt_path}") torch.set_num_threads(N_CPU) num_gpus = torch.cuda.device_count() device_name = torch.cuda.get_device_name() print(f"Found {num_gpus} GPUs. Using CUDA device: {device_name}") options = {"statsd_host": "127.0.0.1", "statsd_port": 8125} initialize(**options) self.worker = CodecWorker() self.worker.start() self.modal_f_video_generator = modal.Cls.lookup( f"videos-v2-{'dev' if DEPLOYMENT_TYPE == 'dev' else 'prod'}", "DummyV0Stub", )().write_video def print_info(self, queue_item_id: str, message: str): print(f"DecoderStub ({queue_item_id}): {message}") @modal.method() @distributed_trace("decoder_generate", DECODER_DDOG_SERVICE, env_name=DEPLOYMENT_TYPE) def generate(self, item_json: str, stream_key: Optional[str] = None): start_time = time.time() if VERBOSE_MESSAGE: print_gpu_memory_usage(self.__class__.__name__) torch.cuda.reset_max_memory_allocated() item = QueueItem(**json.loads(item_json)) item_id = item.id device_name = torch.cuda.get_device_name() dd_tags = [ f"model:{MODEL}", f"env:{DEPLOYMENT_TYPE}", f"env_name:{WORKER_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}", f"custom_lyrics:{not bool(item.metadata.get('gpt_description_prompt'))}", ] token_queue = queue.Queue() # 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 loader_thread = threading.Thread( target=fill_local_token_queue_codec, args=(token_queue, partition), daemon=True, ) loader_thread.start() ids = [item_id] first_log_time = None error = None attempts = 2 self.print_info(item_id, "Starting decoder job") try: received_tokens = [] while attempts > 0: error = None attempts -= 1 audios: list[Audio] = [] audio_idx = 0 first_ffmpeg_read_time_mp3 = None first_ffmpeg_read_time_webm = None # TODO: here is the ffmpeg/audio stuff 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) mp3_local_chunk = b"" webm_local_chunk = b"" partially_written_chunk_mp3 = b"" partially_written_chunk_webm = b"" # not currently necessary to input tokens list since its empty, but could be filled in the future with tracer.trace("add_request"): job = self.worker.add_request(CodecRequest(item_id)) # add existing tokens if this is a retry for token in received_tokens: job.add_token(token) while True: # check for new decoded audio if ( audio_idx >= len(job.generated_audios) and len(partially_written_chunk_mp3) == 0 and len(partially_written_chunk_webm) == 0 and job.decoding_completed # we close both together, so checking one is enough and not mp3_proc.stdin.closed ): mp3_proc.stdin.close() webm_proc.stdin.close() # if there is audio to write, and we have space in the ffmpeg buffers, write it to ffmpeg while ( audio_idx < len(job.generated_audios) or len(partially_written_chunk_mp3) > 0 or len(partially_written_chunk_webm) > 0 ): made_progress = False with tracer.trace("write_audio_to_ffmpeg"): # grab new audios until we can fill the partially written chunks # or run out of new audios while ( audio_idx < len(job.generated_audios) and len(partially_written_chunk_mp3) == 0 and len(partially_written_chunk_webm) == 0 ): if audio_idx == 0: statsd.distribution( "decoder.first_generated_audio_time", time.time() - start_time, tags=dd_tags, ) audio = job.generated_audios[audio_idx] audio_idx += 1 audios.append(audio) partially_written_chunk_mp3 = np.ascontiguousarray( audio.array_float.T ).tobytes() partially_written_chunk_webm = partially_written_chunk_mp3 # now we are either out of chunks or have a non-empty partially written chunk if ( len(partially_written_chunk_mp3) == 0 and len(partially_written_chunk_webm) == 0 ): break try: with tracer.trace("write_audio_to_ffmpeg_mp3"): 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: # no space in the buffer pass with tracer.trace("write_audio_to_ffmpeg_webm"): 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: # no space in the buffer pass except BrokenPipeError: self.print_info( item_id, f"Broken pipe, skipping. Wrote {len(audios)} chunks.", ) error = "broken_pipe" break if not made_progress: break mp3_done = False webm_done = False # read audio from ffmpeg if mp3_poll.poll(1): with tracer.trace("read_audio_from_ffmpeg_mp3"): # modal write has 256KiB limit mp3_read_res = mp3_proc.stdout.read(64 * 1024) if mp3_read_res: mp3_local_chunk += mp3_read_res else: # if poll() returns an event but we read 0 bytes, the pipe is closed # at the other end mp3_done = True if first_ffmpeg_read_time_mp3 is None: first_ffmpeg_read_time_mp3 = time.time() statsd.distribution( "decoder.first_ffmpeg_read_time", first_ffmpeg_read_time_mp3 - start_time, tags=dd_tags + ["format:mp3"], ) if first_log_time is None: self.print_info( item_id, f"First audio chunk reading mp3: {time.time() - start_time:.2f}s", ) if webm_poll.poll(1): with tracer.trace("read_audio_from_ffmpeg_webm"): # modal write has 256KiB limit webm_read_res = webm_proc.stdout.read(64 * 1024) if webm_read_res: webm_local_chunk += webm_read_res else: webm_done = True if first_ffmpeg_read_time_webm is None: first_ffmpeg_read_time_webm = time.time() statsd.distribution( "decoder.first_ffmpeg_read_time", first_ffmpeg_read_time_webm - start_time, tags=dd_tags + ["format:webm"], ) if first_log_time is None: self.print_info( item_id, f"First audio chunk reading webm: {time.time() - start_time:.2f}s", ) decoded_duration = sum([audio.duration_s for audio in audios]) should_buffer = first_log_time is None and decoded_duration < 0.6 if ( (len(mp3_local_chunk) > 1024 * 16) and (len(webm_local_chunk) > 1024 * 16) and not should_buffer ): with tracer.trace("write_audio_to_modal_mp3"): if WRITE_TO_REDIS: try: mp3_appchunk_queue.put( mp3_local_chunk, partition=partition, partition_ttl=TOKEN_TIMEOUT_DURATION, block=False, timeout=QUEUE_ADD_TIMEOUT, ) webm_appchunk_queue.put( webm_local_chunk, partition=partition, partition_ttl=TOKEN_TIMEOUT_DURATION, block=False, timeout=QUEUE_ADD_TIMEOUT, ) except Exception: self.print_info(item_id, "Failed to write chunk to Modal.") pass mp3_local_chunk = b"" webm_local_chunk = b"" if first_log_time is None: first_log_time = time.time() elapsed_time = first_log_time - start_time self.print_info( item_id, f"First audio chunk writing: {elapsed_time:.2f}s elapsed since request received.", ) statsd.distribution( "decoder.first_write_time_seconds", elapsed_time, tags=dd_tags, ) statsd.distribution( "decoder.start_time", start_time, tags=dd_tags, ) item.notify_progress( { "id": item.id, "type": "streaming", }, ) events_queue.put( { "type": "gen_streaming", "data": { "elapsed_time": elapsed_time, }, }, partition=item_id, partition_ttl=EVENT_TIMEOUT_DURATION, ) 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"DecoderStub ({item.id}): Start audio write time from gen request: {elapsed_millis // 1000}s" ) if elapsed_millis > 600_000: print( f"DecoderStub ({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, ) if error is not None: break # pull new input tokens from thread queue while True: try: token_batch = token_queue.get_nowait() except queue.Empty: break are_all_codes = all( not isinstance(token, TokenSignalCode) for token in token_batch ) if are_all_codes: # if all are codes we convert them to torch in a batch for speed with tracer.trace("add_tokens_to_job_batch"): tokens_batch = np.stack(token_batch) for token in tokens_batch: job.add_token(token) received_tokens.append(token) else: for token in token_batch: if not isinstance(token, TokenSignalCode): with tracer.trace("add_token_to_job"): if not isinstance(token, np.ndarray): raise ValueError( f"Expected np.ndarray, got {type(token)}" ) # print(f"Adding token to job {item_id}") token = torch.from_numpy(token) job.add_token(token) received_tokens.append(token) else: if token == TokenSignalCode.GPT_ERROR: error = "gpt_error" attempts = 0 elif ( token == TokenSignalCode.DECODER_INPUT_TIMED_OUT and not job.request.input_tokens_finished ): error = "decoder_input_timed_out" attempts = 0 elif token == TokenSignalCode.STREAM_COMPLETE: self.print_info(item_id, "Received EOS token.") else: raise ValueError("This shouldn't happen.") job.request.input_tokens_finished = True if mp3_done and webm_done: break time.sleep(0.1) if error is None: break elif attempts == 0: self.print_info( item_id, f"Failed to process after few attempts, error: {error}", ) else: self.print_info( item_id, f"Retrying due to {error}, {attempts} attempts left.", ) time.sleep(5) mp3_proc.stdin.close() mp3_proc.stdout.close() mp3_proc.terminate() webm_proc.stdin.close() webm_proc.stdout.close() webm_proc.terminate() statsd.distribution( "decoder.step_batch_size", np.mean(self.worker.batch_size_history), tags=dd_tags ) # try to empty the generated audios still? while audio_idx < len(job.generated_audios): audio = job.generated_audios[audio_idx] audio_idx += 1 audios.append(audio) self.print_info(item_id, "Decoder job complete") if VERBOSE_MESSAGE: max_mem = torch.cuda.max_memory_allocated() / 1e9 self.print_info(item_id, f"Max memory allocated: {max_mem:.2f} GB") loop_end_time = time.time() self.worker.remove_job(job.request.id) if len(audios) == 0: self.print_info(item_id, f"No audio generated. Recived tokens: {len(received_tokens)}") error = "decoder_error" if error is None: audio = Audio.concatenate(audios) if audio.duration_s != len(job.tokens) / EMBEDDING_RATE: print( f"Warning: Audio duration {audio.duration_s} does not match expected duration {len(job.tokens) / EMBEDDING_RATE} for item_id {item_id}" ) if WRITE_TO_REDIS: with tracer.trace("write_audio_to_modal"): try: mp3_appchunk_queue.put( mp3_local_chunk, partition=partition, partition_ttl=TOKEN_TIMEOUT_DURATION, block=False, timeout=QUEUE_ADD_TIMEOUT, ) webm_appchunk_queue.put( webm_local_chunk, partition=partition, partition_ttl=TOKEN_TIMEOUT_DURATION, block=False, timeout=QUEUE_ADD_TIMEOUT, ) # write empty bytes to modal to signal end of stream mp3_appchunk_queue.put( b"", partition=partition, partition_ttl=TOKEN_TIMEOUT_DURATION, block=False, timeout=QUEUE_ADD_TIMEOUT, ) webm_appchunk_queue.put( b"", partition=partition, partition_ttl=TOKEN_TIMEOUT_DURATION, block=False, timeout=QUEUE_ADD_TIMEOUT, ) 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"DecoderStub ({item.id}): Finish audio write time from gen request: {elapsed_millis // 1000}s" ) if elapsed_millis > 600_000: print( f"DecoderStub ({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, ) except Exception: self.print_info(item_id, "Failed to write chunk to Modal.") pass if first_log_time is not None: realtime_factor = audio.duration_s / (loop_end_time - first_log_time) self.print_info( item_id, f"Finished writing chunks. Realtime factor: {realtime_factor:.2f}x", ) statsd.increment("decoder.complete", 1, tags=dd_tags) statsd.distribution("decoder.duration_seconds", audio.duration_s, tags=dd_tags) statsd.distribution("decoder.realtime_factor", realtime_factor, tags=dd_tags) with tracer.trace("write_audio_to_s3"): self.worker._write_audio_only(item, audio, gain_adjust=MODEL_GAIN_ADJUST) self.print_info(item_id, "Finished writing to S3.") else: audio = None raise Exception(f"Failed to generate audio for item_id {item_id}") # check if the error queue is empty other_worker_errors = error_queue.get(partition=item_id, block=False) has_errors = other_worker_errors or error if has_errors: self.print_info(item_id, f"Errors: {has_errors}, {other_worker_errors}") item.notify_progress( { "id": item_id, "model": item.model_name, "n_audios": len(ids), "ok": has_errors is None, "gen_duration": time.time() - start_time, "ids": ids, "durations": [round(audio.duration_s, 3) if audio is not None else 0], } ) print(f"Item: {item_id}, finished notifying progress.") video_time = time.time() # Start the video jobs. self.modal_f_video_generator.spawn(item.model_dump_json(), f"image_{item.id}.png") self.print_info( item_id, f"Spawning videos {time.time() - video_time}", ) return item_id except Exception as e: self.print_info(item_id, f"Exception occurs: {e}") finish_time = time.time() item.notify_progress( { "id": item_id, "model": item.model_name, "n_audios": len(ids), "ok": 0, "gen_duration": finish_time - start_time, "error_type": "generation_failure", "error_message": "Failed audio generation.", } ) events_queue.put( { "type": "gen_failed", "data": { "ok": 0, }, }, partition=item_id, partition_ttl=600, ) traceback.print_exc() raise e def test_cancel_job(): import time # Create a job item_id = str(uuid4()) job_data = json.dumps( dict( id=item_id, prompt_text="This is a test job to be cancelled." * 20, metadata={"tags": "test"}, ) ) model = ChirpV2Stub() job = model.generate.spawn(job_data) # Cancel the job appinbound_event_queue.put( Event(item_id, EventType.CANCEL), partition=item_id, partition_ttl=TOKEN_TIMEOUT_DURATION ) time.sleep(1000) def test_stress_test(): import random random.seed(1) NUM_TESTS = 5_000 inputs = [] print(f"Running {NUM_TESTS} generation stress tests") genres = ["edm", "rap", "rock", "pop", "country", "jazz", "classical", "metal", "blues"] for i in range(NUM_TESTS): uid = str(uuid4()) inputs.append( json.dumps( dict( id=uid, prompt_text=f"""Hello world {i}, """ * random.randint(1, 200), metadata={"tags": f"{genres[i % len(genres)]}"}, ), ), ) model = ChirpV2Stub() for input in inputs: model.generate.spawn(input) time.sleep(30) def test_cover_test(): import random random.seed(1) NUM_TESTS = 500 inputs = [] print(f"Running {NUM_TESTS} cover tests") prompt_text_cover = """ [Instrumental intro]\n\n[Verse 1]\nOh Yoni, today's your special day, \nMarking thirty trips around the sun, hooray! \nYet your heart's still a sprightly Pikachu, \nForever young, in all you do. \n\nWith that twinkle in your eye, \nYou're a real adult, oh my! \nBut don't let taxes bring you down, \nWhen there's pokemon cards abound. \n\n[Pre-Chorus]\nYou always play two games, side by side, \nSwitching focus like a tide, \nIn a pixelated world so vast, \nAdventures race by, never too fast. \n\n[melodic break]\n\n[Chorus]\nYour birthday falls on Thanksgiving, \nAnd I'm grateful you're my friend! \nTurkey and pie don't steal your show, \nToday you're a real adult person, you know! \n\n[Verse 2]\nSo open up a pack, and let your spirit soar, \nWill you find a Charizard, or maybe something more? \nLet each shiny card you unveil, \nTell a story, spin a tale. \n\n[Pre Chorus]\nThirty years and still so young, \nWith every beat of life's song sung. \nYour adventure's just begun, \nLike Ash, your journey's never done. \n\n[melodic break]\n\n[Chorus]\nYour birthday falls on Thanksgiving, \nAnd I'm grateful you're my friend! \nTurkey and pie don't steal your show, \nToday you're a real adult person, you know! \n\n[Bridge] \nSo here's to thirty and to thirty more, \nTo catching dreams and opening doors, \nWith every candle blown, make a wish, \nMay your deck be full of everything you cherish! \n\n[Final Chorus]\nYour birthday falls on Thanksgiving, \nAnd I'm grateful you're my friend! \nTurkey and pie don't steal your show, \nToday you're a real adult person, you know! \n\n[Outro]\nHappy birthday Yoni\n(Happy birthday Yoni)\nHappy birthday Yoni\n(Happy birthday Yoni) """ for i in range(NUM_TESTS): uid = str(uuid4()) inputs.append( json.dumps( dict( id=uid, prompt_text=prompt_text_cover, metadata={ "lang": "English", "tags": "hard rock, distorted guitar, metal drums, screamo, upbeat", "task": "cover", "type": "gen", "is_bot": False, "artist_end_s": 120, "feature_flags": 3, "cover_clip_id": "ed1a70e7-150c-4189-933b-2eda8cfd6dd1", "edited_clip_id": "ed1a70e7-150c-4189-933b-2eda8cfd6dd1", }, ), ), ) model = ChirpV2Stub() from suno_utils.worker.loader import S3Loader item = QueueItem(**json.loads(inputs[0])) cover_audio = S3Loader()._load_special_audio_prompt( item=item, model_version="4.0.0.0", is_cover=True ) history = HistoryPrompt(cover_audio=cover_audio) for input in inputs: model.generate.spawn(input, history) time.sleep(30) def test_artist_test(): import random random.seed(1) NUM_TESTS = 500 inputs = [] print(f"Running {NUM_TESTS} persona tests") prompt_text_artist = """ [Verse 1]\nWe don't know if you got a dog\nWe don't even know if you've even got a blog\nWe don't know if you've got a love life\nWe don't even know if you're somebody's wife\n\n[Verse 2]\nWe have never seen you cover up a yawn\nDid you wait for someone after prom?\nDid you get a degree? Do you work with your brain?\nAre you workin' in the government? Ain't that a pain?\n\n[Chorus]\nHappy birthday\nTo someone who's probably somethin'\n\n[Verse 3]\nAre you neat or are you cluttered?\nAre you quiet or do you mutter under your breath?\nDo you clear your throat too much\nOr do you have ASMR with a whisperin' touch?\n\n[Verse 4]\nDid you go on a boozy cruise?\nDid you see the blues played by Baby Clues?\nDo you know how to lose someone?\nWell if you don't that's prob'ly a good one\n\n[Bridge]\nWe don't have a picture of your hair\nWe don't know if you are in love\nDo you love to breathe the air? (Do you love to breathe the air?)\nIf you are in love (If you are in love)\nI do too when push comes to shove (When push comes to shove) """ for i in range(NUM_TESTS): uid = str(uuid4()) inputs.append( json.dumps( dict( id=uid, prompt_text=prompt_text_artist, metadata={ "lang": "English", "tags": "alternative rock, playful, whimsical, folk pop, vocal harmonies, humorous, male vocalist, upbeat, electric guitar, piano", "task": "artist_consistency", "type": "gen", "is_bot": False, "artist_end_s": 120, "feature_flags": 3, "artist_clip_id": "ed1a70e7-150c-4189-933b-2eda8cfd6dd1", "edited_clip_id": "ed1a70e7-150c-4189-933b-2eda8cfd6dd1", "persona_id": "de9104d0-dd25-43a6-a2da-10e4f78b45d6", }, ), ), ) model = ChirpV2Stub() from suno_utils.worker.loader import S3Loader item = QueueItem(**json.loads(inputs[0])) artist_audio = S3Loader()._load_special_audio_prompt( item=item, model_version="4.0.0.0", is_artist=True ) history = HistoryPrompt(artist_audio=artist_audio) for input in inputs: model.generate.spawn(input, history) time.sleep(30) @app.local_entrypoint() def main(): # test_artist_test() # only for 30b # test_cover_test() # only for 30b test_stress_test() time.sleep(3600 * 2)