update effects inference

my_submission/hdemucs.py

@@ -65,13 +65,14 @@ class HDemucs(IdentitySeparationModel):
             [
                 mixed_sound_array,
                 mixed_diff.unsqueeze(0),
+                mixed_mean.unsqueeze(0),
             ],
             dim=0,
         )
 
         dialog_idx = self.instruments_idx["dialog"]
-        sep_l, sep_r, sep_diff = self.hdemucs(mixed_sound_array).squeeze()
-        dialog = (sep_l[dialog_idx] + sep_r[dialog_idx]) / 2
+        sep_l, sep_r, sep_diff, sep_mean = self.hdemucs(mixed_sound_array).squeeze()
+        dialog = (sep_l[dialog_idx] + sep_r[dialog_idx] + sep_mean[dialog_idx] ) / 3
 
         mixed_mean[0] -= dialog
         separated = self.hdemucs(mixed_mean.unsqueeze(0)).squeeze()
@@ -89,9 +90,9 @@ class HDemucs(IdentitySeparationModel):
 
         effect_idx = self.instruments_idx["effect"]
         music_idx = self.instruments_idx["music"]
-        # effect_diff = mixed_diff[0] - sep_diff[music_idx]
-        effect_l = sep_diff[effect_idx] + separated[effect_idx]
-        effect_r = separated[effect_idx] - sep_diff[effect_idx]
+        effect_diff = mixed_diff[0] - sep_diff[music_idx]
+        effect_l = effect_diff + separated[effect_idx]
+        effect_r = separated[effect_idx] - effect_diff
         # effect_l = 0.5 * (effect_l + sep_l[effect_idx])
         # effect_r = 0.5 * (effect_r + sep_r[effect_idx])