From 2cc604823b28321c111ddd62f87f22a8ab5f43e2 Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Fri, 10 Feb 2023 23:06:30 -0600
Subject: [PATCH] fix(api): launch with included extras if user does not
 provide their own

---
 api/launch.bat                   |  1 +
 api/launch.sh                    |  7 ++++++-
 api/onnx_web/convert/__main__.py | 24 ++++++++++++------------
 api/schemas/extras.yaml          |  3 ++-
 4 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/api/launch.bat b/api/launch.bat
index 9d851e31..dd1cd87c 100644
--- a/api/launch.bat
+++ b/api/launch.bat
@@ -1,4 +1,5 @@
 echo "Downloading and converting models to ONNX format..."
+IF "%ONNX_WEB_EXTRA_MODELS%"=="" (set ONNX_WEB_EXTRA_MODELS=extras.json)
 python -m onnx_web.convert --diffusion --upscaling --correction --extras=extras.json --token=%HF_TOKEN%
 
 echo "Launching API server..."
diff --git a/api/launch.sh b/api/launch.sh
index 3fd5e5a0..84b377ed 100755
--- a/api/launch.sh
+++ b/api/launch.sh
@@ -16,7 +16,12 @@ else
 fi
 
 echo "Downloading and converting models to ONNX format..."
-python3 -m onnx_web.convert --diffusion --upscaling --correction --extras=extras.json --token=${HF_TOKEN:-}
+python3 -m onnx_web.convert \
+  --diffusion \
+  --upscaling \
+  --correction \
+  --extras=${ONNX_WEB_EXTRA_MODELS:-extras.json} \
+  --token=${HF_TOKEN:-}
 
 echo "Launching API server..."
 flask --app=onnx_web.serve run --host=0.0.0.0
diff --git a/api/onnx_web/convert/__main__.py b/api/onnx_web/convert/__main__.py
index 86f13659..38188e97 100644
--- a/api/onnx_web/convert/__main__.py
+++ b/api/onnx_web/convert/__main__.py
@@ -45,19 +45,19 @@ base_models: Models = {
             "stable-diffusion-onnx-v1-5",
             model_source_huggingface + "runwayml/stable-diffusion-v1-5",
         ),
-        # (
-        #     "stable-diffusion-onnx-v1-inpainting",
-        #     model_source_huggingface + "runwayml/stable-diffusion-inpainting",
-        # ),
+        (
+            "stable-diffusion-onnx-v1-inpainting",
+            model_source_huggingface + "runwayml/stable-diffusion-inpainting",
+        ),
         # v2.x
-        # (
-        #     "stable-diffusion-onnx-v2-1",
-        #     model_source_huggingface + "stabilityai/stable-diffusion-2-1",
-        # ),
-        # (
-        #     "stable-diffusion-onnx-v2-inpainting",
-        #     model_source_huggingface + "stabilityai/stable-diffusion-2-inpainting",
-        # ),
+        (
+            "stable-diffusion-onnx-v2-1",
+            model_source_huggingface + "stabilityai/stable-diffusion-2-1",
+        ),
+        (
+            "stable-diffusion-onnx-v2-inpainting",
+            model_source_huggingface + "stabilityai/stable-diffusion-2-inpainting",
+        ),
         # TODO: should have its own converter
         (
             "upscaling-stable-diffusion-x4",
diff --git a/api/schemas/extras.yaml b/api/schemas/extras.yaml
index 56230cc4..b611d0e3 100644
--- a/api/schemas/extras.yaml
+++ b/api/schemas/extras.yaml
@@ -6,8 +6,9 @@ $defs:
     type: array
     items:
       oneOf:
-        - type: string
+        - type: boolean
         - type: number
+        - type: string
 
   base_model:
     type: object