Update util.py

util.py

@@ -73,8 +73,8 @@ class NemoAudioPlayer:
         ).eval()
         
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        print(f"Moving NeMo codec to device: {self.device}")
-        self.nemo_codec_model.to(self.device)
+        # print(f"Moving NeMo codec to device: {self.device}")
+        # self.nemo_codec_model.to(self.device)
         
         self.text_tokenizer_name = text_tokenizer_name
         if self.text_tokenizer_name:
@@ -145,7 +145,12 @@ class NemoAudioPlayer:
         return text
 
     def get_waveform(self, out_ids):
+
         """Convert model output to audio waveform"""
+
+        print(f"Moving NeMo codec to device: {self.device}")
+        self.nemo_codec_model.to(self.device)
+
         out_ids = out_ids.flatten()
 
         # Validate output
@@ -223,7 +228,7 @@ class KaniModel:
         self.model = AutoModelForCausalLM.from_pretrained(
             self.conf.model_name,
             dtype=torch.bfloat16,
-            device_map=self.conf.device_map,
+            # device_map=self.conf.device_map,
             trust_remote_code=True
         )
 
@@ -264,6 +269,7 @@ class KaniModel:
             rp: float,
             max_tok: int) -> torch.tensor:
         """Generate tokens using the model"""
+        self.model.to(self.device)
         input_ids = input_ids.to(self.device)
         attention_mask = attention_mask.to(self.device)