modify app

app.py

@@ -10,8 +10,10 @@ from config import args
 mastering_transfer = MasteringStyleTransfer(args)
 
 def process_audio(input_audio, reference_audio):
+    input_tensor = mastering_transfer.preprocess_audio(input_audio, args.sample_rate)
+    reference_tensor = mastering_transfer.preprocess_audio(reference_audio, args.sample_rate)
     output_audio, predicted_params, _, _, _, sr = mastering_transfer.process_audio(
-        input_audio, reference_audio, reference_audio, {}, False
+        input_tensor, reference_tensor, reference_tensor, {}, False
     )
     
     param_output = mastering_transfer.get_param_output_string(predicted_params)
@@ -44,16 +46,19 @@ def perform_ito(input_audio, reference_audio, ito_reference_audio, num_steps, op
     
     return "ito_output_mastered.wav", ito_param_output, steps_taken, 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")
+        with gr.Row():
+            input_audio = gr.Audio(label="Input Audio")
+            reference_audio = gr.Audio(label="Reference Audio")
+        
         process_button = gr.Button("Process")
-        output_audio = gr.Audio(label="Output Audio")
-        param_output = gr.Textbox(label="Predicted Parameters", lines=10)
+        
+        with gr.Row():
+            output_audio = gr.Audio(label="Output Audio")
+            param_output = gr.Textbox(label="Predicted Parameters", lines=10)
 
         process_button.click(
             process_audio, 
@@ -62,24 +67,30 @@ with gr.Blocks() as demo:
         )
 
         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)
-        ito_steps_taken = gr.Number(label="ITO Steps Taken")
-        ito_log = gr.Textbox(label="ITO Log", lines=20)
+        with gr.Row():
+            with gr.Column(scale=2):
+                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)
+                ito_steps_taken = gr.Number(label="ITO Steps Taken")
+            
+            with gr.Column(scale=1):
+                ito_log = gr.Textbox(label="ITO Log", lines=30)
 
         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(
+            ito_output, ito_params, steps_taken, log = perform_ito(
                 input_audio, reference_audio, ito_reference_audio, num_steps, optimizer, learning_rate, af_weights
             )
-            return ito_output, ito_params, steps_taken
+            return ito_output, ito_params, steps_taken, log
 
         ito_button.click(
             run_ito,

inference.py

@@ -60,66 +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
-    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
+    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
 
-        # 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)
+            # 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, "\n".join(ito_log)
+        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