modify app

app.py

@@ -42,8 +42,8 @@ def loudness_normalize(audio, sample_rate, target_loudness=-12.0):
     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
+    output_audio, predicted_params, sr = mastering_transfer.process_audio(
+        input_audio, reference_audio, reference_audio
     )
     
     param_output = mastering_transfer.get_param_output_string(predicted_params)
@@ -54,6 +54,8 @@ def process_audio(input_audio, reference_audio):
 
     # # Normalize output audio
     # output_audio = loudness_normalize(output_audio, sr)
+    print(output_audio.shape)
+    print(f"sr: {sr}")
     
     # Denormalize the audio to int16
     output_audio = denormalize_audio(output_audio, dtype=np.int16)
@@ -66,9 +68,6 @@ def process_audio(input_audio, reference_audio):
     # 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)
-
-    print(output_audio.shape)
-    print(param_output)
     
     return (sr, output_audio), param_output
 

inference.py

@@ -108,7 +108,7 @@ class MasteringStyleTransfer:
             # Log top 5 parameter differences
             if step == 0:
                 initial_params = current_params
-            top_5_diff = self.get_top_5_diff_string(initial_params, current_params)
+            top_5_diff = self.get_top_n_diff_string(initial_params, current_params, top_n=5)
             log_entry = f"Step {step + 1}, Loss: {total_loss.item():.4f}\n{top_5_diff}\n"
 
             if divergence_counter >= 10:
@@ -122,17 +122,6 @@ class MasteringStyleTransfer:
 
         return min_loss_output, min_loss_params, min_loss_embedding, min_loss_step + 1
 
-    def get_top_5_diff_string(self, initial_params, current_params):
-        diff_dict = {}
-        for key in initial_params.keys():
-            diff = abs(current_params[key] - initial_params[key])
-            diff_dict[key] = diff
-        
-        sorted_diff = sorted(diff_dict.items(), key=lambda x: x[1], reverse=True)
-        top_5_diff = sorted_diff[:5]
-        
-        return "\n".join([f"{key}: {value:.4f}" for key, value in top_5_diff])
-
     def preprocess_audio(self, audio, target_sample_rate=44100):
         sample_rate, data = audio
 
@@ -166,7 +155,7 @@ class MasteringStyleTransfer:
 
         return data_tensor.to(self.device)
 
-    def process_audio(self, input_audio, reference_audio, ito_reference_audio, params, perform_ito, log_ito=False):
+    def process_audio(self, input_audio, reference_audio, ito_reference_audio):
         input_tensor = self.preprocess_audio(input_audio, self.args.sample_rate)
         reference_tensor = self.preprocess_audio(reference_audio, self.args.sample_rate)
         ito_reference_tensor = self.preprocess_audio(ito_reference_audio, self.args.sample_rate)
@@ -175,24 +164,7 @@ class MasteringStyleTransfer:
 
         output_audio, predicted_params = self.mastering_style_transfer(input_tensor, reference_feature)
 
-        if perform_ito:
-            ito_log = []
-            for i in range(self.args.max_iter_ito):
-                loss, ito_predicted_params = self.ito_step(input_audio, ito_reference_audio, predicted_params)
-                if log_ito:
-                    top_10_diff = self.get_top_10_diff(predicted_params, ito_predicted_params)
-                    log_entry = f"Iteration {i+1}, Loss: {loss:.4f}\nTop 10 parameter differences:\n{top_10_diff}\n"
-                    ito_log.append(log_entry)
-                predicted_params = ito_predicted_params
-            
-            ito_output_audio = self.converter.convert(input_audio, predicted_params)
-            ito_log = "\n".join(ito_log) if log_ito else None
-        else:
-            ito_output_audio = None
-            ito_predicted_params = None
-            ito_log = None
-
-        return output_audio, predicted_params, ito_output_audio, ito_predicted_params, ito_log, self.args.sample_rate
+        return output_audio, predicted_params, self.args.sample_rate
 
     def print_param_difference(self, initial_params, ito_params):
         all_diffs = []
