Update llama/m2ugen.py

llama/m2ugen.py

@@ -43,10 +43,10 @@ class M2UGen(nn.Module):
         # https://huggingface.co/m-a-p/MERT-v1-330M
         # And set the mert_path argument to directory with the model files
         print(f'Initialize MERT...')
-        self.mert_model = AutoModel.from_pretrained(self.args.mert_path, trust_remote_code=True)  # .to(self.device)
+        self.mert_model = AutoModel.from_pretrained(self.args.mert_path, trust_remote_code=True).to("cuda:0")
         self.mert_processor = Wav2Vec2FeatureExtractor.from_pretrained(self.args.mert_path, trust_remote_code=True)
-        self.mu_mert_agg = nn.Conv1d(in_channels=25, out_channels=1, kernel_size=1)
-        self.mu_mert_proj = nn.Linear(1024, 4096)
+        self.mu_mert_agg = nn.Conv1d(in_channels=25, out_channels=1, kernel_size=1).to("cuda:0")
+        self.mu_mert_proj = nn.Linear(1024, 4096).to("cuda:0")
 
         if legacy_bridge:
             bridge_norm_layer = nn.LayerNorm
@@ -57,20 +57,20 @@ class M2UGen(nn.Module):
 
         self.feature_scaler = 1
 
-        self.mu_mert_norm_1 = bridge_norm_layer(4096)
-        self.mu_mert_f1_1 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias)
-        self.mu_mert_f2_1 = nn.Linear(4096 * self.feature_scaler, 4096, bias=bridge_bias)
-        self.mu_mert_f3_1 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias)
+        self.mu_mert_norm_1 = bridge_norm_layer(4096).to("cuda:0")
+        self.mu_mert_f1_1 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias).to("cuda:0")
+        self.mu_mert_f2_1 = nn.Linear(4096 * self.feature_scaler, 4096, bias=bridge_bias).to("cuda:0")
+        self.mu_mert_f3_1 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias).to("cuda:0")
 
-        self.mu_mert_norm_2 = bridge_norm_layer(4096)
-        self.mu_mert_f1_2 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias)
-        self.mu_mert_f2_2 = nn.Linear(4096 * self.feature_scaler, 4096, bias=bridge_bias)
-        self.mu_mert_f3_2 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias)
+        self.mu_mert_norm_2 = bridge_norm_layer(4096).to("cuda:0")
+        self.mu_mert_f1_2 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias).to("cuda:0")
+        self.mu_mert_f2_2 = nn.Linear(4096 * self.feature_scaler, 4096, bias=bridge_bias).to("cuda:0")
+        self.mu_mert_f3_2 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias).to("cuda:0")
 
-        self.mu_mert_norm_3 = bridge_norm_layer(4096)
-        self.mu_mert_f1_3 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias)
-        self.mu_mert_f2_3 = nn.Linear(4096 * self.feature_scaler, 4096, bias=bridge_bias)
-        self.mu_mert_f3_3 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias)
+        self.mu_mert_norm_3 = bridge_norm_layer(4096).to("cuda:0")
+        self.mu_mert_f1_3 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias).to("cuda:0")
+        self.mu_mert_f2_3 = nn.Linear(4096 * self.feature_scaler, 4096, bias=bridge_bias).to("cuda:0")
+        self.mu_mert_f3_3 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias).to("cuda:0")
         print(f'MERT initialized...')
 
         # 2. ViT Encoder
@@ -78,26 +78,26 @@ class M2UGen(nn.Module):
         # https://huggingface.co/google/vit-base-patch16-224-in21k
         # And set the vit_path argument to directory with the model files
         print(f'Initialize ViT...')
-        self.vit_model = ViTModel.from_pretrained(self.args.vit_path)  # .to(self.device)
+        self.vit_model = ViTModel.from_pretrained(self.args.vit_path).to("cuda:0")  # .to(self.device)
         self.vit_model.eval()
         self.vit_processor = ViTImageProcessor.from_pretrained(self.args.vit_path, do_rescale=False)
-        self.iu_vit_agg = nn.Conv1d(in_channels=197, out_channels=1, kernel_size=1)
-        self.iu_vit_proj = nn.Linear(768, 4096)
-
-        self.iu_vit_norm_1 = bridge_norm_layer(4096)
-        self.iu_vit_f1_1 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias)
-        self.iu_vit_f2_1 = nn.Linear(4096 * self.feature_scaler, 4096, bias=bridge_bias)
-        self.iu_vit_f3_1 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias)
-
-        self.iu_vit_norm_2 = bridge_norm_layer(4096)
-        self.iu_vit_f1_2 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias)
-        self.iu_vit_f2_2 = nn.Linear(4096 * self.feature_scaler, 4096, bias=bridge_bias)
-        self.iu_vit_f3_2 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias)
-
-        self.iu_vit_norm_3 = bridge_norm_layer(4096)
-        self.iu_vit_f1_3 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias)
-        self.iu_vit_f2_3 = nn.Linear(4096 * self.feature_scaler, 4096, bias=bridge_bias)
-        self.iu_vit_f3_3 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias)
+        self.iu_vit_agg = nn.Conv1d(in_channels=197, out_channels=1, kernel_size=1).to("cuda:0")
+        self.iu_vit_proj = nn.Linear(768, 4096).to("cuda:0")
+
+        self.iu_vit_norm_1 = bridge_norm_layer(4096).to("cuda:0")
+        self.iu_vit_f1_1 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias).to("cuda:0")
+        self.iu_vit_f2_1 = nn.Linear(4096 * self.feature_scaler, 4096, bias=bridge_bias).to("cuda:0")
+        self.iu_vit_f3_1 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias).to("cuda:0")
+
+        self.iu_vit_norm_2 = bridge_norm_layer(4096).to("cuda:0")
+        self.iu_vit_f1_2 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias).to("cuda:0")
+        self.iu_vit_f2_2 = nn.Linear(4096 * self.feature_scaler, 4096, bias=bridge_bias).to("cuda:0")
+        self.iu_vit_f3_2 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias).to("cuda:0")
+
+        self.iu_vit_norm_3 = bridge_norm_layer(4096).to("cuda:0")
+        self.iu_vit_f1_3 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias).to("cuda:0")
+        self.iu_vit_f2_3 = nn.Linear(4096 * self.feature_scaler, 4096, bias=bridge_bias).to("cuda:0")
+        self.iu_vit_f3_3 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias).to("cuda:0")
         print(f'ViT initialized...')
 
         # 3. ViViT Encoder
@@ -105,26 +105,26 @@ class M2UGen(nn.Module):
         # https://huggingface.co/google/vivit-b-16x2-kinetics400
         # And set the vivit_path argument to directory with the model files
         print(f'Initialize ViViT...')
-        self.vivit_model = VivitModel.from_pretrained(self.args.vivit_path)  # .to(self.device)
+        self.vivit_model = VivitModel.from_pretrained(self.args.vivit_path).to("cuda:0")  # .to(self.device)
         self.vivit_model.eval()
         self.vivit_processor = VivitImageProcessor.from_pretrained(self.args.vivit_path)
-        self.iu_vivit_agg = nn.Conv1d(in_channels=3137, out_channels=1, kernel_size=1)
-        self.iu_vivit_proj = nn.Linear(768, 4096)
-
-        self.iu_vivit_norm_1 = bridge_norm_layer(4096)
-        self.iu_vivit_f1_1 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias)
-        self.iu_vivit_f2_1 = nn.Linear(4096 * self.feature_scaler, 4096, bias=bridge_bias)
-        self.iu_vivit_f3_1 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias)
-
-        self.iu_vivit_norm_2 = bridge_norm_layer(4096)
-        self.iu_vivit_f1_2 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias)
-        self.iu_vivit_f2_2 = nn.Linear(4096 * self.feature_scaler, 4096, bias=bridge_bias)
-        self.iu_vivit_f3_2 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias)
-
-        self.iu_vivit_norm_3 = bridge_norm_layer(4096)
-        self.iu_vivit_f1_3 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias)
-        self.iu_vivit_f2_3 = nn.Linear(4096 * self.feature_scaler, 4096, bias=bridge_bias)
-        self.iu_vivit_f3_3 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias)
+        self.iu_vivit_agg = nn.Conv1d(in_channels=3137, out_channels=1, kernel_size=1).to("cuda:0")
+        self.iu_vivit_proj = nn.Linear(768, 4096).to("cuda:0")
+
+        self.iu_vivit_norm_1 = bridge_norm_layer(4096).to("cuda:0")
+        self.iu_vivit_f1_1 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias).to("cuda:0")
+        self.iu_vivit_f2_1 = nn.Linear(4096 * self.feature_scaler, 4096, bias=bridge_bias).to("cuda:0")
+        self.iu_vivit_f3_1 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias).to("cuda:0")
+
+        self.iu_vivit_norm_2 = bridge_norm_layer(4096).to("cuda:0")
+        self.iu_vivit_f1_2 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias).to("cuda:0")
+        self.iu_vivit_f2_2 = nn.Linear(4096 * self.feature_scaler, 4096, bias=bridge_bias).to("cuda:0")
+        self.iu_vivit_f3_2 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias).to("cuda:0")
+
+        self.iu_vivit_norm_3 = bridge_norm_layer(4096).to("cuda:0")
+        self.iu_vivit_f1_3 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias).to("cuda:0")
+        self.iu_vivit_f2_3 = nn.Linear(4096 * self.feature_scaler, 4096, bias=bridge_bias).to("cuda:0")
+        self.iu_vivit_f3_3 = nn.Linear(4096, 4096 * self.feature_scaler, bias=bridge_bias).to("cuda:0")
         print(f'ViViT initialized...')
 
         # 4. llama
@@ -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)
+        self.llama = Transformer(self.model_args).to("cuda:1")
         torch.set_default_tensor_type(torch.FloatTensor)
 
         if load_llama:
@@ -207,7 +207,7 @@ class M2UGen(nn.Module):
         # 5. projector
         self.output_projector = ProjectionLayer(4096, self.model_args.output_dim_tokens,
                                                 num_input_tokens=self.model_args.num_gen_audio_tokens,
-                                                num_output_tokens=self.model_args.num_output_tokens)
+                                                num_output_tokens=self.model_args.num_output_tokens).to("cuda:1")
 
         # 6. Generator
         if self.args.music_decoder.lower() == "audioldm2":
@@ -217,7 +217,7 @@ class M2UGen(nn.Module):
             print(f'Initialize AudioLDM2...')
             dtype = torch.float16 if torch.cuda.is_available() else torch.float32
             self.generation_model = AudioLDM2Pipeline.from_pretrained(self.args.music_decoder_path, torch_dtype=dtype)
-            self.generation_model.to("cuda")
+            self.generation_model.to("cuda:1")
             print(f'AudioLDM2 initialized...')
         else:
             # The model files for MusicGen can be downloaded here in case of network issues:
@@ -225,7 +225,7 @@ class M2UGen(nn.Module):
             # And set the music_decoder_path argument to directory with the model files
             print(f'Initialize MusicGen...')
             self.generation_processor = AutoProcessor.from_pretrained(self.args.music_decoder_path)
-            self.generation_model = MusicgenForConditionalGeneration.from_pretrained(self.args.music_decoder_path)
+            self.generation_model = MusicgenForConditionalGeneration.from_pretrained(self.args.music_decoder_path).to("cuda:1")
             self.generation_model.eval()
             print(f'MusicGen initialized...')
         self.music_decoder = self.args.music_decoder.lower()
@@ -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)
+        self.prefix_query = nn.Embedding(self.query_layer * self.query_len, self.model_args.dim).to("cuda:1")
 
         # 5. knn
         self.knn = knn
@@ -672,10 +672,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).cuda().long()
+        tokens = torch.full((bsz, total_len), 0).to("cuda:1").long()
 
         for k, t in enumerate(prompts):
-            tokens[k, : len(t)] = torch.tensor(t).cuda().long()
+            tokens[k, : len(t)] = torch.tensor(t).to("cuda:1").long()
         input_text_mask = tokens != 0
         start_pos = min_prompt_size
         prev_pos = 0