(self, inputs, lens=None)
| 57 | |
| 58 | |
| 59 | def forward(self, inputs, lens=None): |
| 60 | specs = self.stft(inputs) |
| 61 | real = specs[:,:self.fft_len//2+1] |
| 62 | imag = specs[:,self.fft_len//2+1:] |
| 63 | spec_mags = torch.sqrt(real**2+imag**2+1e-8) |
| 64 | spec_mags = spec_mags |
| 65 | spec_phase = torch.atan2(imag, real) |
| 66 | spec_phase = spec_phase |
| 67 | |
| 68 | |
| 69 | cspecs = torch.cat([real[:,1:,:],imag[:,1:,:]],1) |
| 70 | #print(cspecs.unsqueeze(-1).shape) |
| 71 | |
| 72 | out = self.ana_conv(cspecs.unsqueeze(-1)).squeeze(-1)#self.cln(cspecs) |
| 73 | real, imag = torch.chunk(out, 2, 1) |
| 74 | out = torch.stack([real, imag],1) |
| 75 | |
| 76 | out = self.enh_block(out) |
| 77 | real = out[:,0,:,:] |
| 78 | imag = out[:,1,:,:] |
| 79 | out = torch.cat([real,imag],1) |
| 80 | out = self.sys_conv(out.unsqueeze(-1)).squeeze(-1) |
| 81 | mask_real, mask_imag = torch.chunk(out, 2, 1) |
| 82 | mask_real = F.pad(mask_real, [0,0,1,0]) |
| 83 | mask_imag = F.pad(mask_imag, [0,0,1,0]) |
| 84 | mask_mags = (mask_real**2+mask_imag**2)**0.5 |
| 85 | real_phase = mask_real/(mask_mags+1e-8) |
| 86 | imag_phase = mask_imag/(mask_mags+1e-8) |
| 87 | mask_phase = torch.atan2( |
| 88 | imag_phase, |
| 89 | real_phase |
| 90 | ) |
| 91 | |
| 92 | #mask_mags = torch.clamp_(mask_mags,0,100) |
| 93 | mask_mags = torch.tanh(mask_mags) |
| 94 | est_mags = mask_mags*spec_mags |
| 95 | est_phase = spec_phase + mask_phase |
| 96 | real = est_mags*torch.cos(est_phase) |
| 97 | imag = est_mags*torch.sin(est_phase) |
| 98 | |
| 99 | out_spec = torch.cat([real, imag], 1) |
| 100 | out_wav = self.istft(out_spec) |
| 101 | |
| 102 | out_wav = torch.squeeze(out_wav, 1) |
| 103 | #out_wav = torch.tanh(out_wav) |
| 104 | out_wav = torch.clamp_(out_wav,-1,1) |
| 105 | return out_spec, out_wav |
| 106 | |
| 107 | |
| 108 | class re_block(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected