Update inference/style_transfer.py

inference/style_transfer.py

@@ -32,7 +32,8 @@ class Mixing_Style_Transfer_Inference:
             self.device = torch.device("cuda:0")
         else:
             self.device = torch.device("cpu")
-        
+        print(f"using device: {self.device} for inference")
+            
         # inference computational hyperparameters
         self.args = args
         self.segment_length = args.segment_length
@@ -176,13 +177,14 @@ class Mixing_Style_Transfer_Inference:
             fin_data_out_mix = sum(inst_outputs)
             
             # loudness adjusting for mastering purpose
-            meter = pyloudnorm.Meter(44100)
-            loudness_out = meter.integrated_loudness(fin_data_out_mix.transpose(-1, -2))
-            reference_aud = load_wav_segment(reference_track_path, axis=1)
-            loudness_ref = meter.integrated_loudness(reference_aud)
-            # adjust output loudness to that of the reference
-            fin_data_out_mix = pyloudnorm.normalize.loudness(fin_data_out_mix, loudness_out, loudness_ref)
-            fin_data_out_mix = np.clip(fin_data_out_mix, -1., 1.)
+            if self.args.match_output_loudness:
+                meter = pyloudnorm.Meter(44100)
+                loudness_out = meter.integrated_loudness(fin_data_out_mix.transpose(-1, -2))
+                reference_aud = load_wav_segment(reference_track_path, axis=1)
+                loudness_ref = meter.integrated_loudness(reference_aud)
+                # adjust output loudness to that of the reference
+                fin_data_out_mix = pyloudnorm.normalize.loudness(fin_data_out_mix, loudness_out, loudness_ref)
+                fin_data_out_mix = np.clip(fin_data_out_mix, -1., 1.)
             
             # save output
             fin_output_path = os.path.join(cur_out_dir, f"mixture_{output_name_tag}.wav")
@@ -382,10 +384,11 @@ def set_up(start_point_in_second=0, duration_in_second=30):
     # FX normalization
     inference_args.add_argument('--normalize_input', type=str2bool, default=True)
     inference_args.add_argument('--normalization_order', type=str2bool, default=['loudness', 'eq', 'compression', 'imager', 'loudness']) # Effects to be normalized, order matters
+    inference_args.add_argument('--match_output_loudness', type=str2bool, default=False)
     # interpolation
     inference_args.add_argument('--interpolation', type=str2bool, default=False)
     inference_args.add_argument('--interpolate_segments', type=int, default=30)
-
+    
     device_args = parser.add_argument_group('Device args')
     device_args.add_argument('--workers', type=int, default=1)
     device_args.add_argument('--batch_size', type=int, default=1)   # for processing long audio