import datetime import json import os import pathlib import queue import select import subprocess import threading import time from typing import Any from uuid import uuid4 import modal import numpy as np import redis from openai import OpenAI from torchaudio.pipelines import HDEMUCS_HIGH_MUSDB_PLUS from suno_utils.audio import Audio from suno_utils.worker.generate_song_lyrics import ( ModerationError, ModerationSuccess, generate_song_lyrics_with_genre_tags, moderate_user_inputs, ) from suno_utils.worker.schema import QueueItem from suno_utils.worker.loader import S3Loader from suno_utils.worker.modal_base import MODAL_MOUNTS from suno_utils.worker.utils import recursive_ls_dir base_image = ( modal.Image.debian_slim() .apt_install("curl", "ffmpeg", "sox", "unzip", "libsox-fmt-mp3") .run_commands( [ 'curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip"', "unzip -q awscliv2.zip", "./aws/install", ] ) .pip_install( "torch==2.1.0.dev20230831+cu118", "torchaudio==2.2.0.dev20230831+cu118", index_url="https://download.pytorch.org/whl/nightly/cu118", ) .pip_install_private_repos( "github.com/suno-ai/glockenspiel.git@f05e2f251#subdirectory=descript-audio-codec&egg=descript-audio-codec", git_user="mcamac", secrets=[modal.Secret.from_name("victor-modal-github-token")], ) .pip_install( "boto3", "transformers", "tokenizers", "encodec", "ctc_segmentation", "psutil", "redis", "gradio", "pydantic", "nnAudio", ) .pip_install_from_pyproject( str(pathlib.Path(__file__).parent.parent.parent / "pyproject.toml"), ) .run_commands( "FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE pip install flash-attn --no-build-isolation", ) ) MOUNT_PATH = "/suno/models" FT_GPT_CKPT_PATH = "s3://suno-data/georg/checkpoints/chirp_v2_5/7b_ft_fix.pt" concurrency = 2 POLL_TIMEOUT_MS = 100 IS_7B = "7b" in FT_GPT_CKPT_PATH if IS_7B: from suno_utils.gpt import chirp_v2_5 as chirp_v2 else: from suno_utils.gpt import chirp_v2 class ChirpV2Worker(S3Loader, threading.Thread): gpu_id: int def __init__(self, gpu_id, in_q: queue.Queue, out_qs: dict[queue.Queue], timeout_s=5): S3Loader.__init__(self) threading.Thread.__init__(self) self.gpu_id = gpu_id self.in_q = in_q self.out_qs = out_qs self.timeout_s = timeout_s self.preload() def run(self): """ Main thread loop. Waits for queue to fill or timeout, then processes the batch. """ print("Running main thread") first_request_time = None while True: cur_time = datetime.datetime.now() timeout = first_request_time is not None and ( cur_time - first_request_time > datetime.timedelta(seconds=self.timeout_s) ) # wait for queue to fill, then process the batch if self.in_q.full() or timeout: if timeout: print("Timeout triggered, in queue not full") print(f"Processing batch with size {self.in_q.qsize()}") batch = [] while not self.in_q.empty(): batch.append(self.in_q.get(block=False)) self.process_batch(batch) first_request_time = None else: if first_request_time is None and not self.in_q.empty(): first_request_time = datetime.datetime.now() time.sleep(0.1) def preload(self): start_time = time.time() chirp_v2.preload_models( gpt_ckpt_path=FT_GPT_CKPT_PATH, cache_dir=MOUNT_PATH, local_whisper="/suno/models/whisper" ) finish_time = time.time() print(f"Preloading took {finish_time - start_time}s") @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=FT_GPT_CKPT_PATH, fetch_only=True, cache_dir=dir_path) recursive_ls_dir(dir_path) HDEMUCS_HIGH_MUSDB_PLUS.get_model() import whisper whisper.load_model("small") whisper.load_model("small.en") def process_batch(self, batch: list[QueueItem]) -> list[tuple[list[Audio], list[Any]]]: # build a config arr_text = [] arr_text_tags = [] arr_text_neg_tags = [] arr_cfg_coef = [] for item in batch: text = item.prompt_text or "" text_tags = item.metadata.get("tags", None) if text_tags == "random": text_tags = None # detected_lang = chirp_v2._get_text_lang(text) cfg_coef = 1.25 # cfg_coef = float(options.get("lyrics_strength", 1.2)) # condition on non-english, higher cfg 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 # if text_tags and ("rap" in text_tags or "hip-hop" in text_tags): # cfg_coef = 1.4 text_neg_tags = None # if "instrumental" not in (text_tags or "") and len(text.split()) > 5: # text_neg_tags = "instrumental noise" arr_text.append(text) arr_text_tags.append(text_tags) arr_text_neg_tags.append(text_neg_tags) arr_cfg_coef.append(cfg_coef) cfg = chirp_v2.GenerationConfig( text=arr_text, text_tags=arr_text_tags, text_neg_tags=arr_text_neg_tags, n_batch=len(batch), max_gen_duration_s=80, cfg_coef_neg_tags=0, # cfg_coef=arr_cfg_coef, # cfg_coef_tags_max_steps=150, # cfg_coef_tags=3.0, # 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, # rep_penality=0, stream=True, ) for audios in chirp_v2.generate_audio_stream(cfg, return_raw_arrays=True): if isinstance(audios, dict): # finished, out dict with raw arrays out = audios for i, item in enumerate(batch): print("finished batch", item.id) raw_arrays = out["raw_arrays"][i] self.out_qs[item.id].put((None, raw_arrays)) break for i, item in enumerate(batch): if not audios[i].is_zero(): self.out_qs[item.id].put((audios[i], None)) aws_secret = modal.Secret.from_name("studio-aws") SECRETS = [ aws_secret, modal.Secret.from_dict( { "SUNO_ASSETS_PATH": "/suno/models/assets", "XDG_CACHE_HOME": "/suno/models/", } ), modal.Secret.from_name("openai-secret"), ] def download_model_wrapper_a4c(): ChirpV2Worker.download_models() image = base_image.run_function(download_model_wrapper_a4c, secrets=SECRETS) ENV_NAME = "vt_chirpv3_batching" STUB_NAME = f"streaming-{ENV_NAME}" stub = modal.Stub(STUB_NAME, image=image) CONCURRENCY_LIMITS = { "priority": 80, "staging": 350, "vt_chirpv3_batching": 5, } @stub.cls( cpu=2.0, # memory=16384, gpu=modal.gpu.A10G(count=1), secrets=SECRETS + [modal.Secret.from_name("redis-test")], timeout=200, container_idle_timeout=400, mounts=MODAL_MOUNTS, memory=60000, retries=modal.Retries( max_retries=3, backoff_coefficient=2.0, initial_delay=10.0, ), keep_warm=1, concurrency_limit=CONCURRENCY_LIMITS.get(ENV_NAME, 80), allow_concurrent_inputs=concurrency, ) class ChirpV2Stub: def __enter__(self): import torch num_gpus = torch.cuda.device_count() print(f"Found {num_gpus} GPUs.") recursive_ls_dir("/suno/models") self.in_q = queue.Queue(maxsize=concurrency) self.out_qs = {} self.worker = ChirpV2Worker(0, self.in_q, self.out_qs) self.worker.start() self.openai_client = OpenAI() self.redis = redis.Redis.from_url(os.environ.get("REDIS_TEST_URL")) @modal.method() def generate(self, queue_item: str): start_time = time.time() print(queue_item) start_time = time.time() item = QueueItem(**json.loads(queue_item)) item_id = item.id print(item.metadata) is_square = True ids = [] if item.prompt_audio is not None: raise NotImplementedError("prompt_audio not supported with batching") if user_prompt := item.metadata.get("gpt_description_prompt"): try: title, lyrics, genre_tags = generate_song_lyrics_with_genre_tags( self.openai_client, chirp_v2.text_lang_model, user_prompt ) except ModerationError as e: import traceback self.worker.notify_finish( item, { "id": item.id, "type": "error", "error_message": str(e), }, ) traceback.print_exc() return item.prompt_text = lyrics item.metadata["tags"] = " ".join(genre_tags) print("generated", item.metadata["tags"], lyrics) self.worker.notify_progress( item, { "id": item.id, "title": title, "type": "lyrics", "text": item.prompt_text, "tags": item.metadata["tags"], }, ) else: tags = item.metadata.get("tags") or "" moderation_result = moderate_user_inputs(self.openai_client, item.prompt_text or "", tags) print(moderation_result) if not isinstance(moderation_result, ModerationSuccess): self.worker.notify_finish( item, { "id": item.id, "type": "error", "error_type": "moderation_failure", "error_message": moderation_result.err_msg, }, ) return f = modal.Function.lookup("stable-diffusion-xl-beta-3", "StableDiffusion.generate_image") fn_call = f.spawn( item.prompt_text or "", item_id, num_images=1, tags=item.metadata.get("tags", None), ) try: ids.append(item_id) self.out_qs[item_id] = queue.Queue() print(f"Adding job to queue: {item_id}, {time.time() - start_time:.2f}s elapsed.") self.in_q.put(item) items = [] error = None with subprocess.Popen( "ffmpeg -f f32le -acodec pcm_f32le -ar 48000 -ac 2 -i pipe: -f mp3 pipe:".split(), stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, ) as proc: os.set_blocking(proc.stdout.fileno(), False) poll = select.poll() poll.register(proc.stdout, select.POLLIN) local_chunk = b"" items.append(item.json()) # wait for results has_logged = False audios = [] while True: audio, raw_arrays = self.out_qs[item_id].get(block=True) if audio is None: # finished print(f"Finished {item.id}, writing to npz") break else: audios.append(audio) try: proc.stdin.write(np.ascontiguousarray(audio.array_float.T).tobytes()) except BrokenPipeError: print(f"Broken pipe, skipping {item.id}. Wrote {len(audios)} chunks.") error = "broken_pipe" break while poll.poll(POLL_TIMEOUT_MS): raw_bytes = proc.stdout.read() local_chunk += raw_bytes if len(local_chunk) > 1024 * 16: if not has_logged: has_logged = True cur_time = time.time() print( "Writing to redis: " f"{item.id}. {cur_time - start_time:.2f}s elapsed " "since request recieved." ) self.redis.rpush(item.id, local_chunk) self.redis.expire(item.id, 60 * 10) local_chunk = b"" self.worker._write_npz(item, raw_arrays) # flush remaining data if local_chunk: self.redis.rpush(item.id, local_chunk) # write empty bytes to redis to signal end of stream self.redis.rpush(item.id, b"") proc.stdin.close() proc.stdout.close() proc.terminate() audio = Audio.concatenate(audios) print(f"Job complete: {item_id}") # clear out queue self.out_qs.pop(item.id) self.worker._write_audio_only(item, audio) image_url = fn_call.get(timeout=None) print("image located", image_url) self.worker.notify_finish( item, { "id": item_id, "model": "chirp-v3", "n_audios": len(ids), "ok": error is None, "gen_duration": time.time() - start_time, "ids": ids, "durations": [audio.duration_s], }, queue_name="results:q", ) video_time = time.time() f = modal.Function.lookup("dummy-v2-staging", "DummyV0Stub.write_video") # Start the video jobs. for item in items: f.spawn(item, image_url.split("/")[-1], is_square) print("Spawning videos", time.time() - video_time) return item_id except Exception as e: import traceback finish_time = time.time() self.worker.notify_finish( item, { "id": item_id, "model": "chirp-v3", "n_audios": len(ids), "ok": 0, "gen_duration": finish_time - start_time, }, queue_name="results:q", ) traceback.print_exc() raise e @stub.local_entrypoint() def main(): inputs = [] genres = ["edm", "rap", "rock", "pop", "country", "jazz", "classical", "metal", "blues"] for i in range(10): uid = str(uuid4()) inputs.append( json.dumps( dict( id=uid, prompt_text=f"""Hello world {i}, """ * 20, metadata={"tags": f"{genres[i % len(genres)]}"}, ), ), ) model = ChirpV2Stub() print(uid) for input in inputs: model.generate.spawn(input) time.sleep(200)