update loss

app.py

@@ -222,10 +222,24 @@ with gr.Blocks() as demo:
             num_steps = gr.Slider(minimum=1, maximum=100, value=10, 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")
             loss_function = gr.Radio(["AudioFeatureLoss", "CLAPFeatureLoss"], label="Loss Function", value="AudioFeatureLoss")
-            clap_target_type = gr.Radio(["Audio", "Text"], label="CLAP Target Type", value="Audio", visible=False)
-            clap_text_prompt = gr.Textbox(label="CLAP Text Prompt", visible=False)
+            
+            # af_weights = gr.Textbox(label="AudioFeatureLoss Weights (comma-separated)", value="0.1,0.001,1.0,1.0,0.1")
+            # clap_target_type = gr.Radio(["Audio", "Text"], label="CLAP Target Type", value="Audio", visible=False)
+            # clap_text_prompt = gr.Textbox(label="CLAP Text Prompt", visible=False)
+
+            # Audio Feature Loss weights
+            with gr.Column(visible=True) as audio_feature_weights:
+                af_weights = gr.Textbox(
+                    label="AudioFeatureLoss Weights (comma-separated)",
+                    value="0.1,0.001,1.0,1.0,0.1",
+                    info="RMS, Crest Factor, Stereo Width, Stereo Imbalance, Bark Spectrum"
+                )
+
+            # CLAP Loss options
+            with gr.Column(visible=False) as clap_options:
+                clap_target_type = gr.Radio(["Audio", "Text"], label="CLAP Target Type", value="Audio")
+                clap_text_prompt = gr.Textbox(label="CLAP Text Prompt", visible=False)
 
     def update_clap_options(loss_function):
         if loss_function == "CLAPFeatureLoss":
@@ -236,9 +250,18 @@ with gr.Blocks() as demo:
     loss_function.change(
         update_clap_options,
         inputs=[loss_function],
-        outputs=[clap_target_type, clap_text_prompt]
+        outputs=[audio_feature_weights, clap_options]
     )
 
+    def update_clap_text_prompt(clap_target_type):
+        return gr.update(visible=clap_target_type == "Text")
+
+    clap_target_type.change(
+        update_clap_text_prompt,
+        inputs=[clap_target_type],
+        outputs=[clap_text_prompt]
+    )
+    
     ito_button = gr.Button("Perform ITO")
 
     with gr.Row():

inference.py

@@ -93,10 +93,10 @@ class MasteringStyleTransfer:
                 losses = af_loss(output_audio, reference_tensor)
             elif ito_config['loss_function'] == 'CLAPFeatureLoss':
                 if ito_config['clap_target_type'] == 'Audio':
-                    target = ito_reference_tensor
+                    target = reference_tensor
                 else:
                     target = ito_config['clap_text_prompt']
-                losses = self.clap_loss(est_targets, target, self.args.sample_rate)
+                losses = self.clap_loss(output_audio, target, self.args.sample_rate)
             total_loss = sum(losses.values())
 
             if total_loss < min_loss: