Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
107 lines
6.2 KiB
Python
107 lines
6.2 KiB
Python
#!/usr/bin/env python3
|
|
"""Run the pinned official Real-time RE-USE model without changing its math."""
|
|
import argparse
|
|
import json
|
|
import logging
|
|
from pathlib import Path
|
|
import statistics
|
|
import sys
|
|
import time
|
|
|
|
import soundfile as sf
|
|
import torch
|
|
from safetensors.torch import load_model
|
|
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument('--source', type=Path, required=True)
|
|
parser.add_argument('--audio', type=Path, required=True)
|
|
parser.add_argument('--out-dir', type=Path, required=True)
|
|
parser.add_argument('--exit-layer', type=int, default=8)
|
|
parser.add_argument('--lookahead-frames', type=int, default=0)
|
|
parser.add_argument('--repeats', type=int, default=2)
|
|
parser.add_argument('--threads', type=int, default=8)
|
|
parser.add_argument('--no-tf32', action='store_true', help='Separate strict-F32 parity diagnostic; not the ordinary performance reference')
|
|
parser.add_argument('--mode', choices=['offline', 'streaming', 'both'], default='both')
|
|
parser.add_argument('--reference-file', type=Path, help='Export aligned network inputs and the final waveform after timing')
|
|
parser.add_argument('--log', type=Path, required=True)
|
|
args = parser.parse_args()
|
|
args.out_dir.mkdir(parents=True, exist_ok=True)
|
|
args.log.parent.mkdir(parents=True, exist_ok=True)
|
|
logging.basicConfig(level=logging.INFO, handlers=[logging.FileHandler(args.log, mode='w'), logging.StreamHandler()])
|
|
sys.path.insert(0, str(args.source.resolve()))
|
|
from models.streaming_generator_SEMamba_time_d1_random_layer_ahead_sep_conv import SEMamba_decoder_list
|
|
from models.stfts import mag_phase_stft, mag_phase_istft
|
|
from online_inference import StreamingSEMamba
|
|
|
|
if not 3 <= args.exit_layer <= 12 or not 0 <= args.lookahead_frames <= 2:
|
|
parser.error('exit-layer must be 3-12 and lookahead-frames 0-2')
|
|
torch.set_num_threads(args.threads)
|
|
if args.no_tf32:
|
|
torch.backends.cuda.matmul.allow_tf32 = False
|
|
torch.backends.cudnn.allow_tf32 = False
|
|
config = json.loads((args.source / 'config.json').read_text())
|
|
model = SEMamba_decoder_list(cfg=config)
|
|
load_model(model, str(args.source / 'model.safetensors'), strict=True)
|
|
model = model.cuda().eval()
|
|
streaming = StreamingSEMamba(model).eval()
|
|
audio, sr = sf.read(args.audio, dtype='float32', always_2d=True)
|
|
if audio.shape[1] != 1:
|
|
raise ValueError('reference requires mono audio')
|
|
waveform = torch.from_numpy(audio.T.copy()).cuda()
|
|
n_fft = 320 * sr // 8000
|
|
hop = 160 * sr // 8000
|
|
n_fft += n_fft % 2
|
|
hop += hop % 2
|
|
modes = ['offline', 'streaming'] if args.mode == 'both' else [args.mode]
|
|
results = {}
|
|
with torch.inference_mode():
|
|
for mode in modes:
|
|
timings = []
|
|
for iteration in range(args.repeats + 1):
|
|
torch.cuda.synchronize()
|
|
torch.cuda.reset_peak_memory_stats()
|
|
started = time.perf_counter()
|
|
mag, phase, _ = mag_phase_stft(waveform, n_fft, hop, n_fft, 'relu_log1p', center=True, addeps=False)
|
|
if mode == 'offline':
|
|
enhanced_mag, enhanced_phase, _ = model(mag, phase, args.exit_layer, args.lookahead_frames)
|
|
zeros = (torch.expm1(torch.relu(enhanced_mag)) == 0).sum(1) / enhanced_mag.shape[1]
|
|
enhanced_mag[:, :, (zeros > .5)[0]] = 0
|
|
else:
|
|
streaming.reset(mag.shape[1] // 2 + 1)
|
|
mag = torch.nn.functional.pad(mag, (0, args.lookahead_frames))
|
|
phase = torch.nn.functional.pad(phase, (0, args.lookahead_frames))
|
|
enhanced_mag, enhanced_phase = streaming(mag[:, :, :args.lookahead_frames + 1], phase[:, :, :args.lookahead_frames + 1], args.exit_layer, args.lookahead_frames)
|
|
magnitudes, phases = [enhanced_mag], [enhanced_phase]
|
|
for frame in range(args.lookahead_frames + 1, mag.shape[-1]):
|
|
enhanced_mag, enhanced_phase = streaming(mag[:, :, frame:frame+1], phase[:, :, frame:frame+1], args.exit_layer, args.lookahead_frames)
|
|
if ((torch.expm1(torch.relu(enhanced_mag)) == 0).sum(1) / enhanced_mag.shape[1]).item() > .5:
|
|
enhanced_mag = 0 * enhanced_mag
|
|
magnitudes.append(enhanced_mag)
|
|
phases.append(enhanced_phase)
|
|
enhanced_mag = torch.stack(magnitudes, -1)
|
|
enhanced_phase = torch.stack(phases, -1)
|
|
output = mag_phase_istft(enhanced_mag, enhanced_phase, n_fft, hop, n_fft, 'relu_log1p')
|
|
output = torch.nn.functional.pad(output, (0, waveform.shape[-1] - output.shape[-1]), value=1e-8)
|
|
torch.cuda.synchronize()
|
|
elapsed = (time.perf_counter() - started) * 1000
|
|
logging.info('session.wall_ms %.9f mode=%s iteration=%d', elapsed, mode, iteration)
|
|
timings.append(elapsed)
|
|
sf.write(args.out_dir / f'python-{mode}.wav', output.cpu().numpy().T, sr, subtype='FLOAT')
|
|
results[mode] = dict(warm_wall_ms=timings[1:], warm_rtf=statistics.mean(timings[1:]) / 1000 / (len(audio) / sr),
|
|
peak_allocated_mib=torch.cuda.max_memory_allocated()/1048576,
|
|
peak_reserved_mib=torch.cuda.max_memory_reserved()/1048576)
|
|
(args.out_dir / 'python-results.json').write_text(json.dumps(results, indent=2))
|
|
logging.info('%s', json.dumps(results))
|
|
|
|
if args.reference_file:
|
|
from safetensors.torch import save_file
|
|
with torch.inference_mode():
|
|
mag, phase, _ = mag_phase_stft(waveform, n_fft, hop, n_fft, 'relu_log1p', center=True, addeps=False)
|
|
enhanced_mag, enhanced_phase, _ = model(mag, phase, args.exit_layer, args.lookahead_frames)
|
|
zeros = (torch.expm1(torch.relu(enhanced_mag)) == 0).sum(1) / enhanced_mag.shape[1]
|
|
enhanced_mag[:, :, (zeros > .5)[0]] = 0
|
|
expected = mag_phase_istft(enhanced_mag, enhanced_phase, n_fft, hop, n_fft, 'relu_log1p')
|
|
expected = torch.nn.functional.pad(expected, (0, waveform.shape[-1] - expected.shape[-1]), value=1e-8)
|
|
save_file({'features': torch.cat([mag.transpose(1, 2), phase.transpose(1, 2)], 0).contiguous().cpu(),
|
|
'waveform': waveform.flatten().cpu(), 'expected': expected.flatten().cpu()}, str(args.reference_file),
|
|
metadata={'sample_rate': str(sr), 'exit_layer': str(args.exit_layer), 'lookahead_frames': str(args.lookahead_frames)})
|