import sys
import soundfile as sf
from huggingface_hub import hf_hub_download
import resemble_enhance_mlx as re

ckpt = hf_hub_download("multimodalart/resemble-enhance-mlx", "resemble-enhance-mlx.safetensors")
model = re.load(ckpt, lambd=0.5)

src = sys.argv[1] if len(sys.argv) > 1 else "input.wav"
wav, sr = sf.read(src, dtype="float32")
if wav.ndim > 1:
    wav = wav.mean(axis=1)

out, out_sr = re.enhance(model, wav, sr)
sf.write("enhanced.wav", out, out_sr)
print(f"wrote enhanced.wav ({len(out) / out_sr:.1f}s @ {out_sr} Hz)")