| import sys,time,argparse |
| from pathlib import Path |
| import numpy as np, torch, torchaudio, soundfile as sf, axengine |
| ROOT=Path(__file__).resolve().parent; sys.path.insert(0,str(ROOT)) |
| T=128; F_HR=432; HOP=512; NFFT=1024; SR=48000 |
| WIN=torch.hann_window(NFFT) |
| def stft(x): |
| X=torch.stft(x.reshape(1,-1),NFFT,HOP,window=WIN,center=True,return_complex=True) |
| A=X+1e-4; X=A.abs()**0.2*torch.exp(1j*A.angle()) |
| r=torch.view_as_real(X).permute(0,3,1,2); return r[:,:,:-1,:] |
| def istft(sp): |
| sp=torch.nn.functional.pad(sp,[0,0,0,1],value=0).permute(0,2,3,1).contiguous() |
| X=torch.view_as_complex(sp); X=X.abs()**5.0*torch.exp(1j*X.angle()) |
| return torch.istft(X,NFFT,HOP,window=WIN,center=True,length=245760) |
| def main(): |
| p=argparse.ArgumentParser(); p.add_argument('input'); p.add_argument('output'); p.add_argument('--models',default=str(ROOT/'models')); p.add_argument('--seed',type=int,default=42); p.add_argument('--overlap',type=int,default=64); p.add_argument('--sr',type=int,default=8); a=p.parse_args() |
| ce=axengine.InferenceSession(str(Path(a.models)/f'universr_split_{a.sr}k_condenc.axmodel')) |
| mc=axengine.InferenceSession(str(Path(a.models)/f'universr_split_{a.sr}k_main_cond_actual.axmodel')) |
| mu=axengine.InferenceSession(str(Path(a.models)/f'universr_split_{a.sr}k_main_uncond.axmodel')) |
| x,sr=sf.read(a.input); x=torch.tensor(x,dtype=torch.float32).reshape(1,-1) |
| x=torchaudio.functional.resample(x,sr,SR); x=torch.nn.functional.pad(x,(0,max(0,245760-x.shape[-1])))[...,:245760] |
| Y=stft(x.unsqueeze(1)); total=Y.shape[-1]; lr=80 if a.sr==8 else 170; hf=512-432 |
| pad=(T-total%T)%T |
| Yp=torch.nn.functional.pad(Y,(0,pad),mode='constant') if pad else Y |
| total=Yp.shape[-1] |
| starts=list(range(0,total-T+1,T-a.overlap)) |
| spec=torch.zeros_like(Yp); wsum=torch.zeros(total) |
| rng=np.random.RandomState(a.seed); t0=time.time() |
| for s in starts: |
| y=Yp[:,:,:lr,s:s+T].numpy().astype(np.float32) |
| z=ce.run(None,{'y':y})[0].astype(np.float32) |
| zz=rng.randn(1,2,F_HR,T).astype(np.float32) |
| n=128 |
| for i in range(4): |
| t=np.array([i/4],np.float32); dt=0.25 |
| def f(q,tt): |
| vu=mu.run(None,{'x':q,'t':tt})[0]; vc=mc.run(None,{'x':q,'t':tt,'y_cond':z})[0] |
| return (1-1.5)*vu+1.5*vc |
| k1=f(zz,t); zm=zz+0.5*dt*k1; k2=f(zm,np.array([i/4+0.125],np.float32)); zz=zz+dt*k2 |
| cut=max(0,lr-hf); full=np.concatenate([y,zz[:,:,cut:,:]],axis=2) |
| full=torch.from_numpy(full) |
| w=torch.ones(T) |
| if s>0: w[:a.overlap]=torch.linspace(0,1,a.overlap) |
| if s+T>total: w[-(s+T-total):]=0 |
| spec[:,:,:,s:s+T]+=full*w; wsum[s:s+T]+=w |
| spec=spec/wsum.clamp_min(1e-8); out=istft(spec).squeeze()[:245760] |
| sf.write(a.output,out.numpy(),SR,subtype='PCM_16'); print('time',time.time()-t0,'shape',out.shape) |
| if __name__=='__main__': main() |
|
|
