Update infer/utils_infer.py

infer/utils_infer.py

@@ -137,9 +137,13 @@ asr_pipe = None
 
 def initialize_asr_pipeline(device: str = device, dtype=None):
     if dtype is None:
-        dtype = (
-            torch.float16 if "cuda" in device and torch.cuda.get_device_properties(device).major >= 6 else torch.float32
-        )
+        if "cuda" in device and torch.cuda.get_device_properties(device).major >= 6:
+            dtype = torch.float16
+        elif "cpu" in device:
+            dtype = torch.bfloat16
+        else:
+            dtype = torch.float32
+
     global asr_pipe
     asr_pipe = pipeline(
         "automatic-speech-recognition",
@@ -149,6 +153,7 @@ def initialize_asr_pipeline(device: str = device, dtype=None):
     )
 
 
+
 # transcribe
 
 
@@ -170,10 +175,16 @@ def transcribe(ref_audio, language=None):
 
 def load_checkpoint(model, ckpt_path, device: str, dtype=None, use_ema=True):
     if dtype is None:
-        dtype = (
-            torch.float16 if "cuda" in device and torch.cuda.get_device_properties(device).major >= 6 else torch.float32
-        )
-    model = model.to(dtype)
+        if "cuda" in device and torch.cuda.get_device_properties(device).major >= 6:
+            dtype = torch.float16
+        elif "cpu" in device:
+            dtype = torch.bfloat16
+        else:
+            dtype = torch.float32
+
+    # Move the model to the desired device and dtype
+    model = model.to(device=device, dtype=dtype)
+    #model = model.to(dtype)
 
     ckpt_type = ckpt_path.split(".")[-1]
     if ckpt_type == "safetensors":