Update llama/m2ugen.py

llama/m2ugen.py

@@ -332,7 +332,7 @@ class M2UGen(nn.Module):
                 sub_x = all_layer_hidden_states.mean(-2).unsqueeze(0)
                 aggoutputs += sub_x
             aggoutputs /= len(all_inputs)
-            sub_x = self.mu_mert_agg(aggoutputs.to(self.device)).squeeze()
+            sub_x = self.mu_mert_agg(aggoutputs.to("cuda:0")).squeeze()
             del aggoutputs
             xs.append(sub_x)
         x = torch.stack(xs, dim=0)
@@ -345,7 +345,7 @@ class M2UGen(nn.Module):
             with torch.no_grad():
                 outputs = self.vit_model(**inputs)
             last_hidden_states = outputs.last_hidden_state
-            sub_x = self.iu_vit_agg(last_hidden_states.to(self.device)).squeeze()
+            sub_x = self.iu_vit_agg(last_hidden_states.to("cuda:0")).squeeze()
             xs.append(sub_x)
         return torch.stack(xs, dim=0)
 
@@ -356,7 +356,7 @@ class M2UGen(nn.Module):
             with torch.no_grad():
                 outputs = self.vivit_model(**inputs)
             last_hidden_states = outputs.last_hidden_state
-            sub_x = self.iu_vivit_agg(last_hidden_states.to(self.device)).squeeze()
+            sub_x = self.iu_vivit_agg(last_hidden_states.to("cuda:0")).squeeze()
             xs.append(sub_x)
         return torch.stack(xs, dim=0)
 
@@ -489,20 +489,21 @@ class M2UGen(nn.Module):
     @torch.inference_mode()
     def forward_inference(self, tokens, start_pos: int, audio_feats=None, image_feats=None, video_feats=None):
         _bsz, seqlen = tokens.shape
-        h = self.llama.tok_embeddings(tokens)
-        freqs_cis = self.llama.freqs_cis.to(h.device)
+        h = self.llama.tok_embeddings(tokens).to("cuda:1")
+        freqs_cis = self.llama.freqs_cis.to("cuda:1")
         freqs_cis = freqs_cis[start_pos:start_pos + seqlen]
 
-        feats = torch.zeros((1, 1, 4096)).to(self.device)
+        feats = torch.zeros((1, 1, 4096)).to("cuda:0")
         if audio_feats is not None:
             feats += audio_feats
         if video_feats is not None:
             feats += video_feats
         if image_feats is not None:
             feats += image_feats
-
+        feats = feats.to("cuda:1")
+        
         mask = None
-        mask = torch.full((1, 1, seqlen, seqlen), float("-inf"), device=h.device)
+        mask = torch.full((1, 1, seqlen, seqlen), float("-inf"), device="cuda:1")
         mask = torch.triu(mask, diagonal=start_pos + 1).type_as(h)
 
         music_output_embedding = []
@@ -603,11 +604,10 @@ class M2UGen(nn.Module):
     @torch.inference_mode()
     def generate_music(self, embeddings, audio_length_in_s, music_caption):
         gen_prefix = ''.join([f'[AUD{i}]' for i in range(len(self.audio_tokens))])
-        gen_prefx_ids = self.tokenizer(gen_prefix, add_special_tokens=False, return_tensors="pt").input_ids.to(
-            self.device)
+        gen_prefx_ids = self.tokenizer(gen_prefix, add_special_tokens=False, return_tensors="pt").input_ids.to("cuda:1")
         gen_prefix_embs = self.llama.tok_embeddings(gen_prefx_ids)
         if self.music_decoder == "audioldm2":
-            gen_emb = self.output_projector(embeddings.float().to("cuda"), gen_prefix_embs).squeeze(dim=0) / 10
+            gen_emb = self.output_projector(embeddings.float().to("cuda:1"), gen_prefix_embs).squeeze(dim=0) / 10
             prompt_embeds, generated_prompt_embeds = gen_emb[:, :128 * 1024], gen_emb[:, 128 * 1024:]
             prompt_embeds = prompt_embeds.reshape(prompt_embeds.shape[0], 128, 1024)
             generated_prompt_embeds = generated_prompt_embeds.reshape(generated_prompt_embeds.shape[0], 8, 768)
@@ -623,8 +623,7 @@ class M2UGen(nn.Module):
             print("Generating Music...")
             gen_emb = 0.1 * self.output_projector(embeddings.float().to("cuda"), gen_prefix_embs) / 10
             gen_inputs = self.generation_processor(text=music_caption, padding='max_length',
-                                                   max_length=128, truncation=True, return_tensors="pt").to(
-                self.device)
+                                                   max_length=128, truncation=True, return_tensors="pt").to("cuda:1")
             #gen_emb = self.generation_model.generate(**gen_inputs, guidance_scale=3.5, encoder_only=True)
             audio_outputs = self.generation_model.generate(**gen_inputs, guidance_scale=3.5,
                                                            max_new_tokens=int(256 / 5 * audio_length_in_s))