diff --git a/component/tagparser/nameparser.go b/component/tagparser/nameparser.go index f6dd70619..77bbe3f20 100644 --- a/component/tagparser/nameparser.go +++ b/component/tagparser/nameparser.go @@ -36,7 +36,7 @@ func LibraryTag(filePath string) string { } func isPytorch(filename string) bool { - return (strings.HasPrefix(filename, "pytorch_model") && strings.HasSuffix(filename, ".bin")) || strings.HasSuffix(filename, ".pt") + return (strings.HasPrefix(filename, "pytorch_model") && strings.HasSuffix(filename, ".bin")) || strings.HasSuffix(filename, ".pt") || strings.HasSuffix(filename, ".pth") } func isTensorflow(filename string) bool { diff --git a/component/tagparser/nameparser_test.go b/component/tagparser/nameparser_test.go index 2ec734fa4..8a5553616 100644 --- a/component/tagparser/nameparser_test.go +++ b/component/tagparser/nameparser_test.go @@ -16,6 +16,7 @@ func TestLibraryTag(t *testing.T) { {name: "pytorch", args: args{filePath: "pytorch_model.bin"}, want: "pytorch"}, {name: "pytorch", args: args{filePath: "pytorch_model_001.bin"}, want: "pytorch"}, {name: "pytorch", args: args{filePath: "model.pt"}, want: "pytorch"}, + {name: "pytorch", args: args{filePath: "model.pth"}, want: "pytorch"}, {name: "not pytorch", args: args{filePath: "1-pytorch_model_001.bin"}, want: ""}, {name: "not pytorch", args: args{filePath: "pytorch_model-bin"}, want: ""}, diff --git a/configs/inference/audio-fish.json b/configs/inference/audio-fish.json index 38478db4e..06ae1e505 100644 --- a/configs/inference/audio-fish.json +++ b/configs/inference/audio-fish.json @@ -1,8 +1,8 @@ { - "engine_name": "hf-inference-toolkit", + "engine_name": "audio-fish", "enabled": 1, "container_port": 8000, - "model_format": "safetensors", + "model_format": "pytorch", "engine_images": [ { "compute_type": "gpu",