import torch import numpy as np import gradio as gr from PIL import Image import matplotlib from omegaconf import OmegaConf from einops import repeat import librosa from ldm.models.diffusion.ddim import DDIMSampler from vocoder.bigvgan.models import VocoderBigVGAN from ldm.util import instantiate_from_config from ldm.data.extract_mel_spectrogram import TRANSFORMS_16000 SAMPLE_RATE = 16000 cmap_transform = matplotlib.cm.viridis torch.set_grad_enabled(False) device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") def initialize_model(config, ckpt): config = OmegaConf.load(config) model = instantiate_from_config(config.model) model.load_state_dict(torch.load(ckpt,map_location='cpu')["state_dict"], strict=False) model = model.to(device) print(model.device,device,model.cond_stage_model.device) sampler = DDIMSampler(model) return sampler def make_batch_sd( mel, mask, device, num_samples=1): mel = torch.from_numpy(mel)[None,None,...].to(dtype=torch.float32) mask = torch.from_numpy(mask)[None,None,...].to(dtype=torch.float32) masked_mel = (1 - mask) * mel mel = mel * 2 - 1 mask = mask * 2 - 1 masked_mel = masked_mel * 2 -1 batch = { "mel": repeat(mel.to(device=device), "1 ... -> n ...", n=num_samples), "mask": repeat(mask.to(device=device), "1 ... -> n ...", n=num_samples), "masked_mel": repeat(masked_mel.to(device=device), "1 ... -> n ...", n=num_samples), } return batch def gen_mel(input_audio): sr,ori_wav = input_audio print(sr,ori_wav.shape,ori_wav) ori_wav = ori_wav.astype(np.float32, order='C') / 32768.0 # order='C'是以C语言格式存储，不用管 if len(ori_wav.shape)==2:# stereo ori_wav = librosa.to_mono(ori_wav.T)# gradio load wav shape could be (wav_len,2) but librosa expects (2,wav_len) print(sr,ori_wav.shape,ori_wav) ori_wav = librosa.resample(ori_wav,orig_sr = sr,target_sr = SAMPLE_RATE) mel_len,hop_size = 848,256 input_len = mel_len * hop_size if len(ori_wav) < input_len: input_wav = np.pad(ori_wav,(0,mel_len*hop_size),constant_values=0) else: input_wav = ori_wav[:input_len] mel = TRANSFORMS_16000(input_wav) return mel def show_mel_fn(input_audio): crop_len = 500 # the full mel cannot be showed due to gradio's Image bug when using tool='sketch' crop_mel = gen_mel(input_audio)[:,:crop_len] color_mel = cmap_transform(crop_mel) return Image.fromarray((color_mel*255).astype(np.uint8)) def inpaint(sampler, batch, seed, ddim_steps, num_samples=1, W=512, H=512): model = sampler.model prng = np.random.RandomState(seed) start_code = prng.randn(num_samples, model.first_stage_model.embed_dim, H // 8, W // 8) start_code = torch.from_numpy(start_code).to(device=device, dtype=torch.float32) c = model.get_first_stage_encoding(model.encode_first_stage(batch["masked_mel"])) cc = torch.nn.functional.interpolate(batch["mask"], size=c.shape[-2:]) c = torch.cat((c, cc), dim=1) # (b,c+1,h,w) 1 is mask shape = (c.shape[1]-1,)+c.shape[2:] samples_ddim, _ = sampler.sample(S=ddim_steps, conditioning=c, batch_size=c.shape[0], shape=shape, verbose=False) x_samples_ddim = model.decode_first_stage(samples_ddim) mask = batch["mask"]# [-1,1] mel = torch.clamp((batch["mel"]+1.0)/2.0,min=0.0, max=1.0) mask = torch.clamp((batch["mask"]+1.0)/2.0,min=0.0, max=1.0) predicted_mel = torch.clamp((x_samples_ddim+1.0)/2.0,min=0.0, max=1.0) inpainted = (1-mask)*mel+mask*predicted_mel inpainted = inpainted.cpu().numpy().squeeze() inapint_wav = vocoder.vocode(inpainted) return inpainted,inapint_wav def predict(input_audio,mel_and_mask,ddim_steps,seed): show_mel = np.array(mel_and_mask['image'].convert("L"))/255 # 由于展示的mel只展示了一部分，所以需要重新从音频生成mel mask = np.array(mel_and_mask["mask"].convert("L"))/255 mel_bins,mel_len = 80,848 input_mel = gen_mel(input_audio)[:,:mel_len]# 由于展示的mel只展示了一部分，所以需要重新从音频生成mel mask = np.pad(mask,((0,0),(0,mel_len-mask.shape[1])),mode='constant',constant_values=0)# 将mask填充到原来的mel的大小 print(mask.shape,input_mel.shape) with torch.no_grad(): batch = make_batch_sd(input_mel,mask,device,num_samples=1) inpainted,gen_wav = inpaint( sampler=sampler, batch=batch, seed=seed, ddim_steps=ddim_steps, num_samples=1, H=mel_bins, W=mel_len ) inpainted = inpainted[:,:show_mel.shape[1]] color_mel = cmap_transform(inpainted) input_len = int(input_audio[1].shape[0] * SAMPLE_RATE / input_audio[0]) gen_wav = (gen_wav * 32768).astype(np.int16)[:input_len] return Image.fromarray((color_mel*255).astype(np.uint8)),(SAMPLE_RATE,gen_wav) sampler = initialize_model('./configs/inpaint/txt2audio_args.yaml', './useful_ckpts/inpaint7_epoch00047.ckpt') vocoder = VocoderBigVGAN('./vocoder/logs/bigv16k53w',device=device) block = gr.Blocks().queue() with block: with gr.Row(): gr.Markdown("## Make-An-Audio Inpainting") with gr.Row(): with gr.Column(): input_audio = gr.inputs.Audio() show_button = gr.Button("Show Mel") run_button = gr.Button("Predict Masked Place") with gr.Accordion("Advanced options", open=False): ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=150, value=100, step=1) seed = gr.Slider( label="Seed", minimum=0, maximum=2147483647, step=1, randomize=True, ) with gr.Column(): show_inpainted = gr.Image(type="pil").style(width=848,height=80) outaudio = gr.Audio() show_mel = gr.Image(type="pil",tool='sketch')#.style(width=848,height=80) # 加上这个没办法展示完全图片 show_button.click(fn=show_mel_fn, inputs=[input_audio], outputs=show_mel) run_button.click(fn=predict, inputs=[input_audio,show_mel,ddim_steps,seed], outputs=[show_inpainted,outaudio]) block.launch()