Spaces:
Runtime error
Runtime error
File size: 8,553 Bytes
06f2f15 e0ce7d2 06f2f15 57ede1f 06f2f15 11ef812 06f2f15 57ede1f 11ef812 57ede1f 06f2f15 |
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 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 |
#%%
import os
os.system("git clone https://github.com/v-iashin/SpecVQGAN")
os.system("conda env create -f ./SpecVQGAN/conda_env.yml")
os.system("conda activate specvqgan")
os.system("pip install pytorch-lightning==1.2.10 omegaconf==2.0.6 streamlit==0.80 matplotlib==3.4.1 albumentations==0.5.2 SoundFile torch torchvision librosa gdown")
# %%
import sys
sys.path.append('./SpecVQGAN')
import time
from pathlib import Path
import IPython.display as display_audio
import soundfile
import torch
from IPython import display
from matplotlib import pyplot as plt
from torch.utils.data.dataloader import default_collate
from torchvision.utils import make_grid
from tqdm import tqdm
from feature_extraction.demo_utils import (ExtractResNet50, check_video_for_audio,
extract_melspectrogram, load_model,
show_grid, trim_video)
from sample_visualization import (all_attention_to_st, get_class_preditions,
last_attention_to_st, spec_to_audio_to_st,
tensor_to_plt)
from specvqgan.data.vggsound import CropImage
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# load model
model_name = '2021-07-30T21-34-25_vggsound_transformer'
log_dir = './logs'
os.chdir("./SpecVQGAN/")
config, sampler, melgan, melception = load_model(model_name, log_dir, device)
# %%
def extract_thumbnails(video_path):
# Trim the video
start_sec = 0 # to start with 01:35 use 95 seconds
video_path = trim_video(video_path, start_sec, trim_duration=10)
# Extract Features
extraction_fps = 21.5
feature_extractor = ExtractResNet50(extraction_fps, config.data.params, device)
visual_features, resampled_frames = feature_extractor(video_path)
# Show the selected frames to extract features for
if not config.data.params.replace_feats_with_random:
fig = show_grid(make_grid(resampled_frames))
fig.show()
# Prepare Input
batch = default_collate([visual_features])
batch['feature'] = batch['feature'].to(device)
c = sampler.get_input(sampler.cond_stage_key, batch)
return c, video_path
# %%
import numpy as np
def generate_audio(video_path, temperature = 1.0):
# Define Sampling Parameters
W_scale = 1
mode = 'full'
top_x = sampler.first_stage_model.quantize.n_e // 2
update_every = 0 # use > 0 value, e.g. 15, to see the progress of generation (slows down the sampling speed)
full_att_mat = True
c, video_path = extract_thumbnails(video_path)
# Start sampling
with torch.no_grad():
start_t = time.time()
quant_c, c_indices = sampler.encode_to_c(c)
# crec = sampler.cond_stage_model.decode(quant_c)
patch_size_i = 5
patch_size_j = 53
B, D, hr_h, hr_w = sampling_shape = (1, 256, 5, 53*W_scale)
z_pred_indices = torch.zeros((B, hr_h*hr_w)).long().to(device)
if mode == 'full':
start_step = 0
else:
start_step = (patch_size_j // 2) * patch_size_i
z_pred_indices[:, :start_step] = z_indices[:, :start_step]
pbar = tqdm(range(start_step, hr_w * hr_h), desc='Sampling Codebook Indices')
for step in pbar:
i = step % hr_h
j = step // hr_h
i_start = min(max(0, i - (patch_size_i // 2)), hr_h - patch_size_i)
j_start = min(max(0, j - (patch_size_j // 2)), hr_w - patch_size_j)
i_end = i_start + patch_size_i
j_end = j_start + patch_size_j
local_i = i - i_start
local_j = j - j_start
patch_2d_shape = (B, D, patch_size_i, patch_size_j)
pbar.set_postfix(
Step=f'({i},{j}) | Local: ({local_i},{local_j}) | Crop: ({i_start}:{i_end},{j_start}:{j_end})'
)
patch = z_pred_indices \
.reshape(B, hr_w, hr_h) \
.permute(0, 2, 1)[:, i_start:i_end, j_start:j_end].permute(0, 2, 1) \
.reshape(B, patch_size_i * patch_size_j)
# assuming we don't crop the conditioning and just use the whole c, if not desired uncomment the above
cpatch = c_indices
logits, _, attention = sampler.transformer(patch[:, :-1], cpatch)
# remove conditioning
logits = logits[:, -patch_size_j*patch_size_i:, :]
local_pos_in_flat = local_j * patch_size_i + local_i
logits = logits[:, local_pos_in_flat, :]
logits = logits / temperature
logits = sampler.top_k_logits(logits, top_x)
# apply softmax to convert to probabilities
probs = torch.nn.functional.softmax(logits, dim=-1)
# sample from the distribution
ix = torch.multinomial(probs, num_samples=1)
z_pred_indices[:, j * hr_h + i] = ix
if update_every > 0 and step % update_every == 0:
z_pred_img = sampler.decode_to_img(z_pred_indices, sampling_shape)
# fliping the spectrogram just for illustration purposes (low freqs to bottom, high - top)
z_pred_img_st = tensor_to_plt(z_pred_img, flip_dims=(2,))
display.clear_output(wait=True)
display.display(z_pred_img_st)
if full_att_mat:
att_plot = all_attention_to_st(attention, placeholders=None, scale_by_prior=True)
display.display(att_plot)
plt.close()
else:
quant_z_shape = sampling_shape
c_length = cpatch.shape[-1]
quant_c_shape = quant_c.shape
c_att_plot, z_att_plot = last_attention_to_st(
attention, local_pos_in_flat, c_length, sampler.first_stage_permuter,
sampler.cond_stage_permuter, quant_c_shape, patch_2d_shape,
placeholders=None, flip_c_dims=None, flip_z_dims=(2,))
display.display(c_att_plot)
display.display(z_att_plot)
plt.close()
plt.close()
plt.close()
# quant_z_shape = sampling_shape
z_pred_img = sampler.decode_to_img(z_pred_indices, sampling_shape)
# showing the final image
z_pred_img_st = tensor_to_plt(z_pred_img, flip_dims=(2,))
display.clear_output(wait=True)
display.display(z_pred_img_st)
if full_att_mat:
att_plot = all_attention_to_st(attention, placeholders=None, scale_by_prior=True)
display.display(att_plot)
plt.close()
else:
quant_z_shape = sampling_shape
c_length = cpatch.shape[-1]
quant_c_shape = quant_c.shape
c_att_plot, z_att_plot = last_attention_to_st(
attention, local_pos_in_flat, c_length, sampler.first_stage_permuter,
sampler.cond_stage_permuter, quant_c_shape, patch_2d_shape,
placeholders=None, flip_c_dims=None, flip_z_dims=(2,)
)
display.display(c_att_plot)
display.display(z_att_plot)
plt.close()
plt.close()
plt.close()
print(f'Sampling Time: {time.time() - start_t:3.2f} seconds')
waves = spec_to_audio_to_st(z_pred_img, config.data.params.spec_dir_path,
config.data.params.sample_rate, show_griffin_lim=False,
vocoder=melgan, show_in_st=False)
print(f'Sampling Time (with vocoder): {time.time() - start_t:3.2f} seconds')
print(f'Generated: {len(waves["vocoder"]) / config.data.params.sample_rate:.2f} seconds')
# Melception opinion on the class distribution of the generated sample
topk_preds = get_class_preditions(z_pred_img, melception)
print(topk_preds)
audio_path = os.path.join(log_dir, Path(video_path).stem + '.wav')
audio = waves['vocoder']
audio = np.repeat([audio], 2, axis=0).T
print(audio.shape)
soundfile.write(audio_path, audio, config.data.params.sample_rate, 'PCM_24')
print(f'The sample has been saved @ {audio_path}')
video_out_path = os.path.join(log_dir, Path(video_path).stem + '_audio.mp4')
print(video_path, audio_path, video_out_path)
os.system("ffmpeg -i %s -i %s -map 0:v -map 1:a -c:v copy -shortest %s" % (video_path, audio_path, video_out_path))
return video_out_path
# return config.data.params.sample_rate, audio
# %%
#generate_audio("../kiss.avi")
#%%
import gradio as gr
iface = gr.Interface(generate_audio, "video", ["playable_video"],
description="Generate audio based on the video input")
iface.launch()
# %%
|