Update model/utils.py

model/utils.py

@@ -557,24 +557,39 @@ def repetition_found(text, length = 2, tolerance = 10):
 
 # load model checkpoint for inference
 
-def load_checkpoint(model, ckpt_path, device, use_ema = True):
-    from ema_pytorch import EMA
+def load_checkpoint(model, ckpt_path, device, dtype=None, use_ema=True):
+    if dtype is None:
+        dtype = (
+            torch.float16 if device == "cuda" and torch.cuda.get_device_properties(device).major >= 6 else torch.float32
+        )
+    model = model.to(dtype)
 
     ckpt_type = ckpt_path.split(".")[-1]
     if ckpt_type == "safetensors":
         from safetensors.torch import load_file
-        checkpoint = load_file(ckpt_path, device=device)
+
+        checkpoint = load_file(ckpt_path)
     else:
-        checkpoint = torch.load(ckpt_path, weights_only=True, map_location=device)
+        checkpoint = torch.load(ckpt_path, weights_only=True)
 
-    if use_ema == True:
-        ema_model = EMA(model, include_online_model = False).to(device)
+    if use_ema:
         if ckpt_type == "safetensors":
-            ema_model.load_state_dict(checkpoint)
-        else:
-            ema_model.load_state_dict(checkpoint['ema_model_state_dict'])
-        ema_model.copy_params_from_ema_to_model()
+            checkpoint = {"ema_model_state_dict": checkpoint}
+        checkpoint["model_state_dict"] = {
+            k.replace("ema_model.", ""): v
+            for k, v in checkpoint["ema_model_state_dict"].items()
+            if k not in ["initted", "step"]
+        }
+
+        # patch for backward compatibility, 305e3ea
+        for key in ["mel_spec.mel_stft.mel_scale.fb", "mel_spec.mel_stft.spectrogram.window"]:
+            if key in checkpoint["model_state_dict"]:
+                del checkpoint["model_state_dict"][key]
+
+        model.load_state_dict(checkpoint["model_state_dict"])
     else:
-        model.load_state_dict(checkpoint['model_state_dict'])
-        
-    return model
+        if ckpt_type == "safetensors":
+            checkpoint = {"model_state_dict": checkpoint}
+        model.load_state_dict(checkpoint["model_state_dict"])
+
+    return model.to(device)