modify app

inference.py

@@ -105,12 +105,11 @@ class MasteringStyleTransfer:
             else:
                 divergence_counter = 0
 
-            # Log top 10 parameter differences
+            # Log top 5 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"
-            yield log_entry, output_audio, current_params, step + 1
+            top_5_diff = self.get_top_5_diff_string(initial_params, current_params)
+            log_entry = f"Step {step + 1}, Loss: {total_loss.item():.4f}\n{top_5_diff}\n"
 
             if divergence_counter >= 10:
                 print(f"Optimization stopped early due to divergence at step {step}")
@@ -119,8 +118,21 @@ class MasteringStyleTransfer:
             total_loss.backward()
             optimizer.step()
 
+            yield log_entry, output_audio.detach(), current_params, step + 1, total_loss.item()
+
         return min_loss_output, min_loss_params, min_loss_embedding, min_loss_step + 1
 
+    def get_top_5_diff_string(self, initial_params, current_params):
+        diff_dict = {}
+        for key in initial_params.keys():
+            diff = abs(current_params[key] - initial_params[key])
+            diff_dict[key] = diff
+        
+        sorted_diff = sorted(diff_dict.items(), key=lambda x: x[1], reverse=True)
+        top_5_diff = sorted_diff[:5]
+        
+        return "\n".join([f"{key}: {value:.4f}" for key, value in top_5_diff])
+
     def preprocess_audio(self, audio, target_sample_rate=44100):
         sample_rate, data = audio