Update audio_diffusion_attacks_forhf/src/music_gen.py

audio_diffusion_attacks_forhf/src/music_gen.py

@@ -12,7 +12,7 @@ class MusicGenEval:
         model_name="facebook/musicgen-stereo-small"
         self.processor = AutoProcessor.from_pretrained(model_name)
         self.model = MusicgenForConditionalGeneration.from_pretrained(model_name)
-        self.model=self.model.to(device='cuda')
+        #Andy commented: self.model=self.model.to(device='cuda')
         self.input_sample_rate=input_sample_rate
         self.audio_steps=audio_steps
         self.mel_loss = losses.MelSpectrogramLoss(n_mels=[5, 10, 20, 40, 80, 160, 320],
@@ -27,8 +27,10 @@ class MusicGenEval:
         protected_audio=protected_audio[:, :, :self.audio_steps]
         input_len=original_audio.shape[-1]
         
-        unprotected_gen=self.generate_audio(original_audio)[0].to(device='cuda')
-        protected_gen=self.generate_audio(protected_audio)[0].to(device='cuda')
+        #Andy edited: unprotected_gen=self.generate_audio(original_audio)[0].to(device='cuda')
+        unprotected_gen=self.generate_audio(original_audio)[0]
+        #Andy edited: protected_gen=self.generate_audio(protected_audio)[0].to(device='cuda')
+        protected_gen=self.generate_audio(protected_audio)[0]
         
         eval_dict={}
         # Difference between original and unprotected gen
@@ -48,7 +50,8 @@ class MusicGenEval:
     def generate_audio(self, audio):
         torch.manual_seed(0)
 
-        transform = torchaudio.transforms.Resample(self.input_sample_rate, 32000).to(device='cuda')
+        #Andy edited: transform = torchaudio.transforms.Resample(self.input_sample_rate, 32000).to(device='cuda')
+        transform = torchaudio.transforms.Resample(self.input_sample_rate, 32000)
         waveform=transform(audio[0]).detach().cpu()
         # waveform.clamp_(0,1)
         a=torch.min(waveform)
@@ -64,10 +67,12 @@ class MusicGenEval:
             return_tensors="pt",
         )
         for d in inputs.data:
-            inputs.data[d]=inputs.data[d].to(device='cuda')
+            #Andy edited: inputs.data[d]=inputs.data[d].to(device='cuda')
+            inputs.data[d]=inputs.data[d]
         audio_values = self.model.generate(**inputs, do_sample=True, guidance_scale=3, max_new_tokens=1024)
         
-        transform = torchaudio.transforms.Resample(32000, self.input_sample_rate).to(device='cuda')
+        #Andy edited: transform = torchaudio.transforms.Resample(32000, self.input_sample_rate).to(device='cuda')
+        transform = torchaudio.transforms.Resample(32000, self.input_sample_rate)
         audio_values=transform(audio_values)
         return audio_values