diff --git a/api/onnx_web/image.py b/api/onnx_web/image.py index 7ec8a331..d0f96aee 100644 --- a/api/onnx_web/image.py +++ b/api/onnx_web/image.py @@ -13,9 +13,26 @@ def blend_imult(a): return 1.0 - blend_mult(a) -def blend_mask_inverse_source(source: Tuple[int, int, int], mask: Tuple[int, int, int], noise: Tuple[int, int, int]) -> Tuple[int, int, int]: +def blend_mask_source(source: Tuple[int, int, int], mask: Tuple[int, int, int], noise: Tuple[int, int, int]) -> Tuple[int, int, int]: + ''' + Blend operation, linear interpolation from noise to source based on mask: `(s * (1 - m)) + (n * m)` + Black = noise + White = source + ''' + return ( + int((source[0] * blend_mult(mask[0])) + + (noise[0] * blend_imult(mask[0]))), + int((source[1] * blend_mult(mask[1])) + + (noise[1] * blend_imult(mask[1]))), + int((source[2] * blend_mult(mask[2])) + + (noise[2] * blend_imult(mask[2]))), + ) + +def blend_source_mask(source: Tuple[int, int, int], mask: Tuple[int, int, int], noise: Tuple[int, int, int]) -> Tuple[int, int, int]: ''' Blend operation, linear interpolation from source to noise based on mask: `(s * (1 - m)) + (n * m)` + Black = source + White = noise ''' return ( int((source[0] * blend_imult(mask[0])) + @@ -27,21 +44,29 @@ def blend_mask_inverse_source(source: Tuple[int, int, int], mask: Tuple[int, int ) -def noise_source_original(source_image: Image, dims: Tuple[int, int], origin: Tuple[int, int]) -> Tuple[float, float, float]: +def noise_source_fill(source_image: Image, dims: Tuple[int, int], origin: Tuple[int, int], fill='white') -> Tuple[float, float, float]: + ''' + Identity transform, source image centered on white canvas. + ''' width, height = dims - noise = Image.new('RGB', (width, height), 'white') + noise = Image.new('RGB', (width, height), fill) noise.paste(source_image, origin) return noise -def noise_source_gaussian(source_image: Image, dims: Tuple[int, int], origin: Tuple[int, int]) -> Tuple[float, float, float]: +def noise_source_gaussian(source_image: Image, dims: Tuple[int, int], origin: Tuple[int, int], rounds=3) -> Tuple[float, float, float]: + ''' + Gaussian blur, source image centered on white canvas. + ''' width, height = dims noise = Image.new('RGB', (width, height), 'white') noise.paste(source_image, origin) - noise.filter(ImageFilter.GaussianBlur(5)) + + for i in range(rounds): + noise.filter(ImageFilter.GaussianBlur(5)) return noise @@ -127,7 +152,7 @@ def expand_image( expand_by: Tuple[int, int, int, int], fill='white', noise_source=noise_source_histogram, - blend_op=blend_mask_inverse_source + blend_op=blend_source_mask ): left, right, top, bottom = expand_by diff --git a/api/onnx_web/serve.py b/api/onnx_web/serve.py index 3a53e978..aa17b7f9 100644 --- a/api/onnx_web/serve.py +++ b/api/onnx_web/serve.py @@ -29,7 +29,7 @@ from struct import pack from os import environ, makedirs, path, scandir from typing import Any, Dict, Tuple, Union -from .image import expand_image +from .image import expand_image, noise_source_gaussian, noise_source_histogram, noise_source_normal, noise_source_fill, noise_source_uniform, blend_source_mask, blend_imult, blend_mask_source, blend_mult import json import numpy as np @@ -72,6 +72,17 @@ pipeline_schedulers = { 'lms-discrete': LMSDiscreteScheduler, 'pndm': PNDMScheduler, } +noise_sources = { + 'fill': noise_source_fill, + 'gaussian': noise_source_gaussian, + 'histogram': noise_source_histogram, + 'normal': noise_source_normal, + 'uniform': noise_source_uniform, +} +blend_modes = { + 'mask-source': blend_mask_source, + 'source-mask': blend_source_mask, +} def get_and_clamp_float(args, key: str, default_value: float, max_value: float, min_value=0.0) -> float: @@ -264,7 +275,27 @@ def run_img2img_pipeline(model, provider, scheduler, prompt, negative_prompt, cf print('saved img2img output: %s' % (output)) -def run_inpaint_pipeline(model, provider, scheduler, prompt, negative_prompt, cfg, steps, seed, output, height, width, source_image, mask_image, left, right, top, bottom): +def run_inpaint_pipeline( + model: str, + provider: str, + scheduler: Any, + prompt: str, + negative_prompt: Union[str, None], + cfg: float, + steps: int, + seed: int, + output: str, + height: int, + width: int, + source_image: Image, + mask_image: Image, + left: int, + right: int, + top: int, + bottom: int, + noise_source: Any, + blend_op: Any +): pipe = load_pipeline(OnnxStableDiffusionInpaintPipeline, model, provider, scheduler) @@ -273,7 +304,12 @@ def run_inpaint_pipeline(model, provider, scheduler, prompt, negative_prompt, cf if left > 0 or right > 0 or top > 0 or bottom > 0: print('expanding image for outpainting') - source_image, mask_image, _full_noise, _full_dims = expand_image(source_image, mask_image, (left, right, top, bottom)) + source_image, mask_image, _full_noise, _full_dims = expand_image( + source_image, + mask_image, + (left, right, top, bottom), + noise_source=noise_source, + blend_op=blend_op) image = pipe( prompt, @@ -445,10 +481,18 @@ def inpaint(): (model, provider, scheduler, prompt, negative_prompt, cfg, steps, height, width, seed) = pipeline_from_request() - left = get_and_clamp_int(request.args, 'left', 0, config_params.get('width').get('max'), 0) - right = get_and_clamp_int(request.args, 'right', 0, config_params.get('width').get('max'), 0) - top = get_and_clamp_int(request.args, 'top', 0, config_params.get('height').get('max'), 0) - bottom = get_and_clamp_int(request.args, 'bottom', 0, config_params.get('height').get('max'), 0) + left = get_and_clamp_int(request.args, 'left', 0, + config_params.get('width').get('max'), 0) + right = get_and_clamp_int(request.args, 'right', + 0, config_params.get('width').get('max'), 0) + top = get_and_clamp_int(request.args, 'top', 0, + config_params.get('height').get('max'), 0) + bottom = get_and_clamp_int( + request.args, 'bottom', 0, config_params.get('height').get('max'), 0) + + noise_source = get_from_map( + request.args, 'noise', noise_sources, 'histogram') + blend_op = get_from_map(request.args, 'blend', blend_modes, 'mask-source') (output_file, output_full) = make_output_path( 'inpaint', seed, (prompt, cfg, steps, height, width, seed, left, right, top, bottom)) @@ -456,8 +500,28 @@ def inpaint(): source_image.thumbnail((width, height)) mask_image.thumbnail((width, height)) - executor.submit_stored(output_file, run_inpaint_pipeline, model, provider, scheduler, prompt, negative_prompt, - cfg, steps, seed, output_full, height, width, source_image, mask_image, left, right, top, bottom) + executor.submit_stored( + output_file, + run_inpaint_pipeline, + model, + provider, + scheduler, + prompt, + negative_prompt, + cfg, + steps, + seed, + output_full, + height, + width, + source_image, + mask_image, + left, + right, + top, + bottom, + noise_source, + blend_op) return jsonify({ 'output': output_file,