fix(api): trim model names relative to model path

ssube · Jan 17, 2023 · 4472a6f · 4472a6f
1 parent 23a9d5a
commit 4472a6f
Showing 1 changed file with 14 additions and 5 deletions.
diff --git a/api/onnx_web/serve.py b/api/onnx_web/serve.py
@@ -203,17 +203,26 @@ def check_paths(context: ServerContext):
         makedirs(context.output_path)
 
 
+def get_model_name(model: str) -> str:
+    base = path.basename(model)
+    (file, _ext) = path.splitext(base)
+    return file
+
+
 def load_models(context: ServerContext):
     global diffusion_models
     global correction_models
     global upscaling_models
 
-    diffusion_models = glob(path.join(context.model_path, 'diffusion-*'))
-    diffusion_models.extend(glob(path.join(context.model_path, 'stable-diffusion-*')))
-
-    correction_models = glob(path.join(context.model_path, 'correction-*'))
-    upscaling_models = glob(path.join(context.model_path, 'upscaling-*'))
+    diffusion_models = [get_model_name(f) for f in glob(
+        path.join(context.model_path, 'diffusion-*'))]
+    diffusion_models.extend([
+        get_model_name(f) for f in glob(path.join(context.model_path, 'stable-diffusion-*'))])
 
+    correction_models = [
+        get_model_name(f) for f in glob(path.join(context.model_path, 'correction-*'))]
+    upscaling_models = [
+        get_model_name(f) for f in glob(path.join(context.model_path, 'upscaling-*'))]
 
 
 def load_params(context: ServerContext):