update loss

app.py

@@ -94,7 +94,8 @@ def process_audio(input_audio, reference_audio):
 
     return (sr, output_audio), param_output, (sr, normalized_input)
 
-def perform_ito(input_audio, reference_audio, ito_reference_audio, num_steps, optimizer, learning_rate, af_weights):
+# def perform_ito(input_audio, reference_audio, ito_reference_audio, num_steps, optimizer, learning_rate, af_weights):
+def perform_ito(input_audio, reference_audio, ito_reference_audio, num_steps, optimizer, learning_rate, af_weights, loss_function, clap_target_type, clap_text_prompt):
     if ito_reference_audio is None:
         ito_reference_audio = reference_audio
     af_weights = [float(w.strip()) for w in af_weights.split(',')]
@@ -104,7 +105,10 @@ def perform_ito(input_audio, reference_audio, ito_reference_audio, num_steps, op
         'learning_rate': learning_rate,
         'num_steps': num_steps,
         'af_weights': af_weights,
-        'sample_rate': args.sample_rate
+        'sample_rate': args.sample_rate,
+        'loss_function': loss_function,
+        'clap_target_type': clap_target_type,
+        'clap_text_prompt': clap_text_prompt
     }
 
     input_tensor = mastering_transfer.preprocess_audio(input_audio, args.sample_rate)
@@ -219,7 +223,22 @@ with gr.Blocks() as demo:
             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)
+
+    def update_clap_options(loss_function):
+        if loss_function == "CLAPFeatureLoss":
+            return gr.update(visible=True), gr.update(visible=True)
+        else:
+            return gr.update(visible=False), gr.update(visible=False)
+
+    loss_function.change(
+        update_clap_options,
+        inputs=[loss_function],
+        outputs=[clap_target_type, clap_text_prompt]
+    )
+
     ito_button = gr.Button("Perform ITO")
 
     with gr.Row():
@@ -243,7 +262,7 @@ with gr.Blocks() as demo:
 
     ito_button.click(
         perform_ito,
-        inputs=[normalized_input, reference_audio, ito_reference_audio, num_steps, optimizer, learning_rate, af_weights],
+        inputs=[normalized_input, reference_audio, ito_reference_audio, num_steps, optimizer, learning_rate, af_weights, loss_function, clap_target_type, clap_text_prompt],
         outputs=[ito_output_audio, ito_param_output, ito_step_slider, ito_log, ito_loss_plot, all_results]
     ).then(
         update_ito_output,

inference.py

@@ -34,7 +34,14 @@ class MasteringStyleTransfer:
         self.fx_normalizer = Audio_Effects_Normalizer(precomputed_feature_path=args.fx_norm_feature_path, \
                                                         STEMS=['mixture'], \
                                                         EFFECTS=['eq', 'imager', 'loudness'])
-
+        # Loss functions
+        self.af_loss = AudioFeatureLoss(
+            weights=ito_config['af_weights'],
+            sample_rate=ito_config['sample_rate'],
+            stem_separation=False,
+            use_clap=False
+        )
+        self.clap_loss = CLAPFeatureLoss(distance_fn='cosine')
 
     def load_effects_encoder(self):
         effects_encoder = Effects_Encoder(self.args.cfg_enc)
@@ -70,13 +77,6 @@ class MasteringStyleTransfer:
         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
         all_results = []
@@ -87,7 +87,15 @@ class MasteringStyleTransfer:
             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)
+            # Compute loss
+            if ito_config['loss_function'] == 'AudioFeatureLoss':
+                losses = self.af_loss(output_audio, reference_tensor)
+            elif ito_config['loss_function'] == 'CLAPFeatureLoss':
+                if ito_config['clap_target_type'] == 'Audio':
+                    target = ito_reference_tensor
+                else:
+                    target = ito_config['clap_text_prompt']
+                losses = self.clap_loss(est_targets, target, self.args.sample_rate)
             total_loss = sum(losses.values())
 
             if total_loss < min_loss:

modules/loss.py

@@ -196,36 +196,50 @@ class CLAPFeatureLoss(nn.Module):
         else:
             raise ValueError(f"Unsupported distance function: {distance_fn}")
 
-    def forward(self, input_audio, target_audio, sample_rate):
+    def forward(self, input_audio, target, sample_rate):
+        # Process input audio
+        input_embed = self.process_audio(input_audio, sample_rate)
+
+        # Process target (audio or text)
+        if isinstance(target, torch.Tensor):
+            target_embed = self.process_audio(target, sample_rate)
+        elif isinstance(target, str) or (isinstance(target, list) and isinstance(target[0], str)):
+            target_embed = self.process_text(target)
+        else:
+            raise ValueError("Target must be either audio tensor or text (string or list of strings)")
+
+        # Compute loss using the specified distance function
+        loss = self.compute_distance(input_embed, target_embed)
+
+        return loss
+
+    def process_audio(self, audio, sample_rate):
         # Ensure input is in the correct shape (N, C, T)
-        if input_audio.dim() == 2:
-            input_audio = input_audio.unsqueeze(1)
-        if target_audio.dim() == 2:
-            target_audio = target_audio.unsqueeze(1)
+        if audio.dim() == 2:
+            audio = audio.unsqueeze(1)
 
         # Convert to mono if stereo
-        if input_audio.shape[1] > 1:
-            input_audio = input_audio.mean(dim=1, keepdim=True)
-        if target_audio.shape[1] > 1:
-            target_audio = target_audio.mean(dim=1, keepdim=True)
+        if audio.shape[1] > 1:
+            audio = audio.mean(dim=1, keepdim=True)
 
         # Resample if necessary
         if sample_rate != self.target_sample_rate:
-            input_audio = self.resample(input_audio, sample_rate)
-            target_audio = self.resample(target_audio, sample_rate)
+            audio = self.resample(audio, sample_rate)
 
         # Quantize audio data
-        input_audio = self.quantize(input_audio)
-        target_audio = self.quantize(target_audio)
+        audio = self.quantize(audio)
 
         # Get CLAP embeddings
-        input_embed = self.model.get_audio_embedding_from_data(x=input_audio, use_tensor=True)
-        target_embed = self.model.get_audio_embedding_from_data(x=target_audio, use_tensor=True)
+        embed = self.model.get_audio_embedding_from_data(x=audio, use_tensor=True)
+        return embed
 
-        # Compute loss using the specified distance function
-        loss = self.compute_distance(input_embed, target_embed)
-
-        return loss
+    def process_text(self, text):
+        # Get CLAP embeddings for text
+        # ensure input is a list of strings
+        if not isinstance(text, list):
+            text = [text]
+        embed = self.model.get_text_embedding(text, use_tensor=True)
+        return embed
 
     def quantize(self, audio):
         audio = audio.squeeze(1)  # Remove channel dimension
@@ -490,4 +504,19 @@ class AudioFeatureLoss(torch.nn.Module):
                 val = torch.nn.functional.mse_loss(input_transform, target_transform)
                 losses[key] = weight * val * self.source_weights[stem_idx]
 
-        return losses
+        return losses
+
+
+
+if __name__ == "__main__":
+    clap_loss = CLAPFeatureLoss(distance_fn='cosine')
+    
+    input_audio = torch.randn(1, 2, 44100)
+    target_audio = torch.randn(1, 2, 44100)
+    target_text = "This is a test"
+    sample_rate = 44100
+    loss = clap_loss(input_audio, target_audio, sample_rate)
+    print(loss)
+    loss = clap_loss(input_audio, target_text, sample_rate)
+    print(loss)
+