comment out speech rnn

my_submission/ensemble.py

@@ -57,10 +57,10 @@ class EnsembleNet(IdentitySeparationModel):
         )
         self.rnn_music = model.to(self.device)
 
-        model, config = load_checkpoint(
-            self.speech_bandsplitRNN_path, self.bandsplitRNN_ckpt_name
-        )
-        self.rnn_speech = model.to(self.device)
+        # model, config = load_checkpoint(
+        #     self.speech_bandsplitRNN_path, self.bandsplitRNN_ckpt_name
+        # )
+        # self.rnn_speech = model.to(self.device)
 
         n_fft = config["model"]["init_args"]["n_fft"]
         hop_length = config["model"]["init_args"]["hop_length"]
@@ -139,8 +139,9 @@ class EnsembleNet(IdentitySeparationModel):
         # bandsplitRNN
         X = self.spec(mixed_sound_array[:2].transpose(0, 1))
         music_mask = self.rnn_music(X.abs())
-        speech_mask = self.rnn_speech(X.abs())
-        mask = torch.stack([music_mask, speech_mask], dim=1)
+        # speech_mask = self.rnn_speech(X.abs())
+        # mask = torch.stack([music_mask, speech_mask], dim=1)
+        mask = music_mask.unsqueeze(1)
         # speech_spec = X * mask
         speech_spec = self.mwf(mask, X)
         separated = self.inv_spec(speech_spec).squeeze().cpu()
@@ -150,19 +151,17 @@ class EnsembleNet(IdentitySeparationModel):
                 separated, (0, mixed_sound_array.shape[-1] - separated.shape[-1])
             )
 
-        separated_music, separated_speech, separated_fx = separated.permute(
-            0, 2, 1
-        ).numpy()
+        separated_music, separated_fx = separated.permute(0, 2, 1).numpy()
 
         # input_length = len(left_mixed_arr)
         separated_music_arrays["music"] += separated_music
         separated_music_arrays["music"] /= 2
 
-        separated_music_arrays["effect"] += separated_fx
-        separated_music_arrays["effect"] /= 2
+        # separated_music_arrays["effect"] += separated_fx
+        # separated_music_arrays["effect"] /= 2
 
-        separated_music_arrays["dialog"] += separated_speech.mean(axis=1, keepdims=True)
-        separated_music_arrays["dialog"] /= 2
+        # separated_music_arrays["dialog"] += separated_speech.mean(axis=1, keepdims=True)
+        # separated_music_arrays["dialog"] /= 2
 
         return separated_music_arrays, output_sample_rates