Update infer/utils_infer.py

infer/utils_infer.py

@@ -101,6 +101,7 @@ def load_vocoder(vocoder_name="vocos", is_local=False, local_path="", device=dev
             repo_id = "charactr/vocos-mel-24khz"
             config_path = hf_hub_download(repo_id=repo_id, cache_dir=hf_cache_dir, filename="config.yaml")
             model_path = hf_hub_download(repo_id=repo_id, cache_dir=hf_cache_dir, filename="pytorch_model.bin")
+        
         vocoder = Vocos.from_hparams(config_path)
         state_dict = torch.load(model_path, map_location="cpu", weights_only=True)
         from vocos.feature_extractors import EncodecFeatures
@@ -111,13 +112,18 @@ def load_vocoder(vocoder_name="vocos", is_local=False, local_path="", device=dev
                 for key, value in vocoder.feature_extractor.encodec.state_dict().items()
             }
             state_dict.update(encodec_parameters)
+        
         vocoder.load_state_dict(state_dict)
-        vocoder = vocoder.eval().to(device)
+        
+        # Convert vocoder to bfloat16 if using a compatible device
+        vocoder = vocoder.eval().to(device).to(torch.bfloat16)
+    
     elif vocoder_name == "bigvgan":
         try:
             from third_party.BigVGAN import bigvgan
         except ImportError:
             print("You need to follow the README to init submodule and change the BigVGAN source code.")
+        
         if is_local:
             """download from https://huggingface.co/nvidia/bigvgan_v2_24khz_100band_256x/tree/main"""
             vocoder = bigvgan.BigVGAN.from_pretrained(local_path, use_cuda_kernel=False)
@@ -126,7 +132,8 @@ def load_vocoder(vocoder_name="vocos", is_local=False, local_path="", device=dev
             vocoder = bigvgan.BigVGAN.from_pretrained(local_path, use_cuda_kernel=False)
 
         vocoder.remove_weight_norm()
-        vocoder = vocoder.eval().to(device)
+        vocoder = vocoder.eval().to(device).to(torch.bfloat16)  # Convert to bfloat16
+    
     return vocoder