Use model.safetensors with Whisper (#1399)
Some checks are pending
Test / check_lint (push) Waiting to run
Some checks are pending
Test / check_lint (push) Waiting to run
This commit is contained in:
parent
7ddca42f4d
commit
e52c128d11
|
|
@ -382,7 +382,7 @@ if __name__ == "__main__":
|
|||
|
||||
# Save weights
|
||||
print("[INFO] Saving")
|
||||
mx.save_safetensors(str(mlx_path / "weights.safetensors"), weights)
|
||||
mx.save_safetensors(str(mlx_path / "model.safetensors"), weights)
|
||||
|
||||
# Save config.json with model_type
|
||||
with open(str(mlx_path / "config.json"), "w") as f:
|
||||
|
|
|
|||
|
|
@ -26,7 +26,10 @@ def load_model(
|
|||
|
||||
model_args = whisper.ModelDimensions(**config)
|
||||
|
||||
wf = model_path / "weights.safetensors"
|
||||
# Prefer model.safetensors, fall back to weights.safetensors, then weights.npz
|
||||
wf = model_path / "model.safetensors"
|
||||
if not wf.exists():
|
||||
wf = model_path / "weights.safetensors"
|
||||
if not wf.exists():
|
||||
wf = model_path / "weights.npz"
|
||||
weights = mx.load(str(wf))
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user