@@ -278,7 +250,7 @@ class MasteringStyleTransfer:
         
         return "\n".join(output)
 
-    def get_top_10_diff_string(self, initial_params, ito_params):
+    def get_top_n_diff_string(self, initial_params, ito_params, top_n=5):
         if initial_params is None or ito_params is None:
             return "Cannot compare parameters"
         
@@ -299,7 +271,7 @@ class MasteringStyleTransfer:
                 normalized_diff = abs(ito_value - initial_value)
                 all_diffs.append((fx_name, 'width', initial_value.item(), ito_value.item(), normalized_diff.item()))
 
-        top_diffs = sorted(all_diffs, key=lambda x: x[4], reverse=True)[:10]
+        top_diffs = sorted(all_diffs, key=lambda x: x[4], reverse=True)[:top_n]
         
         output = ["Top 10 parameter differences (sorted by normalized difference):"]
         for fx_name, param_name, initial_value, ito_value, normalized_diff in top_diffs:
@@ -322,42 +294,3 @@ def reload_weights(model, ckpt_path, device):
         new_state_dict[name] = v
     model.load_state_dict(new_state_dict, strict=False)
 
-
-if __name__ == "__main__":
-    basis_path = '/data2/tony/Mastering_Style_Transfer/results/dasp_tcn_tuneenc_daspman_loudnessnorm/ckpt/1000/'
-
-    parser = argparse.ArgumentParser(description="Mastering Style Transfer")
-    parser.add_argument("--input_path", type=str, required=True, help="Path to input audio file")
-    parser.add_argument("--reference_path", type=str, required=True, help="Path to reference audio file")
-    parser.add_argument("--ito_reference_path", type=str, required=True, help="Path to ITO reference audio file")
-    parser.add_argument("--model_path", type=str, default=f"{basis_path}dasp_tcn_tuneenc_daspman_loudnessnorm_mastering_converter_1000.pt", help="Path to mastering converter model")
-    parser.add_argument("--encoder_path", type=str, default=f"{basis_path}dasp_tcn_tuneenc_daspman_loudnessnorm_effects_encoder_1000.pt", help="Path to effects encoder model")
-    parser.add_argument("--perform_ito", action="store_true", help="Whether to perform ITO")
-    parser.add_argument("--optimizer", type=str, default="RAdam", help="Optimizer for ITO")
-    parser.add_argument("--learning_rate", type=float, default=0.001, help="Learning rate for ITO")
-    parser.add_argument("--num_steps", type=int, default=100, help="Number of optimization steps for ITO")
-    parser.add_argument("--af_weights", nargs='+', type=float, default=[0.1, 0.001, 1.0, 1.0, 0.1], help="Weights for AudioFeatureLoss")
-    parser.add_argument("--sample_rate", type=int, default=44100, help="Sample rate for AudioFeatureLoss")
-    parser.add_argument("--path_to_config", type=str, default='/home/tony/mastering_transfer/networks/configs.yaml', help="Path to network architecture configuration file")
-
-    args = parser.parse_args()
-    
-    # load network configurations
-    with open(args.path_to_config, 'r') as f:
-        configs = yaml.full_load(f)
-    args.cfg_converter = configs['TCN']['param_mapping']
-    args.cfg_enc = configs['Effects_Encoder']['default']
-
-    ito_config = {
-        'optimizer': args.optimizer,
-        'learning_rate': args.learning_rate,
-        'num_steps': args.num_steps,
-        'af_weights': args.af_weights,
-        'sample_rate': args.sample_rate
-    }
-
-    mastering_style_transfer = MasteringStyleTransfer(args)
-    output_audio, predicted_params, ito_output_audio, ito_predicted_params, optimized_reference_feature, sr, ito_steps = mastering_style_transfer.process_audio(
-        args.input_path, args.reference_path, args.ito_reference_path, ito_config, args.perform_ito
-    )
-