mirror of
https://github.com/Nighthawk42/soprano-factory.git
synced 2026-08-30 04:30:21 +00:00
81 lines
2.5 KiB
Python
81 lines
2.5 KiB
Python
# verify_voice.py - Pure ASCII, Standardized for V3
|
|
import argparse
|
|
import torch
|
|
import soundfile as sf
|
|
import torchaudio
|
|
import os
|
|
from model.encoder import Encoder
|
|
from model.decoder import Decoder
|
|
from utils.config import cfg
|
|
|
|
def setup_audio_backend():
|
|
cwd = os.getcwd()
|
|
local_ffmpeg = os.path.join(cwd, "tools", "ffmpeg")
|
|
if os.path.exists(local_ffmpeg) and os.path.isdir(local_ffmpeg):
|
|
if local_ffmpeg not in os.environ["PATH"]:
|
|
os.environ["PATH"] = local_ffmpeg + os.pathsep + os.environ["PATH"]
|
|
|
|
setup_audio_backend()
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Verify Codec Reconstruction")
|
|
parser.add_argument("--input", type=str, required=True, help="Single wav file")
|
|
parser.add_argument("--output", type=str, default="verify_out.wav")
|
|
args = parser.parse_args()
|
|
|
|
device = cfg.device
|
|
print(f"Running verification on {device}")
|
|
|
|
# Initialize Encoder
|
|
encoder = Encoder(
|
|
num_input_mels=cfg.codec['input_mels'],
|
|
encoder_dim=cfg.codec['dim'],
|
|
encoder_layers=cfg.codec['layers'],
|
|
bottleneck_channels=cfg.codec['bottleneck']
|
|
).to(device).float()
|
|
|
|
# Initialize Decoder (5 channels for Stage 0 weights)
|
|
decoder = Decoder(
|
|
input_channels=cfg.codec['bottleneck'],
|
|
decoder_dim=cfg.codec['dim'],
|
|
decoder_layers=cfg.codec['layers']
|
|
).to(device).float()
|
|
|
|
# Load weights
|
|
e_path = os.path.join(cfg.codec['save_dir'], "encoder.pth")
|
|
d_path = os.path.join(cfg.codec['save_dir'], "decoder.pth")
|
|
|
|
if not os.path.exists(e_path) or not os.path.exists(d_path):
|
|
print(f"Error: Weights not found in {cfg.codec['save_dir']}")
|
|
return
|
|
|
|
encoder.load_state_dict(torch.load(e_path, map_location=device, weights_only=True))
|
|
decoder.load_state_dict(torch.load(d_path, map_location=device, weights_only=True))
|
|
|
|
encoder.eval()
|
|
decoder.eval()
|
|
|
|
# Load Audio
|
|
wav, sr = sf.read(args.input)
|
|
wav = torch.from_numpy(wav).float()
|
|
if wav.ndim == 1: wav = wav.unsqueeze(0)
|
|
else: wav = wav.t()
|
|
|
|
if sr != cfg.common['sample_rate']:
|
|
wav = torchaudio.functional.resample(wav, sr, cfg.common['sample_rate'])
|
|
|
|
if wav.shape[0] > 1: wav = wav.mean(dim=0, keepdim=True)
|
|
wav = wav.unsqueeze(0).to(device)
|
|
|
|
# Process
|
|
with torch.no_grad():
|
|
z = encoder(wav)
|
|
out = decoder(z)
|
|
|
|
# Save
|
|
out_np = out.squeeze().cpu().numpy()
|
|
sf.write(args.output, out_np, cfg.common['sample_rate'])
|
|
print(f"Verification complete: {args.output}")
|
|
|
|
if __name__ == "__main__":
|
|
main() |