diff --git a/api/onnx_web/pipeline.py b/api/onnx_web/pipeline.py index 7c3a9781..be8eb7dc 100644 --- a/api/onnx_web/pipeline.py +++ b/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) diff --git a/api/onnx_web/serve.py b/api/onnx_web/serve.py index a46bad4a..64726bb5 100644 --- a/api/onnx_web/serve.py +++ b/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, diff --git a/api/onnx_web/upscale.py b/api/onnx_web/upscale.py index 9551319d..23fb5d77 100644 --- a/api/onnx_web/upscale.py +++ b/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) diff --git a/onnx-web.code-workspace b/onnx-web.code-workspace index 05d416b2..b4296ecf 100644 --- a/onnx-web.code-workspace +++ b/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"