|
import gc
|
|
import os
|
|
import random
|
|
import numpy as np
|
|
from scipy.signal.windows import hann
|
|
import soundfile as sf
|
|
import torch
|
|
from cog import BasePredictor, Input, Path
|
|
import tempfile
|
|
import argparse
|
|
import librosa
|
|
from audiosr import build_model, super_resolution
|
|
from scipy import signal
|
|
import pyloudnorm as pyln
|
|
|
|
|
|
import warnings
|
|
warnings.filterwarnings("ignore")
|
|
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "true"
|
|
torch.set_float32_matmul_precision("high")
|
|
|
|
def match_array_shapes(array_1:np.ndarray, array_2:np.ndarray):
|
|
if (len(array_1.shape) == 1) & (len(array_2.shape) == 1):
|
|
if array_1.shape[0] > array_2.shape[0]:
|
|
array_1 = array_1[:array_2.shape[0]]
|
|
elif array_1.shape[0] < array_2.shape[0]:
|
|
array_1 = np.pad(array_1, ((array_2.shape[0] - array_1.shape[0], 0)), 'constant', constant_values=0)
|
|
else:
|
|
if array_1.shape[1] > array_2.shape[1]:
|
|
array_1 = array_1[:,:array_2.shape[1]]
|
|
elif array_1.shape[1] < array_2.shape[1]:
|
|
padding = array_2.shape[1] - array_1.shape[1]
|
|
array_1 = np.pad(array_1, ((0,0), (0,padding)), 'constant', constant_values=0)
|
|
return array_1
|
|
|
|
|
|
def lr_filter(audio, cutoff, filter_type, order=12, sr=48000):
|
|
audio = audio.T
|
|
nyquist = 0.5 * sr
|
|
normal_cutoff = cutoff / nyquist
|
|
b, a = signal.butter(order//2, normal_cutoff, btype=filter_type, analog=False)
|
|
sos = signal.tf2sos(b, a)
|
|
filtered_audio = signal.sosfiltfilt(sos, audio)
|
|
return filtered_audio.T
|
|
|
|
class Predictor(BasePredictor):
|
|
def setup(self, model_name="basic", device="auto"):
|
|
self.model_name = model_name
|
|
self.device = device
|
|
self.sr = 48000
|
|
print("Loading Model...")
|
|
self.audiosr = build_model(model_name=self.model_name, device=self.device)
|
|
|
|
|
|
print("Model loaded!")
|
|
|
|
def process_audio(self, input_file, chunk_size=5.12, overlap=0.1, seed=None, guidance_scale=3.5, ddim_steps=50):
|
|
audio, sr = librosa.load(input_file, sr=input_cutoff * 2, mono=False)
|
|
audio = audio.T
|
|
sr = input_cutoff * 2
|
|
print(f"audio.shape = {audio.shape}")
|
|
print(f"input cutoff = {input_cutoff}")
|
|
|
|
is_stereo = len(audio.shape) == 2
|
|
audio_channels = [audio] if not is_stereo else [audio[:, 0], audio[:, 1]]
|
|
print("audio is stereo" if is_stereo else "Audio is mono")
|
|
|
|
chunk_samples = int(chunk_size * sr)
|
|
overlap_samples = int(overlap * chunk_samples)
|
|
output_chunk_samples = int(chunk_size * self.sr)
|
|
output_overlap_samples = int(overlap * output_chunk_samples)
|
|
enable_overlap = overlap > 0
|
|
print(f"enable_overlap = {enable_overlap}")
|
|
|
|
def process_chunks(audio):
|
|
chunks = []
|
|
original_lengths = []
|
|
start = 0
|
|
while start < len(audio):
|
|
end = min(start + chunk_samples, len(audio))
|
|
chunk = audio[start:end]
|
|
if len(chunk) < chunk_samples:
|
|
original_lengths.append(len(chunk))
|
|
chunk = np.concatenate([chunk, np.zeros(chunk_samples - len(chunk))])
|
|
else:
|
|
original_lengths.append(chunk_samples)
|
|
chunks.append(chunk)
|
|
start += chunk_samples - overlap_samples if enable_overlap else chunk_samples
|
|
return chunks, original_lengths
|
|
|
|
|
|
chunks_per_channel = [process_chunks(channel) for channel in audio_channels]
|
|
sample_rate_ratio = self.sr / sr
|
|
total_length = len(chunks_per_channel[0][0]) * output_chunk_samples - (len(chunks_per_channel[0][0]) - 1) * (output_overlap_samples if enable_overlap else 0)
|
|
reconstructed_channels = [np.zeros((1, total_length)) for _ in audio_channels]
|
|
|
|
meter_before = pyln.Meter(sr)
|
|
meter_after = pyln.Meter(self.sr)
|
|
|
|
|
|
for ch_idx, (chunks, original_lengths) in enumerate(chunks_per_channel):
|
|
for i, chunk in enumerate(chunks):
|
|
loudness_before = meter_before.integrated_loudness(chunk)
|
|
print(f"Processing chunk {i+1} of {len(chunks)} for {'Left/Mono' if ch_idx == 0 else 'Right'} channel")
|
|
with tempfile.NamedTemporaryFile(suffix=".wav", delete=True) as temp_wav:
|
|
sf.write(temp_wav.name, chunk, sr)
|
|
|
|
out_chunk = super_resolution(
|
|
self.audiosr,
|
|
temp_wav.name,
|
|
seed=seed,
|
|
guidance_scale=guidance_scale,
|
|
ddim_steps=ddim_steps,
|
|
latent_t_per_second=12.8
|
|
)
|
|
|
|
out_chunk = out_chunk[0]
|
|
num_samples_to_keep = int(original_lengths[i] * sample_rate_ratio)
|
|
out_chunk = out_chunk[:, :num_samples_to_keep].squeeze()
|
|
loudness_after = meter_after.integrated_loudness(out_chunk)
|
|
out_chunk = pyln.normalize.loudness(out_chunk, loudness_after, loudness_before)
|
|
|
|
if enable_overlap:
|
|
actual_overlap_samples = min(output_overlap_samples, num_samples_to_keep)
|
|
fade_out = np.linspace(1., 0., actual_overlap_samples)
|
|
fade_in = np.linspace(0., 1., actual_overlap_samples)
|
|
|
|
if i == 0:
|
|
out_chunk[-actual_overlap_samples:] *= fade_out
|
|
elif i < len(chunks) - 1:
|
|
out_chunk[:actual_overlap_samples] *= fade_in
|
|
out_chunk[-actual_overlap_samples:] *= fade_out
|
|
else:
|
|
out_chunk[:actual_overlap_samples] *= fade_in
|
|
|
|
start = i * (output_chunk_samples - output_overlap_samples if enable_overlap else output_chunk_samples)
|
|
end = start + out_chunk.shape[0]
|
|
reconstructed_channels[ch_idx][0, start:end] += out_chunk.flatten()
|
|
|
|
reconstructed_audio = np.stack(reconstructed_channels, axis=-1) if is_stereo else reconstructed_channels[0]
|
|
|
|
if multiband_ensemble:
|
|
low, _ = librosa.load(input_file, sr=48000, mono=False)
|
|
output = match_array_shapes(reconstructed_audio[0].T, low)
|
|
low = lr_filter(low.T, crossover_freq, 'lowpass', order=10)
|
|
high = lr_filter(output.T, crossover_freq, 'highpass', order=10)
|
|
high = lr_filter(high, 23000, 'lowpass', order=2)
|
|
output = low + high
|
|
else:
|
|
output = reconstructed_audio[0]
|
|
|
|
return output
|
|
|
|
|
|
def predict(self,
|
|
input_file: Path = Input(description="Audio to upsample"),
|
|
ddim_steps: int = Input(description="Number of inference steps", default=50, ge=10, le=500),
|
|
guidance_scale: float = Input(description="Scale for classifier free guidance", default=3.5, ge=1.0, le=20.0),
|
|
overlap: float = Input(description="overlap size", default=0.04),
|
|
chunk_size: float = Input(description="chunksize", default=10.24),
|
|
seed: int = Input(description="Random seed. Leave blank to randomize the seed", default=None)
|
|
) -> Path:
|
|
|
|
if seed == 0:
|
|
seed = random.randint(0, 2**32 - 1)
|
|
print(f"Setting seed to: {seed}")
|
|
print(f"overlap = {overlap}")
|
|
print(f"guidance_scale = {guidance_scale}")
|
|
print(f"ddim_steps = {ddim_steps}")
|
|
print(f"chunk_size = {chunk_size}")
|
|
print(f"multiband_ensemble = {multiband_ensemble}")
|
|
print(f"input file = {os.path.basename(input_file)}")
|
|
os.makedirs(output_folder, exist_ok=True)
|
|
waveform = self.process_audio(
|
|
input_file,
|
|
chunk_size=chunk_size,
|
|
overlap=overlap,
|
|
seed=seed,
|
|
guidance_scale=guidance_scale,
|
|
ddim_steps=ddim_steps
|
|
)
|
|
|
|
filename = os.path.splitext(os.path.basename(input_file))[0]
|
|
sf.write(f"{output_folder}/SR_{filename}.wav", data=waveform, samplerate=48000, subtype="PCM_16")
|
|
print(f"file created: {output_folder}/SR_{filename}.wav")
|
|
del self.audiosr, waveform
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
parser = argparse.ArgumentParser(description="Find volume difference of two audio files.")
|
|
parser.add_argument("--input", help="Path to input audio file")
|
|
parser.add_argument("--output", help="Output folder")
|
|
parser.add_argument("--ddim_steps", help="Number of ddim steps", type=int, required=False, default=50)
|
|
parser.add_argument("--chunk_size", help="chunk size", type=float, required=False, default=10.24)
|
|
parser.add_argument("--guidance_scale", help="Guidance scale value", type=float, required=False, default=3.5)
|
|
parser.add_argument("--seed", help="Seed value, 0 = random seed", type=int, required=False, default=0)
|
|
parser.add_argument("--overlap", help="overlap value", type=float, required=False, default=0.04)
|
|
parser.add_argument("--multiband_ensemble", type=bool, help="Use multiband ensemble with input")
|
|
parser.add_argument("--input_cutoff", help="Define the crossover of audio input in the multiband ensemble", type=int, required=False, default=12000)
|
|
|
|
args = parser.parse_args()
|
|
|
|
input_file_path = args.input
|
|
output_folder = args.output
|
|
ddim_steps = args.ddim_steps
|
|
chunk_size = args.chunk_size
|
|
guidance_scale = args.guidance_scale
|
|
seed = args.seed
|
|
overlap = args.overlap
|
|
input_cutoff = args.input_cutoff
|
|
multiband_ensemble = args.multiband_ensemble
|
|
|
|
crossover_freq = input_cutoff - 1000
|
|
|
|
p = Predictor()
|
|
|
|
p.setup(device='auto')
|
|
|
|
|
|
out = p.predict(
|
|
input_file_path,
|
|
ddim_steps=ddim_steps,
|
|
guidance_scale=guidance_scale,
|
|
seed=seed,
|
|
chunk_size=chunk_size,
|
|
overlap=overlap
|
|
)
|
|
|
|
del p
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
