fix(api): include SD upscaling in diffusion prefixes

ssube · Dec 9, 2023 · 293a1bb · 293a1bb
1 parent 46d9fc0
commit 293a1bb
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 14 deletions.
diff --git a/api/onnx_web/convert/__main__.py b/api/onnx_web/convert/__main__.py
@@ -26,6 +26,7 @@
     DEFAULT_OPSET,
     ConversionContext,
     download_progress,
+    fix_diffusion_name,
     remove_prefix,
     source_format,
     tuple_to_correction,
@@ -265,20 +266,6 @@ def fetch_model(
         return source, False
 
 
-DIFFUSION_PREFIX = ["diffusion-", "diffusion/", "diffusion\\", "stable-diffusion-"]
-
-
-def fix_diffusion_name(name: str):
-    if not any([name.startswith(prefix) for prefix in DIFFUSION_PREFIX]):
-        logger.warning(
-            "diffusion models must have names starting with diffusion- to be recognized by the server: %s does not match",
-            name,
-        )
-        return f"diffusion-{name}"
-
-    return name
-
-
 def convert_models(conversion: ConversionContext, args, models: Models):
     model_errors = []
 

diff --git a/api/onnx_web/convert/utils.py b/api/onnx_web/convert/utils.py
@@ -345,3 +345,17 @@ def onnx_export(
             all_tensors_to_one_file=True,
             location=ONNX_WEIGHTS,
         )
+
+
+DIFFUSION_PREFIX = ["diffusion-", "diffusion/", "diffusion\\", "stable-diffusion-", "upscaling-"]
+
+
+def fix_diffusion_name(name: str):
+    if not any([name.startswith(prefix) for prefix in DIFFUSION_PREFIX]):
+        logger.warning(
+            "diffusion models must have names starting with diffusion- to be recognized by the server: %s does not match",
+            name,
+        )
+        return f"diffusion-{name}"
+
+    return name