Spaces:
Sleeping
Sleeping
File size: 2,545 Bytes
541c6b7 |
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 |
from pathlib import Path
import argbind
import torch
from audiotools import AudioSignal
from audiotools.core import util
from audiotools.ml.decorators import Tracker
from train import Accelerator
from train import DAC
from dac.compare.encodec import Encodec
Encodec = argbind.bind(Encodec)
def load_state(
accel: Accelerator,
tracker: Tracker,
save_path: str,
tag: str = "latest",
load_weights: bool = False,
model_type: str = "dac",
bandwidth: float = 24.0,
):
kwargs = {
"folder": f"{save_path}/{tag}",
"map_location": "cpu",
"package": not load_weights,
}
tracker.print(f"Resuming from {str(Path('.').absolute())}/{kwargs['folder']}")
if model_type == "dac":
generator, _ = DAC.load_from_folder(**kwargs)
elif model_type == "encodec":
generator = Encodec(bandwidth=bandwidth)
generator = accel.prepare_model(generator)
return generator
@torch.no_grad()
def process(signal, accel, generator, **kwargs):
signal = signal.to(accel.device)
recons = generator(signal.audio_data, signal.sample_rate, **kwargs)["audio"]
recons = AudioSignal(recons, signal.sample_rate)
recons = recons.normalize(signal.loudness())
return recons.cpu()
@argbind.bind(without_prefix=True)
@torch.no_grad()
def get_samples(
accel,
path: str = "ckpt",
input: str = "samples/input",
output: str = "samples/output",
model_type: str = "dac",
model_tag: str = "latest",
bandwidth: float = 24.0,
n_quantizers: int = None,
):
tracker = Tracker(log_file=f"{path}/eval.txt", rank=accel.local_rank)
generator = load_state(
accel,
tracker,
save_path=path,
model_type=model_type,
bandwidth=bandwidth,
tag=model_tag,
)
generator.eval()
kwargs = {"n_quantizers": n_quantizers} if model_type == "dac" else {}
audio_files = util.find_audio(input)
global process
process = tracker.track("process", len(audio_files))(process)
output = Path(output)
output.mkdir(parents=True, exist_ok=True)
with tracker.live:
for i in range(len(audio_files)):
signal = AudioSignal(audio_files[i])
recons = process(signal, accel, generator, **kwargs)
recons.write(output / audio_files[i].name)
tracker.done("test", f"N={len(audio_files)}")
if __name__ == "__main__":
args = argbind.parse_args()
with argbind.scope(args):
with Accelerator() as accel:
get_samples(accel)
|