Update infer.py

infer.py

@@ -50,9 +50,10 @@ def load_model(model_name, audio_tokenizer_path):
         use_flash_attention_2=True,
         use_cache=True,
     )
-    model = model.cuda()
+    model = model.to("cuda")
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    tokenizer_voila = VoilaTokenizer(model_path=audio_tokenizer_path, device="cuda")
+    tokenizer_voila = VoilaTokenizer(model_path=audio_tokenizer_path, device="cpu")
+    tokenizer_voila.to("cuda")
     return model, tokenizer, tokenizer_voila, model_type
 
 def is_audio_output_task(task_type):
@@ -90,11 +91,11 @@ def eval_model(model, tokenizer, tokenizer_voila, model_type, task_type, history
             yield all_tokens[:,i]
 
     if model_type == "autonomous":
-        input_generator = get_input_generator(torch.as_tensor(streaming_user_input_audio_tokens).cuda())
-        input_ids = [torch.as_tensor([input]).transpose(1,2).cuda() for input in input_ids]            # transpose to [bs, seq, num_codebooks]
+        input_generator = get_input_generator(torch.as_tensor(streaming_user_input_audio_tokens).to('cuda'))
+        input_ids = [torch.as_tensor([input]).transpose(1,2).to('cuda') for input in input_ids]            # transpose to [bs, seq, num_codebooks]
         input_ids = torch.cat(input_ids, dim=2)            # concat to [bs, seq, num_codebooks*2]
     else:
-        input_ids = torch.as_tensor([input_ids]).transpose(1,2).cuda()      # transpose to [bs, seq, num_codebooks]
+        input_ids = torch.as_tensor([input_ids]).transpose(1,2).to('cuda')      # transpose to [bs, seq, num_codebooks]
     gen_params = {
         "input_ids": input_ids,
         "ref_embs": ref_embs,
@@ -110,8 +111,8 @@ def eval_model(model, tokenizer, tokenizer_voila, model_type, task_type, history
         "audio_top_k": 50,
     }
     if model_type == "audio":
-        audio_datas = torch.tensor([audio_datas], dtype=torch.bfloat16).cuda()
-        audio_data_masks = torch.tensor([audio_data_masks]).cuda()
+        audio_datas = torch.tensor([audio_datas], dtype=torch.bfloat16).to('cuda')
+        audio_data_masks = torch.tensor([audio_data_masks]).to('cuda')
         gen_params["audio_datas"] = audio_datas
         gen_params["audio_data_masks"] = audio_data_masks
     elif model_type == "autonomous":
@@ -141,7 +142,7 @@ def eval_model(model, tokenizer, tokenizer_voila, model_type, task_type, history
             'text': tokenizer.decode(text_outputs),
         }
         if is_audio_output_task(task_type):
-            audio_values = tokenizer_voila.decode(torch.tensor(audio_outputs).cuda())
+            audio_values = tokenizer_voila.decode(torch.tensor(audio_outputs).to('cuda'))
             out['audio'] = (audio_values.detach().cpu().numpy(), 16000)
         return out
 
@@ -185,10 +186,12 @@ if __name__ == "__main__":
     # step2: encode ref
     ref_embs, ref_embs_mask = None, None
     if is_audio_output_task(args.task_type):
-        spkr_model = SpeakerEmbedding(device="cuda")
+        spkr_model = SpeakerEmbedding(device="cpu")
+        spkr_model.model.to("cuda")
+        spkr_model.device = "cuda"
         wav, sr = torchaudio.load(args.ref_audio)
         ref_embs = spkr_model(wav, sr)
-        ref_embs_mask = torch.tensor([1]).cuda()
+        ref_embs_mask = torch.tensor([1]).to('cuda')
 
     out = eval_model(model, tokenizer, tokenizer_voila, model_type, args.task_type, history, ref_embs, ref_embs_mask)
     print(f"Output str: {out['text']}")