"""Generating testing samples for model comparisons. aws s3 cp audio_files.json s3://suno-data-uploads/studio/uploads/neon230817/audio_files.json aws cloudfront create-invalidation --distribution-id E3ROE628KJYDYL --paths "/neon230817/audio_files.json" """ import os import torch from suno_utils.audio import Audio from suno_utils.gpt.chirp_v2_5 import ( GenerationConfig, decode_stream, preload_models, prep_gconf, ) from suno_utils.gpt.generation_engine import ( align_codes, make_request, unshift_arrays_v2, ) from suno_utils.gpt.engine import Engine from tqdm import tqdm import numpy as np from suno_utils.worker.settings import s3_client lyrics_dune = """ feel the breeze that's calling me golden skies above the sea endless summer chasing dreams yeah you and i we werere wild and free we chillin out just riding the vibe under the stars you by my side it feels so right the night is ours you and me are shooting stars we'll never fade the nights into days a summer will'll always stay on the sand dunes we run we fly burning bright beneath the moonlit sky hold me close don't let this go drifting closer side by side feel us fading with the tide every glance every touch makes me feel like it's enough we were chillining out just riding the vibe under the stars you by my side feels so right the night is ours you and me i' shooting stars we'll never fade the nights into days our summer will always stay on the sand dunes we run we fly building bright beneath the moonlit sky every come don let this go the d on the sand dunes we'll be writing our story every moment is a memory of glory on the sand dun run we fly when the winds blow in we sway we sway to the shifting sands we sway we sway we'll never fade the nights and the days a summer will always stay the moonit sky hearts don't let this go on the sand dun well be writing our story every moment is a memory of glory on the sand dun above we fly on dun """ # dict of titie; lyrics; tags test_exmaples = { "ts_pop": { "text": lyrics_dune, "tags": "female, k-pop", }, } if __name__ == "__main__": print(f"working with GPU:{os.environ['CUDA_VISIBLE_DEVICES']} ") # gpt_ckpt_path = "/app/suno/data/dpo/models/model_30b_ft_t3.pt" gpt_ckpt_path = "/app/suno/checkpoints/2024-09-23_22-47-31/last_ckpt_infer.pt" gpt_output_dir = "c_30b_t4_v12_beta5_sft01" # gpt_ckpt_path = "/app/suno/checkpoints/2024-09-12_03-47-19/last_ckpt_infer.pt" # gpt_output_dir = "c_7b_base_t1_v18" gpt_output_path = os.path.join("/home/tony/Work/gpt_samples/", gpt_output_dir) os.makedirs(gpt_output_path, exist_ok=True) print(f"Sending outputs to {gpt_output_path}") MAX_STREAMS = 4 if "30b" not in gpt_output_path else 3 MAX_SEQ = ( 32 if "30b" not in gpt_output_path else 12 ) # 13b 8 head max is 40 and OOM... # Import models and setup preload_models(load_gpt=False, load_semantic=False, load_codec_device="cuda") engine = Engine( gpt_ckpt_path, "/app/suno/data/dpo/models/tokenizer_60k.json", max_sequences=MAX_SEQ, compile=False, ) test_npz = np.load("audios/4bd01278-b5a7-491d-8fa7-58cbefa09873.npz") in_cover_arr = test_npz["v3.0_raw"] in_cover_arr = in_cover_arr[:3000, :] def generate_audio_with_engine(text, text_tags, title, random_seed=42): general_config = dict( cfg_coef=1.0, # no text cfg for dpo stream min_eos_p=0.1, eos_pad_duration_s=0, # cfg_coef_tags=0.0, cfg_coef_tags_max_steps=None, # collect the data for now cfg_coef_tags=2, cfg_coef_neg_tags=-1, text_neg_tags="repetitive, loop" if "30b" in gpt_output_path else "repetitive, loop, noisy, distorted", n_repeat_tags=1, use_whisper=False, text_start_control_tags="{start:0} " if "30b" in gpt_output_path else "{start} ", # text_end_control_tags="{end}", min_text_offset=0, # this is required for 30b random_seed=random_seed, n_batch=1, cover_arr=in_cover_arr[:3000, :], ) if "short" in gpt_output_path: general_config["text_start_control_tags"] = "{start:0;vocals:start}" general_config["text_end_control_tags"] = "{end}" cfg = GenerationConfig( text=text, text_tags=text_tags, max_gen_duration_s=125, **general_config ) preped_cfg = prep_gconf(cfg) model_conf = engine.model.config requests = [] for i in range(1, MAX_STREAMS + 1): request = make_request( f"{i}", preped_cfg, engine.model.config, engine.tokenizer ) requests.append(request) jobs = engine.run_request(requests, tqdm_enabled=True) for i, job in enumerate(jobs): stream = engine.token_generator(job) audios = [] codes = [] raw_codes = [] tensor_codes = [] for code in stream: raw_codes.append(code) torch_codes = torch.stack(raw_codes).cpu() # print(torch_codes.shape) unshifted_codes = unshift_arrays_v2(torch_codes, model_conf) # print(unshifted_codes.shape) for code in align_codes(tqdm(raw_codes), model_conf): codes.append(code.numpy().astype(np.int16)) tensor_codes.append(code) for audio in decode_stream(tensor_codes): audios.append(audio[0]) with open(os.path.join(gpt_output_path, f"{title}_{i}.npz"), "wb") as f: np.savez(f, codes=unshifted_codes) audio_continued = Audio.concatenate(audios) # audio_continued.play() # save the file to mp3 mp3_output_path = os.path.join( gpt_output_path, f"{title}_{i}_seed{random_seed}.mp3" ) audio_continued.to_hq_mp3(mp3_output_path) s3_client.upload_file( mp3_output_path, "suno-data-uploads", f"studio/uploads/neon230817/{gpt_output_dir}/{title}_{i}_seed{random_seed}.mp3", ExtraArgs={ "ContentType": "audio/mp3", }, ) return print("Start generating samples") for random_seed in range(42, 46): for title, example in tqdm(test_exmaples.items()): generate_audio_with_engine( example["text"], example["tags"], title=title, random_seed=random_seed ) print("Finished generating samples. Have a good day!")