From ff70d36a058ddc416777242a5b70d85984a75143 Mon Sep 17 00:00:00 2001 From: Sean Sube Date: Sun, 15 Jan 2023 19:14:58 -0600 Subject: [PATCH] lint(api): add class for params, size, other common data --- api/onnx_web/image.py | 33 +++++++------- api/onnx_web/pipeline.py | 95 ++++++++++++++++++---------------------- api/onnx_web/serve.py | 46 ++++++++++--------- api/onnx_web/utils.py | 35 ++++++++++++++- 4 files changed, 116 insertions(+), 93 deletions(-) diff --git a/api/onnx_web/image.py b/api/onnx_web/image.py index cfac3ee0..9b0a5e4b 100644 --- a/api/onnx_web/image.py +++ b/api/onnx_web/image.py @@ -4,8 +4,13 @@ from typing import Tuple import numpy as np +from .utils import ( + Border, + Point, +) -def mask_filter_none(mask_image: Image, dims: Tuple[int, int], origin: Tuple[int, int], fill='white') -> Image: + +def mask_filter_none(mask_image: Image, dims: Point, origin: Point, fill='white') -> Image: width, height = dims noise = Image.new('RGB', (width, height), fill) @@ -14,7 +19,7 @@ def mask_filter_none(mask_image: Image, dims: Tuple[int, int], origin: Tuple[int return noise -def mask_filter_gaussian_multiply(mask_image: Image, dims: Tuple[int, int], origin: Tuple[int, int], rounds=3) -> Image: +def mask_filter_gaussian_multiply(mask_image: Image, dims: Point, origin: Point, rounds=3) -> Image: ''' Gaussian blur with multiply, source image centered on white canvas. ''' @@ -27,7 +32,7 @@ def mask_filter_gaussian_multiply(mask_image: Image, dims: Tuple[int, int], orig return noise -def mask_filter_gaussian_screen(mask_image: Image, dims: Tuple[int, int], origin: Tuple[int, int], rounds=3) -> Image: +def mask_filter_gaussian_screen(mask_image: Image, dims: Point, origin: Point, rounds=3) -> Image: ''' Gaussian blur, source image centered on white canvas. ''' @@ -40,7 +45,7 @@ def mask_filter_gaussian_screen(mask_image: Image, dims: Tuple[int, int], origin return noise -def noise_source_fill_edge(source_image: Image, dims: Tuple[int, int], origin: Tuple[int, int], fill='white') -> Image: +def noise_source_fill_edge(source_image: Image, dims: Point, origin: Point, fill='white') -> Image: ''' Identity transform, source image centered on white canvas. ''' @@ -52,7 +57,7 @@ def noise_source_fill_edge(source_image: Image, dims: Tuple[int, int], origin: T return noise -def noise_source_fill_mask(source_image: Image, dims: Tuple[int, int], origin: Tuple[int, int], fill='white') -> Image: +def noise_source_fill_mask(source_image: Image, dims: Point, origin: Point, fill='white') -> Image: ''' Fill the whole canvas, no source or noise. ''' @@ -63,7 +68,7 @@ def noise_source_fill_mask(source_image: Image, dims: Tuple[int, int], origin: T return noise -def noise_source_gaussian(source_image: Image, dims: Tuple[int, int], origin: Tuple[int, int], rounds=3) -> Image: +def noise_source_gaussian(source_image: Image, dims: Point, origin: Point, rounds=3) -> Image: ''' Gaussian blur, source image centered on white canvas. ''' @@ -76,7 +81,7 @@ def noise_source_gaussian(source_image: Image, dims: Tuple[int, int], origin: Tu return noise -def noise_source_uniform(source_image: Image, dims: Tuple[int, int], origin: Tuple[int, int]) -> Image: +def noise_source_uniform(source_image: Image, dims: Point, origin: Point) -> Image: width, height = dims size = width * height @@ -98,7 +103,7 @@ def noise_source_uniform(source_image: Image, dims: Tuple[int, int], origin: Tup return noise -def noise_source_normal(source_image: Image, dims: Tuple[int, int], origin: Tuple[int, int]) -> Image: +def noise_source_normal(source_image: Image, dims: Point, origin: Point) -> Image: width, height = dims size = width * height @@ -120,7 +125,7 @@ def noise_source_normal(source_image: Image, dims: Tuple[int, int], origin: Tupl return noise -def noise_source_histogram(source_image: Image, dims: Tuple[int, int], origin: Tuple[int, int]) -> Image: +def noise_source_histogram(source_image: Image, dims: Point, origin: Point) -> Image: r, g, b = source_image.split() width, height = dims size = width * height @@ -154,18 +159,16 @@ def noise_source_histogram(source_image: Image, dims: Tuple[int, int], origin: T def expand_image( source_image: Image, mask_image: Image, - expand_by: Tuple[int, int, int, int], + expand: Border, fill='white', noise_source=noise_source_histogram, mask_filter=mask_filter_none, ): - left, right, top, bottom = expand_by - - full_width = left + source_image.width + right - full_height = top + source_image.height + bottom + full_width = expand.left + source_image.width + expand.right + full_height = expand.top + source_image.height + expand.bottom dims = (full_width, full_height) - origin = (top, left) + origin = (expand.top, expand.left) full_source = Image.new('RGB', dims, fill) full_source.paste(source_image, origin) diff --git a/api/onnx_web/pipeline.py b/api/onnx_web/pipeline.py index 8e789e17..00056643 100644 --- a/api/onnx_web/pipeline.py +++ b/api/onnx_web/pipeline.py @@ -7,7 +7,7 @@ from diffusers import ( ) from os import environ from PIL import Image -from typing import Any, Union +from typing import Any import numpy as np @@ -18,7 +18,10 @@ from .upscale import ( upscale_resrgan, ) from .utils import ( - safer_join + safer_join, + BaseParams, + Border, + Size, ) last_pipeline_instance = None @@ -28,9 +31,9 @@ last_pipeline_scheduler = None # from https://www.travelneil.com/stable-diffusion-updates.html -def get_latents_from_seed(seed: int, width: int, height: int) -> np.ndarray: +def get_latents_from_seed(seed: int, size: Size) -> np.ndarray: # 1 is batch size - latents_shape = (1, 4, height // 8, width // 8) + latents_shape = (1, 4, size.height // 8, size.width // 8) # Gotta use numpy instead of torch, because torch's randn() doesn't support DML rng = np.random.default_rng(seed) image_latents = rng.standard_normal(latents_shape).astype(np.float32) @@ -67,82 +70,70 @@ def load_pipeline(pipeline: DiffusionPipeline, model: str, provider: str, schedu return pipe -def run_txt2img_pipeline(model, provider, scheduler, prompt, negative_prompt, cfg, steps, seed, output, height, width): +def run_txt2img_pipeline(params: BaseParams, size: Size): pipe = load_pipeline(OnnxStableDiffusionPipeline, - model, provider, scheduler) + params.model, params.provider, params.scheduler) - latents = get_latents_from_seed(seed, width, height) - rng = np.random.RandomState(seed) + latents = get_latents_from_seed(params.seed, size.width, size.height) + rng = np.random.RandomState(params.seed) image = pipe( - prompt, - height, - width, + params.prompt, + size.width, + size.height, generator=rng, - guidance_scale=cfg, + guidance_scale=params.cfg, latents=latents, - negative_prompt=negative_prompt, - num_inference_steps=steps, + negative_prompt=params.negative_prompt, + num_inference_steps=params.steps, ).images[0] image = upscale_resrgan(image, model_path) - image.save(output) + image.save(params.output.file) - print('saved txt2img output: %s' % (output)) + print('saved txt2img output: %s' % (params.output.file)) -def run_img2img_pipeline(model, provider, scheduler, prompt, negative_prompt, cfg, steps, seed, output, strength, input_image): +def run_img2img_pipeline(params: BaseParams, strength, input_image): pipe = load_pipeline(OnnxStableDiffusionImg2ImgPipeline, - model, provider, scheduler) + params.model, params.provider, params.scheduler) - rng = np.random.RandomState(seed) + rng = np.random.RandomState(params.seed) image = pipe( - prompt, + params.prompt, generator=rng, - guidance_scale=cfg, + guidance_scale=params.cfg, image=input_image, - negative_prompt=negative_prompt, - num_inference_steps=steps, + negative_prompt=params.negative_prompt, + num_inference_steps=params.steps, strength=strength, ).images[0] image = upscale_resrgan(image, model_path) - image.save(output) + image.save(params.output.file) - print('saved img2img output: %s' % (output)) + print('saved img2img output: %s' % (params.output.file)) 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, + params: BaseParams, + size: Size, source_image: Image, mask_image: Image, - left: int, - right: int, - top: int, - bottom: int, + expand: Border, noise_source: Any, mask_filter: Any ): pipe = load_pipeline(OnnxStableDiffusionInpaintPipeline, - model, provider, scheduler) + params.model, params.provider, params.scheduler) - latents = get_latents_from_seed(seed, width, height) - rng = np.random.RandomState(seed) + latents = get_latents_from_seed(params.seed, size) + rng = np.random.RandomState(params.seed) print('applying mask filter and generating noise source') source_image, mask_image, noise_image, _full_dims = expand_image( source_image, mask_image, - (left, right, top, bottom), + expand, noise_source=noise_source, mask_filter=mask_filter) @@ -152,18 +143,18 @@ def run_inpaint_pipeline( noise_image.save(safer_join(output_path, 'last-noise.png')) image = pipe( - prompt, + params.prompt, generator=rng, - guidance_scale=cfg, - height=height, + guidance_scale=params.cfg, + height=size.height, image=source_image, latents=latents, mask_image=mask_image, - negative_prompt=negative_prompt, - num_inference_steps=steps, - width=width, + negative_prompt=params.negative_prompt, + num_inference_steps=params.steps, + width=size.width, ).images[0] - image.save(output) + image.save(params.output.file) - print('saved inpaint output: %s' % (output)) + print('saved inpaint output: %s' % (params.output.file)) diff --git a/api/onnx_web/serve.py b/api/onnx_web/serve.py index b189b088..d280876c 100644 --- a/api/onnx_web/serve.py +++ b/api/onnx_web/serve.py @@ -46,7 +46,10 @@ from .utils import ( get_and_clamp_float, get_and_clamp_int, get_from_map, - safer_join + safer_join, + BaseParams, + OutputPath, + Size, ) import json @@ -111,11 +114,14 @@ def serve_bundle_file(filename='index.html'): return send_from_directory(path.join('..', bundle_path), filename) -def make_output_path(mode: str, seed: int, params: Tuple[Union[str, int, float]]): +def make_output_path(mode: str, params: BaseParams, size: Size, extras: Tuple[Union[str, int, float]]) -> OutputPath: now = int(time.time()) sha = sha256() sha.update(mode.encode('utf-8')) + # TODO: add params + # TODO: add size + for param in params: if param is None: continue @@ -128,10 +134,10 @@ def make_output_path(mode: str, seed: int, params: Tuple[Union[str, int, float]] else: print('cannot hash param: %s, %s' % (param, type(param))) - output_file = '%s_%s_%s_%s.png' % (mode, seed, sha.hexdigest(), now) + output_file = '%s_%s_%s_%s.png' % (mode, params.seed, sha.hexdigest(), now) output_full = safer_join(output_path, output_file) - return (output_file, output_full) + return OutputPath(output_full, output_file) def url_from_rule(rule): @@ -142,7 +148,7 @@ def url_from_rule(rule): return url_for(rule.endpoint, **options) -def pipeline_from_request(): +def pipeline_from_request() -> Tuple[BaseParams, Size]: user = request.remote_addr # pipeline stuff @@ -189,7 +195,9 @@ def pipeline_from_request(): print("request from %s: %s rounds of %s using %s on %s, %sx%s, %s, %s - %s" % (user, steps, scheduler.__name__, model, provider, width, height, cfg, seed, prompt)) - return (model, provider, scheduler, prompt, negative_prompt, cfg, steps, height, width, seed) + params = BaseParams(model, provider, scheduler, prompt, negative_prompt, cfg, steps, seed) + size = Size(width, height) + return (params, size) def check_paths(): @@ -283,27 +291,17 @@ def img2img(): strength = get_and_clamp_float(request.args, 'strength', 0.5, 1.0) - (model, provider, scheduler, prompt, negative_prompt, cfg, steps, height, - width, seed) = pipeline_from_request() + params, size = pipeline_from_request() - (output_file, output_full) = make_output_path( + output = make_output_path( 'img2img', - seed, ( - model, - provider, - scheduler.__name__, - prompt, - negative_prompt, - cfg, - steps, - strength, - height, - width)) - print("img2img output: %s" % (output_full)) + params, + size, + extras=(strength)) + print("img2img output: %s" % (output.path)) - input_image.thumbnail((width, height)) - executor.submit_stored(output_file, run_img2img_pipeline, model, provider, - scheduler, prompt, negative_prompt, cfg, steps, seed, output_full, strength, input_image) + input_image.thumbnail((size.width, size.height)) + executor.submit_stored(output.file, run_img2img_pipeline, params, output, strength, input_image) return jsonify({ 'output': output_file, diff --git a/api/onnx_web/utils.py b/api/onnx_web/utils.py index 67092523..bf6ce51e 100644 --- a/api/onnx_web/utils.py +++ b/api/onnx_web/utils.py @@ -2,9 +2,7 @@ from os import path from typing import Any, Dict, Tuple -Border = Tuple[int, int, int, int] Point = Tuple[int, int] -Size = Tuple[int, int] def get_and_clamp_float(args, key: str, default_value: float, max_value: float, min_value=0.0) -> float: @@ -26,3 +24,36 @@ def get_from_map(args, key: str, values: Dict[str, Any], default: Any): def safer_join(base, tail): safer_path = path.relpath(path.normpath(path.join('/', tail)), '/') return path.join(base, safer_path) + + +class OutputPath: + def __init__(self, path, file): + self.path = path + self.file = file + + +class BaseParams: + def __init__(self, model, provider, scheduler, prompt, negative_prompt, cfg, steps, seed): + self.model = model + self.provider = provider + self.scheduler = scheduler + self.prompt = prompt + self.negative_prompt = negative_prompt + self.cfg = cfg + self.steps = steps + self.seed = seed + self.output = output + + +class Border: + def __init__(self, left, right, top, bottom): + self.left = left + self.right = right + self.top = top + self.bottom = bottom + + +class Size: + def __init__(self, width, height): + self.width = width + self.height = height