import os os.environ["CUDA_VISIBLE_DEVICES"] = "1" import json import glob import torch import IPython import torchaudio import numpy as np from stable_audio_tools.inference.generation import ( upsample_diffusion, upsample_diffusion_from_codes, denoising_diffusion_from_codes, upsample_diffusion_from_semantic, ) from stable_audio_tools.interface.gradio import load_model from stable_audio_tools.models.utils import apply_normalization from suno_utils.utils.s3 import read_from_s3 from suno_utils.models.dac.nn.quantize_2 import ResidualVectorQuantize #from suno_utils.tasks.dac_2c_12cb import load_model as load_vae_model # VAE from suno_utils.tasks.dac_vae_peaq import ( preload_models as preload_vae_models, load_model as load_vae_model, encode as vae_encode, decode as vae_decode, ) # MERT from suno_utils.tasks.mert_25 import ( preload_models as preload_semantic_models, encode as semantic_encode, ) _ = preload_semantic_models( checkpoint_filepath="s3://suno-data/georg/models/semantic/mert_25.pt", centroids_filepath="s3://suno-data/georg/models/semantic/mert_25_2x4k.npy", device="cuda", ) def decode_vq(model, codes, n_quantizers): z_q = 0 for i, quantizer in enumerate(model.quantizer.quantizers[:n_quantizers]): _z_q = quantizer.embed_code(codes[:, :, i]).transpose(1, 2) _z_q = quantizer.out_proj(_z_q) z_q += _z_q.transpose(1, 2) return z_q ckpt_path = "/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_s1_vae_25hz_1b_epoch=0-step=90000.ckpt" config_dir = "/home/christian/code/neon/stable-audio-tools/stable_audio_tools/configs/model_configs/txt2audio" model_type = "semantic" model_config_path = os.path.join( config_dir, "stable_audio_2_0_semantic_48khz_lg_vae.json" ) # load model from checkpoint if model_config_path is not None: # Load config from json file with open(model_config_path) as f: model_config = json.load(f) else: model_config = None print(model_config) # model_config["model"]["diffusion"]["config"]["use_checkpointing"] = True device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model, model_config = load_model( model_config, ckpt_path, # pretrained_name=pretrained_name, # pretransform_ckpt_path=pretransform_ckpt_path, # model_half=model_half, device="cuda", ) # load VAE vae_model = load_vae_model( checkpoint_filepath="s3://suno-data/christian/mw_vae_peaq_128_fix.pth", device="cuda" ) audio_paths = [ "/home/christian/audio/reference-audio-wav/09 Sounds Like Hallelujah.wav", "/home/christian/audio/reference-audio-wav/02 Freddie Freeloader.wav", "/home/christian/audio/reference-audio-wav/01 No Son Of Mine.wav", "/home/christian/audio/reference-audio-wav/02 Dreams.wav", "/home/christian/audio/reference-audio-wav/01 J.S. Bach Suite No.1, S.1007, G major - I. Prelude.wav", "/home/christian/audio/reference-audio-wav/04 Fuckwithmeyouknowigotit.wav", "/home/christian/audio/reference-audio-wav/03 Always Be.wav", ] audio_paths = glob.glob("/home/christian/audio/reference-audio-wav/*.wav")[:10] file_ext = ".wav" for audio_path in audio_paths: print(audio_path) audio_input, sr = torchaudio.load(audio_path) if audio_input.shape[0] == 1: audio_input = audio_input.repeat(2, 1) if sr != 48000: audio_input = torchaudio.functional.resample(audio_input, sr, 48000) start_frame = audio_input.shape[-1] // 2 end_frame = start_frame + int(10 * 48000) + 1024 audio_input = audio_input[:, start_frame:end_frame] audio_input_24khz = torchaudio.functional.resample(audio_input, 48000, 24000) semantic_codes = semantic_encode([audio_input_24khz.mean(dim=0, keepdim=True)]) semantic_codes = torch.from_numpy(semantic_codes[0][:,0]).long().cuda() print("semantic_codes", semantic_codes.shape) target_cycled_audio = vae_model(audio_input.unsqueeze(0).cuda())["audio"].detach().cpu() target_cycled_audio /= target_cycled_audio.abs().max() with torch.no_grad(): upsampled_latents = upsample_diffusion_from_semantic( model, semantic_codes, steps=1000, cfg_scale=1.0, sample_size=semantic_codes.shape[0], sample_rate=48000, ) pred_zq = upsampled_latents#.squeeze()#.permute(1, 0) print("pred_zq", pred_zq.shape) pred_audio = vae_model.decode(pred_zq)[0].detach().cpu() pred_audio /= pred_audio.abs().max() IPython.display.display(IPython.display.Audio(data=audio_input.cpu().squeeze().numpy(), rate=48000)) IPython.display.display(IPython.display.Audio(data=target_cycled_audio.cpu().squeeze().numpy(), rate=48000)) IPython.display.display(IPython.display.Audio(data=pred_audio.cpu().squeeze().numpy(), rate=48000)) # ----- %load_ext autoreload %autoreload 2 # setup GPT from suno_utils.gpt.chirp_v2_5 import ( GenerationConfig, decode_stream, preload_models, prep_gconf, codec_decode_stream_to_full_audio, generate, ) from suno_utils.gpt.generation import clean_models, _load_model def load_gpt_model( ckpt_path=None, tokenizer_path=None, use_gpu=True, force_reload=False, use_tp=False, ): if torch.cuda.device_count() == 0 or not use_gpu: device = "cpu" else: device = "cuda" model_key = "main_model" if ckpt_path is None: raise ValueError("model not initialized, need checkpoint path. maybe run `preload_models`?") clean_models(model_key=model_key) model, tokenizer = _load_model(ckpt_path, tokenizer_path, device, use_tp=use_tp) return model #gpt_ckpt_path = "/app/suno/data/dpo/models/model_13b_full.pt" # model before fine-tuning #gpt_ckpt_path = "/app/suno/checkpoints/2024-07-12_13-28-28/last_ckpt.pt" # 2b trained only on semantic #gpt_ckpt_path = "/app/suno/checkpoints/2024-07-13_20-42-16/last_ckpt.pt" gpt_ckpt_path = "/app/suno/checkpoints/2024-07-14_21-24-26/last_ckpt.pt" tokenizer_path = "/app/suno/data/dpo/models/tokenizer_60k.json" preload_models( gpt_ckpt_path=gpt_ckpt_path, tokenizer_path=tokenizer_path, load_gpt=True, load_semantic=False, load_codec_device="cuda", ) text_tags_list = [ # "indie-pop, energetic, rock, psychedelic", # "jazz, female vocal, smooth", # "r&b, soul, funk", "pop, high energy, female vocal", "rock, high energy, male vocal", "indie rock, female vocal, psychedelic", # "moody, blues, soulful", # "classical, cinematic, female vocal", # "metal, hardcore, dark", "male vocal, bluegrass, guitar", # "female vocal, voice, singing, guitar, folk, jazz, r&b", ] text_a = """ Almost Heaven, West Virginia Blue Ridge Mountains, Shenandoah River Life is old there, older than the trees Younger than the mountains, growing like a breeze Country roads, take me home To the place I belong West Virginia, mountain mama Take me home, country roads [outro] """ text_b = """ [Verse 1] If I had the superpower to reverse Time back to a spring century before When everyone was hustling to preserve Unless I wielded powers to reverse [Instrumental Solo] """ random_seed = 13 text_tags = "bluegrass, vocal, voice, singing, female singer" #text_tags = "kpop, vocals, female, beat, soul, energetic" #text_tags = "rock, high energy, male vocal" general_config = dict( #cfg_coef=2.5, #min_eos_p=0.1, #eos_pad_duration_s=0, #cfg_coef_tags=1.0, #cfg_coef_tags_max_steps=None, # collect the data for now #n_repeat_tags=3, #use_whisper=False, text_start_control_tags="{start}", # text_end_control_tags="{end}", random_seed=random_seed, n_batch=1, ) gconfig = GenerationConfig( text=text_a, text_tags=text_tags, max_gen_duration_s=60, #temp_semantic=0.9, #top_k_semantic=None, #top_p_semantic=None, **general_config, ) semantic_codes = generate(gconfig) semantic_codes = torch.from_numpy(semantic_codes[0]).squeeze() semantic_codes = semantic_codes.cuda() # evaluate on val example print(semantic_codes.shape) n_chunks = semantic_codes.shape[0] // 250 pred_zqs = [] num_steps = 250 loudnorm = True for chunk_idx in range(n_chunks): with torch.no_grad(): upsampled_latents = upsample_diffusion_from_semantic( model, semantic_codes[chunk_idx * 250 : (chunk_idx + 1) * 250], steps=num_steps, cfg_scale=1.0, sample_size=250, sample_rate=48000, ) pred_zq = upsampled_latents print("pred_zq", pred_zq.shape) pred_zqs.append(pred_zq) pred_zq = torch.cat(pred_zqs, dim=-1) print(pred_zq.shape) pred_audio = vae_model.decode(pred_zq)[0].detach().cpu() #IPython.display.display(IPython.display.Audio(data=target_audio.cpu().squeeze().numpy(), rate=48000)) IPython.display.display(IPython.display.Audio(data=pred_audio.cpu().squeeze().numpy(), rate=48000))