update loss

app.py

@@ -94,8 +94,7 @@ 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, loss_function, clap_target_type, clap_text_prompt):
+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, clap_distance_fn):
     if ito_reference_audio is None:
         ito_reference_audio = reference_audio
     af_weights = [float(w.strip()) for w in af_weights.split(',')]
@@ -108,7 +107,8 @@ def perform_ito(input_audio, reference_audio, ito_reference_audio, num_steps, op
         'sample_rate': args.sample_rate,
         'loss_function': loss_function,
         'clap_target_type': clap_target_type,
-        'clap_text_prompt': clap_text_prompt
+        'clap_text_prompt': clap_text_prompt,
+        'clap_distance_fn': clap_distance_fn
     }
 
     input_tensor = mastering_transfer.preprocess_audio(input_audio, args.sample_rate)
@@ -163,7 +163,7 @@ with gr.Blocks() as demo:
         gr.Markdown("Interactive demo of Inference Time Optimization (ITO) for Music Mastering Style Transfer. \
                     The mastering style transfer is performed by a differentiable audio processing model, and the predicted parameters are shown as the output. \
                     Perform mastering style transfer with an input source audio and a reference mastering style audio. On top of this result, you can perform ITO to optimize the reference embedding $z_{ref}$ to further gain control over the output mastering style.")
-        gr.Image("ito_snow.png", width=300)
+        gr.Image("ito_snow.png", width=100, label="ITO pipeline")
 
     gr.Markdown("## Step 1: Mastering Style Transfer")
 
@@ -219,14 +219,10 @@ with gr.Blocks() as demo:
     with gr.Row():
         ito_reference_audio = gr.Audio(label="ITO Reference Style Audio $x'_{ref}$ (optional)")
         with gr.Column():
-            num_steps = gr.Slider(minimum=1, maximum=100, value=10, step=1, label="Number of Steps")
+            num_steps = gr.Slider(minimum=1, maximum=100, value=10, step=1, label="Number of Steps for additional optimization")
             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")
             loss_function = gr.Radio(["AudioFeatureLoss", "CLAPFeatureLoss"], label="Loss Function", value="AudioFeatureLoss")
-            
-            # 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:
@@ -240,6 +236,7 @@ with gr.Blocks() as demo:
             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)
+                clap_distance_fn = gr.Dropdown([ "cosine", "mse", "l1"], label="CLAP Distance Function", value="cosine")
 
     def update_clap_options(loss_function):
         if loss_function == "CLAPFeatureLoss":
@@ -285,7 +282,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, loss_function, clap_target_type, clap_text_prompt],
+        inputs=[normalized_input, reference_audio, ito_reference_audio, num_steps, optimizer, learning_rate, af_weights, loss_function, clap_target_type, clap_text_prompt, clap_distance_fn],
         outputs=[ito_output_audio, ito_param_output, ito_step_slider, ito_log, ito_loss_plot, all_results]
     ).then(
         update_ito_output,

inference.py

@@ -35,7 +35,7 @@ class MasteringStyleTransfer:
                                                         STEMS=['mixture'], \
                                                         EFFECTS=['eq', 'imager', 'loudness'])
         # Loss functions
-        self.clap_loss = CLAPFeatureLoss(distance_fn='cosine')
+        self.clap_loss = CLAPFeatureLoss()
 
     def load_effects_encoder(self):
         effects_encoder = Effects_Encoder(self.args.cfg_enc)
@@ -97,7 +97,7 @@ class MasteringStyleTransfer:
                     target = reference_tensor
                 else:
                     target = ito_config['clap_text_prompt']
-                losses = self.clap_loss(output_audio, target, self.args.sample_rate)
+                losses = self.clap_loss(output_audio, target, self.args.sample_rate, distance_fn=ito_config['clap_distance_fn'])
                 total_loss = losses
 
             if total_loss < min_loss:

modules/loss.py

@@ -180,23 +180,13 @@ import laion_clap
 import torchaudio
 # CLAP feature loss
 class CLAPFeatureLoss(nn.Module):
-    def __init__(self, distance_fn='mse'):
+    def __init__(self):
         super(CLAPFeatureLoss, self).__init__()
         self.target_sample_rate = 48000  # CLAP expects 48kHz audio
         self.model = laion_clap.CLAP_Module(enable_fusion=False)
         self.model.load_ckpt()  # download the default pretrained checkpoint
-        
-        self.distance_fn = distance_fn
-        if distance_fn == 'mse':
-            self.compute_distance = F.mse_loss
-        elif distance_fn == 'l1':
-            self.compute_distance = F.l1_loss
-        elif distance_fn == 'cosine':
-            self.compute_distance = lambda x, y: 1 - F.cosine_similarity(x, y).mean()
-        else:
-            raise ValueError(f"Unsupported distance function: {distance_fn}")
 
-    def forward(self, input_audio, target, sample_rate):
+    def forward(self, input_audio, target, sample_rate, distance_fn='cosine'):
         # Process input audio
         input_embed = self.process_audio(input_audio, sample_rate)
 
@@ -209,7 +199,7 @@ class CLAPFeatureLoss(nn.Module):
             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)
+        loss = self.compute_distance(input_embed, target_embed, distance_fn)
 
         return loss
 
@@ -230,7 +220,8 @@ class CLAPFeatureLoss(nn.Module):
         audio = self.quantize(audio)
 
         # Get CLAP embeddings
-        embed = self.model.get_audio_embedding_from_data(x=audio, use_tensor=True)
+        with torch.no_grad():
+            embed = self.model.get_audio_embedding_from_data(x=audio, use_tensor=True)
         return embed
 
     def process_text(self, text):
@@ -238,18 +229,29 @@ class CLAPFeatureLoss(nn.Module):
         # ensure input is a list of strings
         if not isinstance(text, list):
             text = [text]
-        embed = self.model.get_text_embedding(text, use_tensor=True)
+        with torch.no_grad():
+            embed = self.model.get_text_embedding(text, use_tensor=True)
         return embed
 
+    def compute_distance(self, x, y, distance_fn):
+        if distance_fn == 'mse':
+            return F.mse_loss(x, y)
+        elif distance_fn == 'l1':
+            return F.l1_loss(x, y)
+        elif distance_fn == 'cosine':
+            return 1 - F.cosine_similarity(x, y).mean()
+        else:
+            raise ValueError(f"Unsupported distance function: {distance_fn}")
+
     def quantize(self, audio):
         audio = audio.squeeze(1)  # Remove channel dimension
         audio = torch.clamp(audio, -1.0, 1.0)
         audio = (audio * 32767.0).to(torch.int16).to(torch.float32) / 32767.0
         return audio
 
-    def resample(self, audio, sample_rate):
+    def resample(self, audio, input_sample_rate):
         resampler = torchaudio.transforms.Resample(
-            orig_freq=sample_rate, new_freq=self.target_sample_rate
+            orig_freq=input_sample_rate, new_freq=self.target_sample_rate
         ).to(audio.device)
         return resampler(audio)
 
@@ -506,18 +508,3 @@ class AudioFeatureLoss(torch.nn.Module):
 
         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)
-
-    print(loss.item())