modify app

app.py

@@ -9,39 +9,40 @@ from config import args
 
 mastering_transfer = MasteringStyleTransfer(args)
 
-def process_audio(input_audio, reference_audio, perform_ito, ito_reference_audio=None):
-    # Process the audio files
-    output_audio, predicted_params, ito_output_audio, ito_predicted_params, ito_log, sr = mastering_transfer.process_audio(
-        input_audio, reference_audio, ito_reference_audio if ito_reference_audio else reference_audio, {}, perform_ito
+def process_audio(input_audio, reference_audio):
+    output_audio, predicted_params, _, _, _, sr = mastering_transfer.process_audio(
+        input_audio, reference_audio, reference_audio, {}, False
     )
     
-    # Generate parameter output strings
     param_output = mastering_transfer.get_param_output_string(predicted_params)
-    ito_param_output = mastering_transfer.get_param_output_string(ito_predicted_params) if ito_predicted_params is not None else "ITO not performed"
     
-    # Generate top 10 differences if ITO was performed
-    top_10_diff = mastering_transfer.get_top_10_diff_string(predicted_params, ito_predicted_params) if ito_predicted_params is not None else "ITO not performed"
-    
-    return "output_mastered.wav", "ito_output_mastered.wav" if ito_output_audio is not None else None, param_output, ito_param_output, top_10_diff, ito_log
+    return "output_mastered.wav", param_output
 
-def process_with_ito(input_audio, reference_audio, perform_ito, use_same_reference, ito_reference_audio):
-    ito_ref = reference_audio if use_same_reference else ito_reference_audio
-    return process_audio(input_audio, reference_audio, perform_ito, ito_ref)
+def perform_ito(input_audio, reference_audio, ito_reference_audio, num_steps, optimizer, learning_rate, af_weights):
+    if ito_reference_audio is None:
+        ito_reference_audio = reference_audio
 
-def process_youtube_with_ito(input_url, reference_url, perform_ito, use_same_reference, ito_reference_url):
-    input_audio = download_youtube_audio(input_url)
-    reference_audio = download_youtube_audio(reference_url)
-    ito_ref = reference_audio if use_same_reference else download_youtube_audio(ito_reference_url)
-    
-    output_audio, predicted_params, ito_output_audio, ito_predicted_params, ito_log, sr = mastering_transfer.process_audio(
-        input_audio, reference_audio, ito_ref, {}, perform_ito, log_ito=True
+    ito_config = {
+        'optimizer': optimizer,
+        'learning_rate': learning_rate,
+        'num_steps': num_steps,
+        'af_weights': af_weights,
+        'sample_rate': args.sample_rate
+    }
+
+    input_tensor = mastering_transfer.preprocess_audio(input_audio, args.sample_rate)
+    reference_tensor = mastering_transfer.preprocess_audio(reference_audio, args.sample_rate)
+    ito_reference_tensor = mastering_transfer.preprocess_audio(ito_reference_audio, args.sample_rate)
+
+    initial_reference_feature = mastering_transfer.get_reference_embedding(reference_tensor)
+
+    ito_output, ito_params, optimized_embedding, steps_taken, ito_log = mastering_transfer.inference_time_optimization(
+        input_tensor, ito_reference_tensor, ito_config, initial_reference_feature
     )
+
+    ito_param_output = mastering_transfer.get_param_output_string(ito_params)
     
-    param_output = mastering_transfer.get_param_output_string(predicted_params)
-    ito_param_output = mastering_transfer.get_param_output_string(ito_predicted_params) if ito_predicted_params is not None else "ITO not performed"
-    top_10_diff = mastering_transfer.get_top_10_diff_string(predicted_params, ito_predicted_params) if ito_predicted_params is not None else "ITO not performed"
-    
-    return "output_mastered_yt.wav", "ito_output_mastered_yt.wav" if ito_output_audio is not None else None, param_output, ito_param_output, top_10_diff, ito_log
+    return "ito_output_mastered.wav", ito_param_output, steps_taken, ito_log
 
 
 with gr.Blocks() as demo:
@@ -50,63 +51,154 @@ with gr.Blocks() as demo:
     with gr.Tab("Upload Audio"):
         input_audio = gr.Audio(label="Input Audio")
         reference_audio = gr.Audio(label="Reference Audio")
-        perform_ito = gr.Checkbox(label="Perform ITO")
-        with gr.Column(visible=False) as ito_options:
-            use_same_reference = gr.Checkbox(label="Use same reference audio for ITO", value=True)
-            ito_reference_audio = gr.Audio(label="ITO Reference Audio", visible=False)
-        
-        def update_ito_options(perform_ito):
-            return gr.Column.update(visible=perform_ito)
-        
-        def update_ito_reference(use_same):
-            return gr.Audio.update(visible=not use_same)
-        
-        perform_ito.change(fn=update_ito_options, inputs=perform_ito, outputs=ito_options)
-        use_same_reference.change(fn=update_ito_reference, inputs=use_same_reference, outputs=ito_reference_audio)
-        
-        submit_button = gr.Button("Process")
+        process_button = gr.Button("Process")
         output_audio = gr.Audio(label="Output Audio")
-        ito_output_audio = gr.Audio(label="ITO Output Audio")
         param_output = gr.Textbox(label="Predicted Parameters", lines=10)
+
+        process_button.click(
+            process_audio, 
+            inputs=[input_audio, reference_audio], 
+            outputs=[output_audio, param_output]
+        )
+
+        gr.Markdown("## Inference Time Optimization (ITO)")
+        ito_reference_audio = gr.Audio(label="ITO Reference Audio (optional)")
+        num_steps = gr.Slider(minimum=1, maximum=1000, value=100, step=1, label="Number of Steps")
+        optimizer = gr.Dropdown(["Adam", "RAdam", "SGD"], value="RAdam", label="Optimizer")
+        learning_rate = gr.Slider(minimum=0.0001, maximum=0.1, value=0.001, step=0.0001, label="Learning Rate")
+        af_weights = gr.Textbox(label="AudioFeatureLoss Weights (comma-separated)", value="0.1,0.001,1.0,1.0,0.1")
+        
+        ito_button = gr.Button("Perform ITO")
+        ito_output_audio = gr.Audio(label="ITO Output Audio")
         ito_param_output = gr.Textbox(label="ITO Predicted Parameters", lines=10)
-        top_10_diff = gr.Textbox(label="Top 10 Parameter Differences", lines=10)
+        ito_steps_taken = gr.Number(label="ITO Steps Taken")
         ito_log = gr.Textbox(label="ITO Log", lines=20)
-        
-        submit_button.click(
-            process_with_ito, 
-            inputs=[input_audio, reference_audio, perform_ito, use_same_reference, ito_reference_audio], 
-            outputs=[output_audio, ito_output_audio, param_output, ito_param_output, top_10_diff, ito_log]
+
+        def run_ito(input_audio, reference_audio, ito_reference_audio, num_steps, optimizer, learning_rate, af_weights):
+            af_weights = [float(w.strip()) for w in af_weights.split(',')]
+            ito_output, ito_params, steps_taken = perform_ito(
+                input_audio, reference_audio, ito_reference_audio, num_steps, optimizer, learning_rate, af_weights
+            )
+            return ito_output, ito_params, steps_taken
+
+        ito_button.click(
+            run_ito,
+            inputs=[input_audio, reference_audio, ito_reference_audio, num_steps, optimizer, learning_rate, af_weights],
+            outputs=[ito_output_audio, ito_param_output, ito_steps_taken, ito_log]
         )
 
-    with gr.Tab("YouTube URLs"):
-        input_url = gr.Textbox(label="Input YouTube URL")
-        reference_url = gr.Textbox(label="Reference YouTube URL")
-        perform_ito_yt = gr.Checkbox(label="Perform ITO")
-        with gr.Column(visible=False) as ito_options_yt:
-            use_same_reference_yt = gr.Checkbox(label="Use same reference audio for ITO", value=True)
-            ito_reference_url = gr.Textbox(label="ITO Reference YouTube URL", visible=False)
+demo.launch()
+
+
+# import gradio as gr
+# import torch
+# import soundfile as sf
+# import numpy as np
+# import yaml
+# from inference import MasteringStyleTransfer
+# from utils import download_youtube_audio
+# from config import args
+
+# mastering_transfer = MasteringStyleTransfer(args)
+
+# def process_audio(input_audio, reference_audio, perform_ito, ito_reference_audio=None):
+#     # Process the audio files
+#     output_audio, predicted_params, ito_output_audio, ito_predicted_params, ito_log, sr = mastering_transfer.process_audio(
+#         input_audio, reference_audio, ito_reference_audio if ito_reference_audio else reference_audio, {}, perform_ito
+#     )
+    
+#     # Generate parameter output strings
+#     param_output = mastering_transfer.get_param_output_string(predicted_params)
+#     ito_param_output = mastering_transfer.get_param_output_string(ito_predicted_params) if ito_predicted_params is not None else "ITO not performed"
+    
+#     # Generate top 10 differences if ITO was performed
+#     top_10_diff = mastering_transfer.get_top_10_diff_string(predicted_params, ito_predicted_params) if ito_predicted_params is not None else "ITO not performed"
+    
+#     return "output_mastered.wav", "ito_output_mastered.wav" if ito_output_audio is not None else None, param_output, ito_param_output, top_10_diff, ito_log
+
+# def process_with_ito(input_audio, reference_audio, perform_ito, use_same_reference, ito_reference_audio):
+#     ito_ref = reference_audio if use_same_reference else ito_reference_audio
+#     return process_audio(input_audio, reference_audio, perform_ito, ito_ref)
+
+# def process_youtube_with_ito(input_url, reference_url, perform_ito, use_same_reference, ito_reference_url):
+#     input_audio = download_youtube_audio(input_url)
+#     reference_audio = download_youtube_audio(reference_url)
+#     ito_ref = reference_audio if use_same_reference else download_youtube_audio(ito_reference_url)
+    
+#     output_audio, predicted_params, ito_output_audio, ito_predicted_params, ito_log, sr = mastering_transfer.process_audio(
+#         input_audio, reference_audio, ito_ref, {}, perform_ito, log_ito=True
+#     )
+    
+#     param_output = mastering_transfer.get_param_output_string(predicted_params)
+#     ito_param_output = mastering_transfer.get_param_output_string(ito_predicted_params) if ito_predicted_params is not None else "ITO not performed"
+#     top_10_diff = mastering_transfer.get_top_10_diff_string(predicted_params, ito_predicted_params) if ito_predicted_params is not None else "ITO not performed"
+    
+#     return "output_mastered_yt.wav", "ito_output_mastered_yt.wav" if ito_output_audio is not None else None, param_output, ito_param_output, top_10_diff, ito_log
+
+
+# with gr.Blocks() as demo:
+#     gr.Markdown("# Mastering Style Transfer Demo")
+
+#     with gr.Tab("Upload Audio"):
+#         input_audio = gr.Audio(label="Input Audio")
+#         reference_audio = gr.Audio(label="Reference Audio")
+#         perform_ito = gr.Checkbox(label="Perform ITO")
+#         with gr.Column(visible=False) as ito_options:
+#             use_same_reference = gr.Checkbox(label="Use same reference audio for ITO", value=True)
+#             ito_reference_audio = gr.Audio(label="ITO Reference Audio", visible=False)
         
-        def update_ito_options_yt(perform_ito):
-            return gr.Column.update(visible=perform_ito)
+#         def update_ito_options(perform_ito):
+#             return gr.Column.update(visible=perform_ito)
         
-        def update_ito_reference_yt(use_same):
-            return gr.Textbox.update(visible=not use_same)
+#         def update_ito_reference(use_same):
+#             return gr.Audio.update(visible=not use_same)
         
-        perform_ito_yt.change(fn=update_ito_options_yt, inputs=perform_ito_yt, outputs=ito_options_yt)
-        use_same_reference_yt.change(fn=update_ito_reference_yt, inputs=use_same_reference_yt, outputs=ito_reference_url)
+#         perform_ito.change(fn=update_ito_options, inputs=perform_ito, outputs=ito_options)
+#         use_same_reference.change(fn=update_ito_reference, inputs=use_same_reference, outputs=ito_reference_audio)
         
-        submit_button_yt = gr.Button("Process")
-        output_audio_yt = gr.Audio(label="Output Audio")
-        ito_output_audio_yt = gr.Audio(label="ITO Output Audio")
-        param_output_yt = gr.Textbox(label="Predicted Parameters", lines=10)
-        ito_param_output_yt = gr.Textbox(label="ITO Predicted Parameters", lines=10)
-        top_10_diff_yt = gr.Textbox(label="Top 10 Parameter Differences", lines=10)
-        ito_log_yt = gr.Textbox(label="ITO Log", lines=20)
-
-        submit_button_yt.click(
-            process_youtube_with_ito, 
-            inputs=[input_url, reference_url, perform_ito_yt, use_same_reference_yt, ito_reference_url], 
-            outputs=[output_audio_yt, ito_output_audio_yt, param_output_yt, ito_param_output_yt, top_10_diff_yt, ito_log_yt]
-        )
+#         submit_button = gr.Button("Process")
+#         output_audio = gr.Audio(label="Output Audio")
+#         ito_output_audio = gr.Audio(label="ITO Output Audio")
+#         param_output = gr.Textbox(label="Predicted Parameters", lines=10)
+#         ito_param_output = gr.Textbox(label="ITO Predicted Parameters", lines=10)
+#         top_10_diff = gr.Textbox(label="Top 10 Parameter Differences", lines=10)
+#         ito_log = gr.Textbox(label="ITO Log", lines=20)
+        
+#         submit_button.click(
+#             process_with_ito, 
+#             inputs=[input_audio, reference_audio, perform_ito, use_same_reference, ito_reference_audio], 
+#             outputs=[output_audio, ito_output_audio, param_output, ito_param_output, top_10_diff, ito_log]
+#         )
+
+#     with gr.Tab("YouTube URLs"):
+#         input_url = gr.Textbox(label="Input YouTube URL")
+#         reference_url = gr.Textbox(label="Reference YouTube URL")
+#         perform_ito_yt = gr.Checkbox(label="Perform ITO")
+#         with gr.Column(visible=False) as ito_options_yt:
+#             use_same_reference_yt = gr.Checkbox(label="Use same reference audio for ITO", value=True)
+#             ito_reference_url = gr.Textbox(label="ITO Reference YouTube URL", visible=False)
+        
+#         def update_ito_options_yt(perform_ito):
+#             return gr.Column.update(visible=perform_ito)
+        
+#         def update_ito_reference_yt(use_same):
+#             return gr.Textbox.update(visible=not use_same)
+        
+#         perform_ito_yt.change(fn=update_ito_options_yt, inputs=perform_ito_yt, outputs=ito_options_yt)
+#         use_same_reference_yt.change(fn=update_ito_reference_yt, inputs=use_same_reference_yt, outputs=ito_reference_url)
+        
+#         submit_button_yt = gr.Button("Process")
+#         output_audio_yt = gr.Audio(label="Output Audio")
+#         ito_output_audio_yt = gr.Audio(label="ITO Output Audio")
+#         param_output_yt = gr.Textbox(label="Predicted Parameters", lines=10)
+#         ito_param_output_yt = gr.Textbox(label="ITO Predicted Parameters", lines=10)
+#         top_10_diff_yt = gr.Textbox(label="Top 10 Parameter Differences", lines=10)
+#         ito_log_yt = gr.Textbox(label="ITO Log", lines=20)
+
+#         submit_button_yt.click(
+#             process_youtube_with_ito, 
+#             inputs=[input_url, reference_url, perform_ito_yt, use_same_reference_yt, ito_reference_url], 
+#             outputs=[output_audio_yt, ito_output_audio_yt, param_output_yt, ito_param_output_yt, top_10_diff_yt, ito_log_yt]
+#         )
     
-demo.launch()
+# demo.launch()

inference.py

@@ -60,59 +60,66 @@ class MasteringStyleTransfer:
             predicted_params = self.mastering_converter.get_last_predicted_params()
         return output_audio, predicted_params
 
-    def inference_time_optimization(self, input_tensor, reference_tensor, ito_config, initial_reference_feature):
-        fit_embedding = torch.nn.Parameter(initial_reference_feature)
-        optimizer = getattr(torch.optim, ito_config['optimizer'])([fit_embedding], lr=ito_config['learning_rate'])
-    
-        af_loss = AudioFeatureLoss(
-            weights=ito_config['af_weights'],
-            sample_rate=ito_config['sample_rate'],
-            stem_separation=False,
-            use_clap=False
-        )
-
-        min_loss = float('inf')
-        min_loss_step = 0
-        min_loss_output = None
-        min_loss_params = None
-        min_loss_embedding = None
-
-        loss_history = []
-        divergence_counter = 0
-
-        for step in range(ito_config['num_steps']):
-            optimizer.zero_grad()
-
-            output_audio = self.mastering_converter(input_tensor, fit_embedding)
-
-            losses = af_loss(output_audio, reference_tensor)
-            total_loss = sum(losses.values())
-
-            loss_history.append(total_loss.item())
-
-            if total_loss < min_loss:
-                min_loss = total_loss.item()
-                min_loss_step = step
-                min_loss_output = output_audio.detach()
-                min_loss_params = self.mastering_converter.get_last_predicted_params()
-                min_loss_embedding = fit_embedding.detach().clone()
-
-            # Check for divergence
-            if len(loss_history) > 10 and total_loss > loss_history[-11]:
-                divergence_counter += 1
-            else:
-                divergence_counter = 0
+def inference_time_optimization(self, input_tensor, reference_tensor, ito_config, initial_reference_feature):
+    fit_embedding = torch.nn.Parameter(initial_reference_feature)
+    optimizer = getattr(torch.optim, ito_config['optimizer'])([fit_embedding], lr=ito_config['learning_rate'])
+
+    af_loss = AudioFeatureLoss(
+        weights=ito_config['af_weights'],
+        sample_rate=ito_config['sample_rate'],
+        stem_separation=False,
+        use_clap=False
+    )
+
+    min_loss = float('inf')
+    min_loss_step = 0
+    min_loss_output = None
+    min_loss_params = None
+    min_loss_embedding = None
+
+    loss_history = []
+    divergence_counter = 0
+    ito_log = []
+
+    for step in range(ito_config['num_steps']):
+        optimizer.zero_grad()
+
+        output_audio = self.mastering_converter(input_tensor, fit_embedding)
+        current_params = self.mastering_converter.get_last_predicted_params()
+
+        losses = af_loss(output_audio, reference_tensor)
+        total_loss = sum(losses.values())
+
+        loss_history.append(total_loss.item())
+
+        if total_loss < min_loss:
+            min_loss = total_loss.item()
+            min_loss_step = step
+            min_loss_output = output_audio.detach()
+            min_loss_params = current_params
+            min_loss_embedding = fit_embedding.detach().clone()
+
+        # Check for divergence
+        if len(loss_history) > 10 and total_loss > loss_history[-11]:
+            divergence_counter += 1
+        else:
+            divergence_counter = 0
 
-            print(total_loss, min_loss)
+        # Log top 10 parameter differences
+        if step == 0:
+            initial_params = current_params
+        top_10_diff = self.get_top_10_diff_string(initial_params, current_params)
+        log_entry = f"Step {step + 1}, Loss: {total_loss.item():.4f}\n{top_10_diff}\n"
+        ito_log.append(log_entry)
 
-            if divergence_counter >= 10:
-                print(f"Optimization stopped early due to divergence at step {step}")
-                break
+        if divergence_counter >= 10:
+            print(f"Optimization stopped early due to divergence at step {step}")
+            break
 
-            total_loss.backward()
-            optimizer.step()
+        total_loss.backward()
+        optimizer.step()
 
-        return min_loss_output, min_loss_params, min_loss_embedding, min_loss_step + 1
+    return min_loss_output, min_loss_params, min_loss_embedding, min_loss_step + 1, "\n".join(ito_log)
 
     def preprocess_audio(self, audio, target_sample_rate=44100):
         sample_rate, data = audio