fix(api): allow PTH tensor files, add helper to check extension

ssube · Oct 7, 2023 · ebdfa78 · ebdfa78
1 parent 1351b2f
commit ebdfa78
Showing 1 changed file with 10 additions and 3 deletions.
diff --git a/api/onnx_web/convert/utils.py b/api/onnx_web/convert/utils.py
@@ -185,16 +185,23 @@ def tuple_to_upscaling(model: Union[ModelDict, LegacyModel]):
 
 
 MODEL_FORMATS = ["onnx", "pth", "ckpt", "safetensors"]
-RESOLVE_FORMATS = ["safetensors", "ckpt", "pt", "bin"]
+RESOLVE_FORMATS = ["safetensors", "ckpt", "pt", "pth", "bin"]
+
+
+def check_ext(name: str, exts: List[str]) -> Tuple[bool, str]:
+    _name, ext = path.splitext(name)
+    ext = ext.strip(".")
+
+    return (name in exts, ext)
 
 
 def source_format(model: Dict) -> Optional[str]:
     if "format" in model:
         return model["format"]
 
     if "source" in model:
-        _name, ext = path.splitext(model["source"])
-        if ext in MODEL_FORMATS:
+        valid, ext = check_ext(model["source"], MODEL_FORMATS)
+        if valid:
             return ext
 
     return None