File size: 8,488 Bytes
f032e68 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 |
import gradio as gr
import spaces
import yaml
import torch
import librosa
from diffusers import DDIMScheduler
from transformers import AutoProcessor, ClapModel
from model.udit import UDiT
from vae_modules.autoencoder_wrapper import Autoencoder
import numpy as np
diffusion_config = './config/SoloAudio.yaml'
diffusion_ckpt = './pretrained_models/soloaudio_v2.pt'
autoencoder_path = './pretrained_models/audio-vae.pt'
uncond_path = './pretrained_models/uncond.npz'
sample_rate = 24000
device = 'cuda' if torch.cuda.is_available() else 'cpu'
with open(diffusion_config, 'r') as fp:
diff_config = yaml.safe_load(fp)
v_prediction = diff_config["ddim"]["v_prediction"]
clapmodel = ClapModel.from_pretrained("laion/larger_clap_general").to(device)
processor = AutoProcessor.from_pretrained('laion/larger_clap_general')
autoencoder = Autoencoder(autoencoder_path, 'stable_vae', quantization_first=True)
autoencoder.eval()
autoencoder.to(device)
unet = UDiT(
**diff_config['diffwrap']['UDiT']
).to(device)
unet.load_state_dict(torch.load(diffusion_ckpt)['model'])
unet.eval()
if v_prediction:
print('v prediction')
scheduler = DDIMScheduler(**diff_config["ddim"]['diffusers'])
else:
print('noise prediction')
scheduler = DDIMScheduler(**diff_config["ddim"]['diffusers'])
# these steps reset dtype of noise_scheduler params
latents = torch.randn((1, 128, 128),
device=device)
noise = torch.randn(latents.shape).to(latents.device)
timesteps = torch.randint(0, scheduler.config.num_train_timesteps,
(noise.shape[0],),
device=latents.device).long()
_ = scheduler.add_noise(latents, noise, timesteps)
def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
"""
Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and
Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4
"""
std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
# rescale the results from guidance (fixes overexposure)
noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
# mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
return noise_cfg
@spaces.GPU
def sample_diffusion(mixture, timbre, ddim_steps=50, eta=0, seed=2023, guidance_scale=False, guidance_rescale=0.0,):
with torch.no_grad():
scheduler.set_timesteps(ddim_steps)
generator = torch.Generator(device=device).manual_seed(seed)
# init noise
noise = torch.randn(mixture.shape, generator=generator, device=device)
pred = noise
for t in scheduler.timesteps:
pred = scheduler.scale_model_input(pred, t)
if guidance_scale:
uncond = torch.tensor(np.load(uncond_path)['arr_0']).unsqueeze(0).to(device)
pred_combined = torch.cat([pred, pred], dim=0)
mixture_combined = torch.cat([mixture, mixture], dim=0)
timbre_combined = torch.cat([timbre, uncond], dim=0)
output_combined = unet(x=pred_combined, timesteps=t, mixture=mixture_combined, timbre=timbre_combined)
output_pos, output_neg = torch.chunk(output_combined, 2, dim=0)
model_output = output_neg + guidance_scale * (output_pos - output_neg)
if guidance_rescale > 0.0:
# avoid overexposed
model_output = rescale_noise_cfg(model_output, output_pos,
guidance_rescale=guidance_rescale)
else:
model_output = unet(x=pred, timesteps=t, mixture=mixture, timbre=timbre)
pred = scheduler.step(model_output=model_output, timestep=t, sample=pred,
eta=eta, generator=generator).prev_sample
pred = autoencoder(embedding=pred).squeeze(1)
return pred
@spaces.GPU
def tse(gt_file_input, text_input, num_infer_steps, eta, seed, guidance_scale, guidance_rescale):
with torch.no_grad():
mixture, _ = librosa.load(gt_file_input, sr=sample_rate)
# Check the length of the audio in samples
current_length = len(mixture)
target_length = sample_rate * 10
# Cut or pad the audio to match the target length
if current_length > target_length:
# Trim the audio if it's longer than the target length
mixture = mixture[:target_length]
elif current_length < target_length:
# Pad the audio with zeros if it's shorter than the target length
padding = target_length - current_length
mixture = np.pad(mixture, (0, padding), mode='constant')
mixture = torch.tensor(mixture).unsqueeze(0).to(device)
mixture = autoencoder(audio=mixture.unsqueeze(1))
text_inputs = processor(
text=[text_input],
max_length=10, # Fixed length for text
padding='max_length', # Pad text to max length
truncation=True, # Truncate text if it's longer than max length
return_tensors="pt"
)
inputs = {
"input_ids": text_inputs["input_ids"][0].unsqueeze(0), # Text input IDs
"attention_mask": text_inputs["attention_mask"][0].unsqueeze(0), # Attention mask for text
}
inputs = {key: value.to(device) for key, value in inputs.items()}
timbre = clapmodel.get_text_features(**inputs)
pred = sample_diffusion(mixture, timbre, num_infer_steps, eta, seed, guidance_scale, guidance_rescale)
return sample_rate, pred.squeeze().cpu().numpy()
# CSS styling (optional)
css = """
#col-container {
margin: 0 auto;
max-width: 1280px;
}
"""
# Gradio Blocks layout
with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
with gr.Column(elem_id="col-container"):
gr.Markdown("""
# SoloAudio: Target Sound Extraction with Language-oriented Audio Diffusion Transformer.
Adjust advanced settings for more control. This space only supports a 10-second audio input now.
Learn more about 🟣**SoloAudio** on the [SoloAudio Homepage](https://wanghelin1997.github.io/SoloAudio-Demo/).
""")
with gr.Tab("Target Sound Extraction"):
# Basic Input: Text prompt
with gr.Row():
gt_file_input = gr.Audio(label="Upload Audio to Extract", type="filepath", value="demo/0_mix.wav")
text_input = gr.Textbox(
label="Text Prompt",
show_label=True,
max_lines=2,
placeholder="Enter your prompt",
container=True,
value="The sound of gunshot",
scale=4
)
# Run button
run_button = gr.Button("Extract", scale=1)
# Output Component
result = gr.Audio(label="Extracted Audio", type="numpy")
# Advanced settings in an Accordion
with gr.Accordion("Advanced Settings", open=False):
# Audio Length
guidance_scale = gr.Slider(minimum=1.0, maximum=10, step=0.1, value=3.0, label="Guidance Scale")
guidance_rescale = gr.Slider(minimum=0.0, maximum=1, step=0.05, value=0., label="Guidance Rescale")
num_infer_steps = gr.Slider(minimum=25, maximum=200, step=5, value=50, label="DDIM Steps")
eta = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0.0, label="Eta")
seed = gr.Slider(minimum=0, maximum=100, step=1, value=0, label="Seed")
# Define the trigger and input-output linking for generation
run_button.click(
fn=tse,
inputs=[gt_file_input, text_input, num_infer_steps, eta, seed, guidance_scale, guidance_rescale],
outputs=[result]
)
text_input.submit(fn=tse,
inputs=[gt_file_input, text_input, num_infer_steps, eta, seed, guidance_scale, guidance_rescale],
outputs=[result]
)
# Launch the Gradio demo
demo.launch() |