modify app

app.py

@@ -21,6 +21,21 @@ def denormalize_audio(audio, dtype=np.int16):
     else:
         raise ValueError("Unsupported dtype. Use np.int16 or np.float32.")
 
+def loudness_normalize(audio, sample_rate, target_loudness=-12.0):
+    # Ensure audio is float32
+    if audio.dtype != np.float32:
+        audio = audio.astype(np.float32)
+    
+    # If audio is mono, reshape to (samples, 1)
+    if audio.ndim == 1:
+        audio = audio.reshape(-1, 1)
+    
+    meter = pyln.Meter(sample_rate)  # create BS.1770 meter
+    loudness = meter.integrated_loudness(audio)
+    
+    loudness_normalized_audio = pyln.normalize.loudness(audio, loudness, target_loudness)
+    return loudness_normalized_audio
+
 def process_audio(input_audio, reference_audio):
     output_audio, predicted_params, _, _, _, sr = mastering_transfer.process_audio(
         input_audio, reference_audio, reference_audio, {}, False
@@ -31,6 +46,9 @@ def process_audio(input_audio, reference_audio):
     # Convert output_audio to numpy array if it's a tensor
     if isinstance(output_audio, torch.Tensor):
         output_audio = output_audio.cpu().numpy()
+
+    # Normalize output audio
+    output_audio = loudness_normalize(output_audio, sr)
     
     # Denormalize the audio to int16
     output_audio = denormalize_audio(output_audio, dtype=np.int16)
@@ -78,10 +96,13 @@ def perform_ito(input_audio, reference_audio, ito_reference_audio, num_steps, op
         if isinstance(current_output, torch.Tensor):
             current_output = current_output.detach().cpu().numpy()
         
+        # Normalize output audio
+        current_output = loudness_normalize(current_output, args.sample_rate)
+        
         # Denormalize the audio to int16
         current_output = denormalize_audio(current_output, dtype=np.int16)
         
-        if output_audio.ndim == 1:
+        if current_output.ndim == 1:
             current_output = current_output.reshape(-1, 1)
         elif current_output.ndim > 2:
             current_output = current_output.squeeze()
@@ -163,6 +184,8 @@ with gr.Blocks() as demo:
 demo.launch()
 
 
+
+
 # import gradio as gr
 # import torch
 # import soundfile as sf