modify fx norm

app.py

@@ -63,31 +63,34 @@ def process_audio_with_youtube(input_audio, input_youtube_url, reference_audio,
     
     return process_audio(input_audio, reference_audio)
 
+def to_numpy_audio(audio):
+    # Convert output_audio to numpy array if it's a tensor
+    if isinstance(audio, torch.Tensor):
+        audio = audio.cpu().numpy()
+    # check dimension
+    if audio.ndim == 1:
+        audio = audio.reshape(-1, 1)
+    elif audio.ndim > 2:
+        audio = audio.squeeze()
+    # Ensure the audio is in the correct shape (samples, channels)
+    if audio.shape[1] > audio.shape[0]:
+        audio = audio.transpose(1,0)
+    return audio
+
 def process_audio(input_audio, reference_audio):
     output_audio, predicted_params, sr, normalized_input = mastering_transfer.process_audio(
         input_audio, reference_audio
     )
     
     param_output = mastering_transfer.get_param_output_string(predicted_params)
-    
-    # Convert output_audio to numpy array if it's a tensor
-    if isinstance(output_audio, torch.Tensor):
-        output_audio = output_audio.cpu().numpy()
-
-    if output_audio.ndim == 1:
-        output_audio = output_audio.reshape(-1, 1)
-    elif output_audio.ndim > 2:
-        output_audio = output_audio.squeeze()
-
-    # Ensure the audio is in the correct shape (samples, channels)
-    if output_audio.shape[1] > output_audio.shape[0]:
-        output_audio = output_audio.transpose(1,0)
 
+    # Convert to numpy audio
+    output_audio = to_numpy_audio(output_audio)
+    normalized_input = to_numpy_audio(normalized_input)
     # Normalize output audio
-    output_audio = loudness_normalize(output_audio, sr)    
+    output_audio = loudness_normalize(output_audio, sr)
     # Denormalize the audio to int16
     output_audio = denormalize_audio(output_audio, dtype=np.int16)
-    normalized_input = denormalize_audio(normalized_input, dtype=np.int16)
 
     return (sr, output_audio), param_output, (sr, normalized_input)
 
@@ -125,18 +128,8 @@ def perform_ito(input_audio, reference_audio, ito_reference_audio, num_steps, op
     current_output = last_result['audio']
     ito_param_output = mastering_transfer.get_param_output_string(last_result['params'])
 
-    # Convert current_output to numpy array if it's a tensor
-    if isinstance(current_output, torch.Tensor):
-        current_output = current_output.cpu().numpy()
-
-    if current_output.ndim == 1:
-        current_output = current_output.reshape(-1, 1)
-    elif current_output.ndim > 2:
-        current_output = current_output.squeeze()
-    # Ensure the audio is in the correct shape (samples, channels)
-    if current_output.shape[1] > current_output.shape[0]:
-        current_output = current_output.transpose(1,0)
-
+    # Convert to numpy audio
+    current_output = to_numpy_audio(current_output)
     # Loudness normalize output audio
     current_output = loudness_normalize(current_output, args.sample_rate)
     # Denormalize the audio to int16
@@ -149,18 +142,8 @@ def update_ito_output(all_results, selected_step):
     current_output = selected_result['audio']
     ito_param_output = mastering_transfer.get_param_output_string(selected_result['params'])
 
-    # Convert current_output to numpy array if it's a tensor
-    if isinstance(current_output, torch.Tensor):
-        current_output = current_output.cpu().numpy()
-
-    if current_output.ndim == 1:
-        current_output = current_output.reshape(-1, 1)
-    elif current_output.ndim > 2:
-        current_output = current_output.squeeze()
-    # Ensure the audio is in the correct shape (samples, channels)
-    if current_output.shape[1] > current_output.shape[0]:
-        current_output = current_output.transpose(1,0)
-
+    # Convert to numpy audio
+    current_output = to_numpy_audio(current_output)
     # Loudness normalize output audio
     current_output = loudness_normalize(current_output, args.sample_rate)
     # Denormalize the audio to int16

inference.py

@@ -153,8 +153,6 @@ class MasteringStyleTransfer:
     def process_audio(self, input_audio, reference_audio):
         input_tensor = self.preprocess_audio(input_audio, self.args.sample_rate, normalize=True)
         reference_tensor = self.preprocess_audio(reference_audio, self.args.sample_rate)
-        print(f"input_tensor: {input_tensor.shape}")
-        print(f"reference_tensor: {reference_tensor.shape}")
 
         reference_feature = self.get_reference_embedding(reference_tensor)