Update llama/m2ugen.py

llama/m2ugen.py

@@ -152,7 +152,7 @@ class M2UGen(nn.Module):
 
         if torch.cuda.is_available():
             torch.set_default_tensor_type(torch.cuda.HalfTensor)
-        self.llama = Transformer(self.model_args).to("cuda:1")
+        self.llama = Transformer(self.model_args).to("cuda:0")
         torch.set_default_tensor_type(torch.FloatTensor)
 
         if load_llama:
@@ -233,7 +233,7 @@ class M2UGen(nn.Module):
         # 4. prefix
         self.query_layer = 20
         self.query_len = 1
-        self.prefix_query = nn.Embedding(self.query_layer * self.query_len, self.model_args.dim).to("cuda:1")
+        self.prefix_query = nn.Embedding(self.query_layer * self.query_len, self.model_args.dim).to("cuda:0")
 
         # 5. knn
         self.knn = knn
@@ -489,8 +489,8 @@ 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).to("cuda:1")
-        freqs_cis = self.llama.freqs_cis.to("cuda:1")
+        h = self.llama.tok_embeddings(tokens).to("cuda:0")
+        freqs_cis = self.llama.freqs_cis.to("cuda:0")
         freqs_cis = freqs_cis[start_pos:start_pos + seqlen]
 
         feats = torch.zeros((1, 1, 4096)).to("cuda:0")
@@ -500,10 +500,9 @@ class M2UGen(nn.Module):
             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="cuda:1")
+        mask = torch.full((1, 1, seqlen, seqlen), float("-inf"), device="cuda:0")
         mask = torch.triu(mask, diagonal=start_pos + 1).type_as(h)
 
         music_output_embedding = []
@@ -669,10 +668,10 @@ class M2UGen(nn.Module):
 
         total_len = min(params.max_seq_len, max_gen_len + max_prompt_size)
 
-        tokens = torch.full((bsz, total_len), 0).to("cuda:1").long()
+        tokens = torch.full((bsz, total_len), 0).to("cuda:0").long()
 
         for k, t in enumerate(prompts):
-            tokens[k, : len(t)] = torch.tensor(t).to("cuda:1").long()
+            tokens[k, : len(t)] = torch.tensor(t).to("cuda:0").long()
         input_text_mask = tokens != 0
         start_pos = min_prompt_size
         prev_pos = 0