From 4a89bee517411db10b86c20046747b40c3162119 Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Sat, 2 Mar 2024 13:11:45 -0600
Subject: [PATCH] fix special param parsing within img2img/inpaint pipelines

---
 api/onnx_web/server/api.py           | 20 ++++++++++----------
 api/onnx_web/server/params.py        |  8 ++++----
 api/tests/test_diffusers/test_run.py |  2 +-
 3 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/api/onnx_web/server/api.py b/api/onnx_web/server/api.py
index b30b3ac8..ac8fe955 100644
--- a/api/onnx_web/server/api.py
+++ b/api/onnx_web/server/api.py
@@ -263,12 +263,13 @@ def img2img(server: ServerContext, pool: DevicePoolExecutor):
     source = Image.open(BytesIO(source_file.read())).convert("RGB")
 
     data = get_request_data()
+    data_params = data.get("params", data)
     source_filter = get_from_list(
-        data["params"], "sourceFilter", list(get_source_filters().keys())
+        data_params, "sourceFilter", list(get_source_filters().keys())
     )
 
     strength = get_and_clamp_float(
-        data["params"],
+        data_params,
         "strength",
         get_config_value("strength"),
         get_config_value("strength", "max"),
@@ -336,11 +337,12 @@ def inpaint(server: ServerContext, pool: DevicePoolExecutor):
     mask.convert(mode="L")
 
     data = get_request_data()
+    data_params = data.get("params", data)
     full_res_inpaint = get_boolean(
-        data["params"], "fullresInpaint", get_config_value("fullresInpaint")
+        data_params, "fullresInpaint", get_config_value("fullresInpaint")
     )
     full_res_inpaint_padding = get_and_clamp_float(
-        data["params"],
+        data_params,
         "fullresInpaintPadding",
         get_config_value("fullresInpaintPadding"),
         get_config_value("fullresInpaintPadding", "max"),
@@ -350,13 +352,11 @@ def inpaint(server: ServerContext, pool: DevicePoolExecutor):
     params = get_request_params(server, JobType.INPAINT.value)
     replace_wildcards(params.image, get_wildcard_data())
 
-    fill_color = get_not_empty(data["params"], "fillColor", "white")
-    mask_filter = get_from_map(data["params"], "filter", get_mask_filters(), "none")
-    noise_source = get_from_map(
-        data["params"], "noise", get_noise_sources(), "histogram"
-    )
+    fill_color = get_not_empty(data_params, "fillColor", "white")
+    mask_filter = get_from_map(data_params, "filter", get_mask_filters(), "none")
+    noise_source = get_from_map(data_params, "noise", get_noise_sources(), "histogram")
     tile_order = get_from_list(
-        data["params"],
+        data_params,
         "tileOrder",
         [TileOrder.grid, TileOrder.kernel, TileOrder.spiral],
     )
diff --git a/api/onnx_web/server/params.py b/api/onnx_web/server/params.py
index 4464237e..a80fc90b 100644
--- a/api/onnx_web/server/params.py
+++ b/api/onnx_web/server/params.py
@@ -476,10 +476,10 @@ def get_request_params(
 
     device, params, size = pipeline_from_json(server, data, default_pipeline)
 
-    border = build_border(data.get("border"))
-    upscale = build_upscale(data.get("upscale"))
-    highres = build_highres(data.get("highres"))
-    experimental = build_experimental(data.get("experimental"))
+    border = build_border(data.get("border", {}))
+    upscale = build_upscale(data.get("upscale", {}))
+    highres = build_highres(data.get("highres", {}))
+    experimental = build_experimental(data.get("experimental", {}))
 
     return RequestParams(
         device,
diff --git a/api/tests/test_diffusers/test_run.py b/api/tests/test_diffusers/test_run.py
index ffea3dc1..ae3ffdb2 100644
--- a/api/tests/test_diffusers/test_run.py
+++ b/api/tests/test_diffusers/test_run.py
@@ -90,7 +90,7 @@ class TestTxt2ImgPipeline(unittest.TestCase):
 
         self.assertTrue(path.exists("../outputs/test-txt2img-basic_0.png"))
 
-        with Image.open("../outputs/test-txt2img-basic.png") as output:
+        with Image.open("../outputs/test-txt2img-basic_0.png") as output:
             self.assertEqual(output.size, (256, 256))
             # TODO: test contents of image