Update audio_diffusion_attacks_forhf/src/test_encoder_attack.py

audio_diffusion_attacks_forhf/src/test_encoder_attack.py

@@ -305,13 +305,14 @@ def poison_audio(waveform, sample_rate, encoders, audio_difference_weights=[1],
                     print('step', step, 'loss', loss.item(), 'latent loss', latent_diff.item(), 'audio loss', waveform_diff.item(), 'c_waveform_loss', c_waveform_loss.item(), 'c_stft_loss', c_stft_loss.item(), 'l1_loss', l1_loss.item())
                 latent_diff=latent_diff.detach().item()
             
-            audio_save=torch.reshape((noise+waveform), (2, waveform.shape[2]))[0, :audio_len].detach().cpu().numpy().flatten()
+            #Andy removed: audio_save=torch.reshape((noise+waveform), (2, waveform.shape[2]))[0, :audio_len].detach().cpu().numpy().flatten()
             #Andy removed: wandb.log({f"perturbed cos_dist_{latent_diff}_diff_weight_{diff_weight}_{audio_name}": wandb.Audio(audio_save, sample_rate=sample_rate)}, step=step)
-            
-            
 
-            #Andy removed: torchaudio.save(os.path.join(audio_folder, f"protected_{audio_name}_{audio_len}_mel_{latent_diff}_diff_weight_{waveform_diff}"), torch.reshape((noise+waveform).detach().cpu(), (2, waveform.shape[2])), sample_rate)
-            return torch.reshape((noise+waveform).detach().cpu(), (2, waveform.shape[2]))
+            #Andy moved from inside the loop:
+            music_gen_eval_dict, unprotected_gen, protected_gen=music_gen_eval.eval(waveform, noise+waveform)
+
+            #Andy edited: torchaudio.save(os.path.join(audio_folder, f"protected_{audio_name}_{audio_len}_mel_{latent_diff}_diff_weight_{waveform_diff}"), torch.reshape((noise+waveform).detach().cpu(), (2, waveform.shape[2])), sample_rate)
+            return (torch.reshape((noise+waveform).detach().cpu(), (2, waveform.shape[2]))), music_gen_eval_dict, unprotected_gen, protected_gen
 
 # encoders = [ArchiSound.from_pretrained('autoencoder1d-AT-v1'),
 #             ArchiSound.from_pretrained("dmae1d-ATC64-v2"),