fix(api): get upscale params from request

api/onnx_web/pipeline.py

@@ -16,6 +16,7 @@ from .image import (
 )
 from .upscale import (
     upscale_resrgan,
+    UpscaleParams,
 )
 from .utils import (
     safer_join,
@@ -75,7 +76,8 @@ def run_txt2img_pipeline(
     ctx: ServerContext,
     params: BaseParams,
     size: Size,
-    output: str
+    output: str,
+    upscale: UpscaleParams
 ):
     pipe = load_pipeline(OnnxStableDiffusionPipeline,
                          params.model, params.provider, params.scheduler)
@@ -93,7 +95,9 @@ def run_txt2img_pipeline(
         negative_prompt=params.negative_prompt,
         num_inference_steps=params.steps,
     ).images[0]
-    image = upscale_resrgan(image, ctx.model_path)
+
+    if upscale.faces or upscale.scale > 1:
+        image = upscale_resrgan(ctx, image, upscale)
 
     dest = safer_join(ctx.output_path, output)
     image.save(dest)
@@ -105,8 +109,9 @@ def run_img2img_pipeline(
     ctx: ServerContext,
     params: BaseParams,
     output: str,
+    upscale: UpscaleParams,
     source_image: Image,
-    strength: float
+    strength: float,
 ):
     pipe = load_pipeline(OnnxStableDiffusionImg2ImgPipeline,
                          params.model, params.provider, params.scheduler)
@@ -122,7 +127,9 @@ def run_img2img_pipeline(
         num_inference_steps=params.steps,
         strength=strength,
     ).images[0]
-    image = upscale_resrgan(image, ctx.model_path)
+
+    if upscale.faces or upscale.scale > 1:
+        image = upscale_resrgan(ctx, image, upscale)
 
     dest = safer_join(ctx.output_path, output)
     image.save(dest)
@@ -135,6 +142,7 @@ def run_inpaint_pipeline(
     params: BaseParams,
     size: Size,
     output: str,
+    upscale: UpscaleParams,
     source_image: Image,
     mask_image: Image,
     expand: Border,
@@ -173,6 +181,9 @@ def run_inpaint_pipeline(
         width=size.width,
     ).images[0]
 
+    if upscale.faces or upscale.scale > 1:
+        image = upscale_resrgan(ctx, image, upscale)
+
     dest = safer_join(ctx.output_path, output)
     image.save(dest)
 

api/onnx_web/serve.py

@@ -39,6 +39,9 @@ from .pipeline import (
     run_inpaint_pipeline,
     run_txt2img_pipeline,
 )
+from .upscale import (
+    UpscaleParams,
+)
 from .utils import (
     get_and_clamp_float,
     get_and_clamp_int,
@@ -186,6 +189,13 @@ def border_from_request() -> Border:
     return Border(left, right, top, bottom)
 
 
+def upscale_from_request() -> UpscaleParams:
+    denoise = get_and_clamp_float(request.args, 'denoise', 0.5, 1.0, 0.0)
+    scale = get_and_clamp_int(request.args, 'scale', 1, 4, 1)
+    faces = request.args.get('faces', 'false') == 'true'
+    platform = 'onnx'
+    return UpscaleParams(scale=scale, faces=faces, platform=platform, denoise=denoise)
+
 def check_paths():
     if not path.exists(model_path):
         raise RuntimeError('model path must exist')
@@ -278,6 +288,7 @@ def img2img():
     source_image = Image.open(BytesIO(source_file.read())).convert('RGB')
 
     params, size = pipeline_from_request()
+    upscale = upscale_from_request()
 
     strength = get_and_clamp_float(
         request.args,
@@ -294,7 +305,7 @@ def img2img():
 
     source_image.thumbnail((size.width, size.height))
     executor.submit_stored(output, run_img2img_pipeline,
-                           context, params, output, source_image, strength)
+                           context, params, output, upscale, source_image, strength)
 
     return jsonify({
         'output': output,
@@ -306,6 +317,7 @@ def img2img():
 @app.route('/api/txt2img', methods=['POST'])
 def txt2img():
     params, size = pipeline_from_request()
+    upscale = upscale_from_request()
 
     output = make_output_name(
         'txt2img',
@@ -314,7 +326,7 @@ def txt2img():
     print("txt2img output: %s" % (output))
 
     executor.submit_stored(
-        output, run_txt2img_pipeline, context, params, size, output)
+        output, run_txt2img_pipeline, context, params, size, output, upscale)
 
     return jsonify({
         'output': output,
@@ -333,6 +345,7 @@ def inpaint():
 
     params, size = pipeline_from_request()
     expand = border_from_request()
+    upscale = upscale_from_request()
 
     mask_filter = get_from_map(request.args, 'filter', mask_filters, 'none')
     noise_source = get_from_map(
@@ -362,6 +375,7 @@ def inpaint():
         params,
         size,
         output,
+        upscale,
         source_image,
         mask_image,
         expand,

api/onnx_web/upscale.py

@@ -129,6 +129,8 @@ def make_resrgan(ctx: ServerContext, params: UpscaleParams, tile=0):
 
 
 def upscale_resrgan(ctx: ServerContext, source_image: Image, params: UpscaleParams) -> Image:
+    print('upscaling image with Real ESRGAN', params)
+
     image = np.array(source_image)
     upsampler = make_resrgan(ctx.model_path)
 
@@ -142,6 +144,8 @@ def upscale_resrgan(ctx: ServerContext, source_image: Image, params: UpscalePara
 
 
 def upscale_gfpgan(ctx: ServerContext, image, upsampler=None) -> Image:
+    print('correcting faces with GFPGAN')
+
     if upsampler is None:
         upsampler = make_resrgan(ctx.model_path, 512)
 

onnx-web.code-workspace

@@ -17,8 +17,10 @@
 			"CUDA",
 			"ddim",
 			"ddpm",
+			"denoise",
 			"directml",
 			"ftfy",
+			"gfpgan",
 			"Heun",
 			"huggingface",
 			"Inpaint",
@@ -35,6 +37,7 @@
 			"pndm",
 			"pretrained",
 			"protobuf",
+			"resrgan",
 			"runwayml",
 			"scandir",
 			"scipy",
@@ -42,6 +45,7 @@
 			"spacy",
 			"spinalcase",
 			"stringcase",
+			"upsampler",
 			"venv",
 			"virtualenv",
 			"zustand"