resemble-enhance icon indicating copy to clipboard operation
resemble-enhance copied to clipboard

inference mode with safetensors and no deespeed

Open rsxdalv opened this issue 10 months ago • 0 comments

Deepspeed was not required for inference, therefore, by this refactoring, users can use the project without deepspeed installed.

Now deepspeed is part of the 'train' installation: pip install resemble-enhance[train]

Additionally, omegaconf and .pt files are preferrable for training, but we can use safetensors and config.json for easier portability.

Finally, since there were 2 models packed into one checkpoint, I split them into denoiser and enhancer: one model.safetensors and config.json for each.

New files: https://huggingface.co/rsxdalv/resemble-enhance/tree/main

Example code:

import torch
import torchaudio

if torch.cuda.is_available():
    device = "cuda"
else:
    device = "cpu"
from resemble_enhance.enhancer.inference import denoise, enhance, download

if path is None:
    return None, None

solver = solver.lower()
nfe = int(nfe)
lambd = 0.9 if denoising else 0.1

dwav, sr = torchaudio.load(path)
dwav = dwav.mean(dim=0)

download(RUN_DIR, safetensors=True)
wav1, new_sr = denoise(dwav, sr, device, run_dir=RUN_DIR / "denoiser")
wav2, new_sr = enhance(
    dwav, sr, device, nfe=nfe, solver=solver, lambd=lambd, tau=tau, run_dir=RUN_DIR / "enhancer", skip_download=True,
)

wav1 = wav1.cpu().numpy()
wav2 = wav2.cpu().numpy()

return (new_sr, wav1), (new_sr, wav2)

rsxdalv avatar Jun 01 '25 18:06 rsxdalv