update

toolbox/torchaudio/models/spectrum_dfnet/modeling_spectrum_dfnet.py

@@ -818,14 +818,14 @@ class SpectrumDfNet(nn.Module):
         # feat_spec shape: [batch_size, 2, time_steps, df_bins]
         feat_spec = feat_spec.detach()
 
-        # # spec shape: [batch_size, spec_bins, time_steps]
-        # spec = torch.unsqueeze(spec_complex, dim=1)
-        # # spec shape: [batch_size, 1, spec_bins, time_steps]
-        # spec = spec.permute(0, 1, 3, 2)
-        # # spec shape: [batch_size, 1, time_steps, spec_bins]
-        # spec = torch.view_as_real(spec)
-        # # spec shape: [batch_size, 1, time_steps, spec_bins, 2]
-        # spec = spec.detach()
+        # spec shape: [batch_size, spec_bins, time_steps]
+        spec = torch.unsqueeze(spec_complex, dim=1)
+        # spec shape: [batch_size, 1, spec_bins, time_steps]
+        spec = spec.permute(0, 1, 3, 2)
+        # spec shape: [batch_size, 1, time_steps, spec_bins]
+        spec = torch.view_as_real(spec)
+        # spec shape: [batch_size, 1, time_steps, spec_bins, 2]
+        spec = spec.detach()
 
         e0, e1, e2, e3, emb, c0, lsnr, h = self.encoder.forward(feat_power, feat_spec)
 
@@ -834,30 +834,30 @@ class SpectrumDfNet(nn.Module):
         if torch.any(mask > 1) or torch.any(mask < 0):
             raise AssertionError
 
-        # spec_m = self.mask.forward(spec, mask)
-        #
-        # # lsnr shape: [batch_size, time_steps, 1]
-        # lsnr = torch.transpose(lsnr, dim0=2, dim1=1)
-        # # lsnr shape: [batch_size, 1, time_steps]
-        #
-        # df_coefs = self.df_decoder.forward(emb, c0)
-        # df_coefs = self.df_out_transform(df_coefs)
-        # # df_coefs shape: [batch_size, df_order, time_steps, df_bins, 2]
-        #
-        # spec_e = self.df_op.forward(spec.clone(), df_coefs)
-        # # spec_e shape: [batch_size, 1, time_steps, spec_bins, 2]
-        #
-        # spec_e[..., self.df_decoder.df_bins:, :] = spec_m[..., self.df_decoder.df_bins:, :]
-        #
-        # spec_e = torch.squeeze(spec_e, dim=1)
-        # spec_e = spec_e.permute(0, 2, 1, 3)
-        # # spec_e shape: [batch_size, spec_bins, time_steps, 2]
+        spec_m = self.mask.forward(spec, mask)
+
+        # lsnr shape: [batch_size, time_steps, 1]
+        lsnr = torch.transpose(lsnr, dim0=2, dim1=1)
+        # lsnr shape: [batch_size, 1, time_steps]
+
+        df_coefs = self.df_decoder.forward(emb, c0)
+        df_coefs = self.df_out_transform(df_coefs)
+        # df_coefs shape: [batch_size, df_order, time_steps, df_bins, 2]
+
+        spec_e = self.df_op.forward(spec.clone(), df_coefs)
+        # spec_e shape: [batch_size, 1, time_steps, spec_bins, 2]
+
+        spec_e[..., self.df_decoder.df_bins:, :] = spec_m[..., self.df_decoder.df_bins:, :]
+
+        spec_e = torch.squeeze(spec_e, dim=1)
+        spec_e = spec_e.permute(0, 2, 1, 3)
+        # spec_e shape: [batch_size, spec_bins, time_steps, 2]
 
         mask = torch.squeeze(mask, dim=1)
         mask = mask.permute(0, 2, 1)
         # mask shape: [batch_size, spec_bins, time_steps]
 
-        return None, mask, lsnr
+        return spec_e, mask, lsnr
 
 
 class SpectrumDfNetPretrainedModel(SpectrumDfNet):