Spaces:
Running
Running
| import librosa | |
| import numpy as np | |
| import torch | |
| from tqdm import tqdm | |
| from tools import rms_normalize, decode_stft, depad_STFT | |
| from model.DiffSynthSampler import DiffSynthSampler | |
| def sample_pipeline_STFT(device, uNet, VAE, MMM, CLAP_tokenizer, | |
| positive_prompts, negative_prompts, batchsize, sample_steps, CFG, seed=None, | |
| freq_resolution=512, time_resolution=256, channels=4, VAE_scale=4, timesteps=1000, noise_strategy="repeat", sampler="ddim", return_latent=True): | |
| "Sample a fix-length audio using a diffusion model, including 'ISTFT+' post-processing." | |
| height = int(freq_resolution/VAE_scale) | |
| width = int(time_resolution/VAE_scale) | |
| VAE_encoder, VAE_quantizer, VAE_decoder = VAE._encoder, VAE._vq_vae, VAE._decoder | |
| text2sound_embedding = \ | |
| MMM.get_text_features(**CLAP_tokenizer([positive_prompts], padding=True, return_tensors="pt"))[0].to(device) | |
| negative_condition = \ | |
| MMM.get_text_features(**CLAP_tokenizer([negative_prompts], padding=True, return_tensors="pt"))[ | |
| 0].to(device) | |
| mySampler = DiffSynthSampler(timesteps, height=height, channels=channels, noise_strategy=noise_strategy, mute=True) | |
| mySampler.activate_classifier_free_guidance(CFG, negative_condition) | |
| mySampler.respace(list(np.linspace(0, timesteps - 1, sample_steps, dtype=np.int32))) | |
| condition = text2sound_embedding.repeat(batchsize, 1) | |
| latent_representations, initial_noise = \ | |
| mySampler.sample(model=uNet, shape=(batchsize, channels, height, width), seed=seed, | |
| return_tensor=True, condition=condition, sampler=sampler) | |
| latent_representations = latent_representations[-1] | |
| quantized_latent_representations, _, (_, _, _) = VAE_quantizer(latent_representations) | |
| if return_latent: | |
| return quantized_latent_representations.detach() | |
| reconstruction_batch = VAE_decoder(quantized_latent_representations).to("cpu").detach().numpy() | |
| rec_signals = [] | |
| for index, STFT in enumerate(reconstruction_batch): | |
| padded_D_rec = decode_stft(STFT) | |
| D_rec = depad_STFT(padded_D_rec) | |
| # get_audio | |
| rec_signal = librosa.istft(D_rec, hop_length=256, win_length=1024) | |
| rec_signals.append(rms_normalize(rec_signal)) | |
| return quantized_latent_representations.detach(), reconstruction_batch, rec_signals | |
| def sample_pipeline_GAN_STFT(device, gan_generator, VAE, MMM, CLAP_tokenizer, | |
| positive_prompts, negative_prompts, batchsize, sample_steps, CFG, seed=None, | |
| freq_resolution=512, time_resolution=256, channels=4, VAE_scale=4, timesteps=1000, noise_strategy="repeat", sampler="ddim", return_latent=True): | |
| "Sample fix-length audio using a GAN, including 'ISTFT+' post-processing." | |
| height = int(freq_resolution/VAE_scale) | |
| width = int(time_resolution/VAE_scale) | |
| VAE_encoder, VAE_quantizer, VAE_decoder = VAE._encoder, VAE._vq_vae, VAE._decoder | |
| text2sound_embedding = \ | |
| MMM.get_text_features(**CLAP_tokenizer([positive_prompts], padding=True, return_tensors="pt"))[0].to(device) | |
| condition = text2sound_embedding.repeat(batchsize, 1) | |
| noise = torch.randn(batchsize, channels, height, width).to(device) | |
| latent_representations = gan_generator(noise, condition) | |
| quantized_latent_representations, _, (_, _, _) = VAE_quantizer(latent_representations) | |
| if return_latent: | |
| return quantized_latent_representations.detach() | |
| reconstruction_batch = VAE_decoder(quantized_latent_representations).to("cpu").detach().numpy() | |
| rec_signals = [] | |
| for index, STFT in enumerate(reconstruction_batch): | |
| padded_D_rec = decode_stft(STFT) | |
| D_rec = depad_STFT(padded_D_rec) | |
| # get_audio | |
| rec_signal = librosa.istft(D_rec, hop_length=256, win_length=1024) | |
| rec_signals.append(rms_normalize(rec_signal)) | |
| return quantized_latent_representations.detach(), reconstruction_batch, rec_signals | |
| def generate_audios_with_diffuSynth_sample(device, uNet, VAE, MMM, CLAP_tokenizer, num_batches, positive_prompts, negative_prompts="", CFG=6, sample_steps=10): | |
| "Sample audios using a diffusion model, including 'ISTFT+' post-processing." | |
| diffuSynth_signals = [] | |
| for _ in tqdm(range(num_batches)): | |
| _, _, signals = sample_pipeline_STFT(device, uNet, VAE, MMM, CLAP_tokenizer, | |
| positive_prompts=positive_prompts, negative_prompts=negative_prompts, | |
| batchsize=8, sample_steps=sample_steps, CFG=CFG, seed=None, return_latent=False) | |
| diffuSynth_signals.extend(signals) | |
| return np.array(diffuSynth_signals) | |