diff --git a/api/onnx_web/chain/base.py b/api/onnx_web/chain/base.py index 56e08233..6ed14294 100644 --- a/api/onnx_web/chain/base.py +++ b/api/onnx_web/chain/base.py @@ -9,6 +9,9 @@ from ..params import ( ImageParams, StageParams, ) +from ..output import ( + save_image, +) from ..utils import ( is_debug, ServerContext, @@ -82,7 +85,7 @@ class ChainPipeline: **kwargs) if is_debug(): - tile.save(path.join(ctx.output_path, 'last-tile.png')) + save_image(ctx, 'last-tile.png', tile) return tile @@ -97,7 +100,7 @@ class ChainPipeline: name, image.width, image.height) if is_debug(): - image.save(path.join(ctx.output_path, 'last-stage.png')) + save_image(ctx, 'last-stage.png', image) end = monotonic() duration = timedelta(seconds=(end - start)) diff --git a/api/onnx_web/chain/blend_inpaint.py b/api/onnx_web/chain/blend_inpaint.py index 8af082d6..ebca759d 100644 --- a/api/onnx_web/chain/blend_inpaint.py +++ b/api/onnx_web/chain/blend_inpaint.py @@ -21,8 +21,10 @@ from ..params import ( SizeChart, StageParams, ) +from ..output import ( + save_image, +) from ..utils import ( - base_join, is_debug, ServerContext, ) @@ -63,9 +65,9 @@ def blend_inpaint( mask_filter=mask_filter) if is_debug(): - source_image.save(base_join(ctx.output_path, 'last-source.png')) - mask_image.save(base_join(ctx.output_path, 'last-mask.png')) - noise_image.save(base_join(ctx.output_path, 'last-noise.png')) + save_image(ctx, 'last-source.png', source_image) + save_image(ctx, 'last-mask.png', mask_image) + save_image(ctx, 'last-noise.png', noise_image) def outpaint(image: Image.Image, dims: Tuple[int, int, int]): left, top, tile = dims @@ -73,8 +75,8 @@ def blend_inpaint( mask = mask_image.crop((left, top, left + tile, top + tile)) if is_debug(): - image.save(base_join(ctx.output_path, 'tile-source.png')) - mask.save(base_join(ctx.output_path, 'tile-mask.png')) + save_image(ctx, 'tile-source.png', image) + save_image(ctx, 'tile-mask.png', mask) # TODO: must use inpainting model here model = '../models/stable-diffusion-onnx-v1-inpainting' diff --git a/api/onnx_web/chain/persist_disk.py b/api/onnx_web/chain/persist_disk.py index 6c25e1db..fb426850 100644 --- a/api/onnx_web/chain/persist_disk.py +++ b/api/onnx_web/chain/persist_disk.py @@ -5,8 +5,10 @@ from ..params import ( ImageParams, StageParams, ) +from ..output import ( + save_image, +) from ..utils import ( - base_join, ServerContext, ) @@ -22,7 +24,6 @@ def persist_disk( output: str, **kwargs, ) -> Image.Image: - dest = base_join(ctx.output_path, output) - source_image.save(dest) + dest = save_image(ctx, output, source_image) logger.info('saved image to %s', dest) return source_image diff --git a/api/onnx_web/chain/upscale_outpaint.py b/api/onnx_web/chain/upscale_outpaint.py index 8332bddb..af74606d 100644 --- a/api/onnx_web/chain/upscale_outpaint.py +++ b/api/onnx_web/chain/upscale_outpaint.py @@ -22,6 +22,9 @@ from ..params import ( SizeChart, StageParams, ) +from ..output import ( + save_image, +) from ..utils import ( base_join, is_debug, @@ -70,9 +73,9 @@ def upscale_outpaint( full_latents = get_latents_from_seed(params.seed, full_size) if is_debug(): - source_image.save(base_join(ctx.output_path, 'last-source.png')) - mask_image.save(base_join(ctx.output_path, 'last-mask.png')) - noise_image.save(base_join(ctx.output_path, 'last-noise.png')) + save_image(ctx, 'last-source.png', source_image) + save_image(ctx, 'last-mask.png', mask_image) + save_image(ctx, 'last-noise.png', noise_image) def outpaint(image: Image.Image, dims: Tuple[int, int, int]): left, top, tile = dims @@ -80,8 +83,8 @@ def upscale_outpaint( mask = mask_image.crop((left, top, left + tile, top + tile)) if is_debug(): - image.save(base_join(ctx.output_path, 'tile-source.png')) - mask.save(base_join(ctx.output_path, 'tile-mask.png')) + save_image(ctx, 'tile-source.png', image) + save_image(ctx, 'tile-mask.png', mask) # TODO: must use inpainting model here model = '../models/stable-diffusion-onnx-v1-inpainting' diff --git a/api/onnx_web/diffusion/load.py b/api/onnx_web/diffusion/load.py index c8e457e4..d4ef9f94 100644 --- a/api/onnx_web/diffusion/load.py +++ b/api/onnx_web/diffusion/load.py @@ -11,7 +11,6 @@ from ..utils import ( run_gc, ) -import gc import numpy as np logger = getLogger(__name__) diff --git a/api/onnx_web/diffusion/run.py b/api/onnx_web/diffusion/run.py index 80c79810..d2694d8b 100644 --- a/api/onnx_web/diffusion/run.py +++ b/api/onnx_web/diffusion/run.py @@ -15,12 +15,15 @@ from ..params import ( Size, StageParams, ) +from ..output import ( + save_image, + save_params, +) from ..upscale import ( run_upscale_correction, UpscaleParams, ) from ..utils import ( - base_join, run_gc, ServerContext, ) @@ -61,8 +64,8 @@ def run_txt2img_pipeline( image = run_upscale_correction( ctx, StageParams(), params, image, upscale=upscale) - dest = base_join(ctx.output_path, output) - image.save(dest) + dest = save_image(ctx, output, image) + save_params(ctx, output, params, size, upscale=upscale) del image del result @@ -97,8 +100,9 @@ def run_img2img_pipeline( image = run_upscale_correction( ctx, StageParams(), params, image, upscale=upscale) - dest = base_join(ctx.output_path, output) - image.save(dest) + dest = save_image(ctx, output, image) + size = Size(*source_image.size) + save_params(ctx, output, params, size, upscale=upscale) del image del result @@ -110,12 +114,12 @@ def run_img2img_pipeline( def run_inpaint_pipeline( ctx: ServerContext, params: ImageParams, - _size: Size, + size: Size, output: str, upscale: UpscaleParams, source_image: Image.Image, mask_image: Image.Image, - expand: Border, + border: Border, noise_source: Any, mask_filter: Any, strength: float, @@ -127,7 +131,7 @@ def run_inpaint_pipeline( stage, params, source_image, - border=expand, + border=border, mask_image=mask_image, fill_color=fill_color, mask_filter=mask_filter, @@ -144,8 +148,8 @@ def run_inpaint_pipeline( image = run_upscale_correction( ctx, stage, params, image, upscale=upscale) - dest = base_join(ctx.output_path, output) - image.save(dest) + dest = save_image(ctx, output, image) + save_params(ctx, output, params, size, upscale=upscale, border=border) del image run_gc() @@ -156,7 +160,7 @@ def run_inpaint_pipeline( def run_upscale_pipeline( ctx: ServerContext, params: ImageParams, - _size: Size, + size: Size, output: str, upscale: UpscaleParams, source_image: Image.Image, @@ -164,8 +168,8 @@ def run_upscale_pipeline( image = run_upscale_correction( ctx, StageParams(), params, source_image, upscale=upscale) - dest = base_join(ctx.output_path, output) - image.save(dest) + dest = save_image(ctx, output, image) + save_params(ctx, output, params, size, upscale=upscale) del image run_gc() diff --git a/api/onnx_web/output.py b/api/onnx_web/output.py new file mode 100644 index 00000000..12e6cd2e --- /dev/null +++ b/api/onnx_web/output.py @@ -0,0 +1,50 @@ +from json import dumps +from PIL import Image +from typing import Any, Optional + +from .params import ( + Border, + ImageParams, + Size, + UpscaleParams, +) +from .utils import ( + base_join, + ServerContext, +) + + +def json_params( + output: str, + params: ImageParams, + size: Size, + upscale: Optional[UpscaleParams] = None, + border: Optional[Border] = None, +) -> Any: + return { + 'border': border.tojson(), + 'output': output, + 'params': params.tojson(), + 'size': upscale.resize(size.add_border(border)).tojson(), + 'upscale': upscale.tojson(), + } + + +def save_image(ctx: ServerContext, output: str, image: Image.Image) -> str: + path = base_join(ctx.output_path, output) + image.save(path) + return path + + +def save_params( + ctx: ServerContext, + output: str, + params: ImageParams, + size: Size, + upscale: Optional[UpscaleParams] = None, + border: Optional[Border] = None, +) -> str: + path = base_join(ctx.output_path, '%s.json' % (output)) + json = json_params(output, params, size, upscale=upscale, border=border) + with open(path, 'w') as f: + f.write(dumps(json)) diff --git a/api/onnx_web/serve.py b/api/onnx_web/serve.py index c16a1d18..7b0250c6 100644 --- a/api/onnx_web/serve.py +++ b/api/onnx_web/serve.py @@ -60,6 +60,11 @@ from .image import ( noise_source_normal, noise_source_uniform, ) +from .output import ( + json_params, + save_image, + save_params, +) from .params import ( Border, ImageParams, @@ -437,11 +442,7 @@ def img2img(): executor.submit_stored(output, run_img2img_pipeline, context, params, output, upscale, source_image, strength) - return jsonify({ - 'output': output, - 'params': params.tojson(), - 'size': upscale.resize(size).tojson(), - }) + return jsonify(json_params(output, params, size, upscale=upscale)) @app.route('/api/txt2img', methods=['POST']) @@ -458,11 +459,7 @@ def txt2img(): executor.submit_stored( output, run_txt2img_pipeline, context, params, size, output, upscale) - return jsonify({ - 'output': output, - 'params': params.tojson(), - 'size': upscale.resize(size).tojson(), - }) + return jsonify(json_params(output, params, size, upscale=upscale)) @app.route('/api/inpaint', methods=['POST']) @@ -529,11 +526,7 @@ def inpaint(): strength, fill_color) - return jsonify({ - 'output': output, - 'params': params.tojson(), - 'size': upscale.resize(size.add_border(expand)).tojson(), - }) + return jsonify(json_params(output, params, size, upscale=upscale, border=expand)) @app.route('/api/upscale', methods=['POST']) @@ -557,11 +550,7 @@ def upscale(): executor.submit_stored(output, run_upscale_pipeline, context, params, size, output, upscale, source_image) - return jsonify({ - 'output': output, - 'params': params.tojson(), - 'size': upscale.resize(size).tojson(), - }) + return jsonify(json_params(output, params, size, upscale=upscale)) @app.route('/api/chain', methods=['POST']) @@ -607,11 +596,7 @@ def chain(): executor.submit_stored(output, pipeline, context, params, fake_source, output=output, size=size) - return jsonify({ - 'output': output, - 'params': params.tojson(), - 'size': size.tojson(), - }) + return jsonify(json_params(output, params, size)) @app.route('/api/ready')