From 44c81d966ffefbfa0f5af0b13d0982bdb45e6ab0 Mon Sep 17 00:00:00 2001 From: Sean Sube Date: Sun, 15 Jan 2023 19:33:40 -0600 Subject: [PATCH] lint(api): finish splitting up modules --- api/onnx_web/pipeline.py | 31 +++++-- api/onnx_web/serve.py | 178 +++++++++++++++------------------------ api/onnx_web/utils.py | 20 ++++- 3 files changed, 109 insertions(+), 120 deletions(-) diff --git a/api/onnx_web/pipeline.py b/api/onnx_web/pipeline.py index 00056643..59c0038a 100644 --- a/api/onnx_web/pipeline.py +++ b/api/onnx_web/pipeline.py @@ -21,12 +21,14 @@ from .utils import ( safer_join, BaseParams, Border, + OutputPath, Size, ) last_pipeline_instance = None last_pipeline_options = (None, None, None) last_pipeline_scheduler = None +model_path = None # from https://www.travelneil.com/stable-diffusion-updates.html @@ -40,6 +42,16 @@ def get_latents_from_seed(seed: int, size: Size) -> np.ndarray: return image_latents +def get_model_path(model: str): + return safer_join(model_path, model) + + +# TODO: hax +def set_model_path(model: str): + global model_path + model_path = model + + def load_pipeline(pipeline: DiffusionPipeline, model: str, provider: str, scheduler): global last_pipeline_instance global last_pipeline_scheduler @@ -70,7 +82,7 @@ def load_pipeline(pipeline: DiffusionPipeline, model: str, provider: str, schedu return pipe -def run_txt2img_pipeline(params: BaseParams, size: Size): +def run_txt2img_pipeline(params: BaseParams, size: Size, output: OutputPath): pipe = load_pipeline(OnnxStableDiffusionPipeline, params.model, params.provider, params.scheduler) @@ -88,12 +100,12 @@ def run_txt2img_pipeline(params: BaseParams, size: Size): num_inference_steps=params.steps, ).images[0] image = upscale_resrgan(image, model_path) - image.save(params.output.file) + image.save(output.file) - print('saved txt2img output: %s' % (params.output.file)) + print('saved txt2img output: %s' % (output.file)) -def run_img2img_pipeline(params: BaseParams, strength, input_image): +def run_img2img_pipeline(params: BaseParams, output: OutputPath, strength: float, input_image: Image): pipe = load_pipeline(OnnxStableDiffusionImg2ImgPipeline, params.model, params.provider, params.scheduler) @@ -117,6 +129,7 @@ def run_img2img_pipeline(params: BaseParams, strength, input_image): def run_inpaint_pipeline( params: BaseParams, size: Size, + output: OutputPath, source_image: Image, mask_image: Image, expand: Border, @@ -138,9 +151,9 @@ def run_inpaint_pipeline( mask_filter=mask_filter) if environ.get('DEBUG') is not None: - source_image.save(safer_join(output_path, 'last-source.png')) - mask_image.save(safer_join(output_path, 'last-mask.png')) - noise_image.save(safer_join(output_path, 'last-noise.png')) + source_image.save(safer_join(output.path, 'last-source.png')) + mask_image.save(safer_join(output.path, 'last-mask.png')) + noise_image.save(safer_join(output.path, 'last-noise.png')) image = pipe( params.prompt, @@ -155,6 +168,6 @@ def run_inpaint_pipeline( width=size.width, ).images[0] - image.save(params.output.file) + image.save(output.file) - print('saved inpaint output: %s' % (params.output.file)) + print('saved inpaint output: %s' % (output.file)) diff --git a/api/onnx_web/serve.py b/api/onnx_web/serve.py index d280876c..1ea12d1a 100644 --- a/api/onnx_web/serve.py +++ b/api/onnx_web/serve.py @@ -24,7 +24,6 @@ from os import environ, makedirs, path, scandir from typing import Tuple, Union from .image import ( - expand_image, # mask filters mask_filter_gaussian_multiply, mask_filter_gaussian_screen, @@ -38,9 +37,11 @@ from .image import ( noise_source_uniform, ) from .pipeline import ( + get_model_path, run_img2img_pipeline, run_inpaint_pipeline, run_txt2img_pipeline, + set_model_path, ) from .utils import ( get_and_clamp_float, @@ -48,6 +49,7 @@ from .utils import ( get_from_map, safer_join, BaseParams, + Border, OutputPath, Size, ) @@ -106,33 +108,41 @@ mask_filters = { } -def get_model_path(model: str): - return safer_join(model_path, model) - - def serve_bundle_file(filename='index.html'): return send_from_directory(path.join('..', bundle_path), filename) +def hash_value(sha, param): + if param is None: + return + elif isinstance(param, float): + sha.update(bytearray(pack('!f', param))) + elif isinstance(param, int): + sha.update(bytearray(pack('!I', param))) + elif isinstance(param, str): + sha.update(param.encode('utf-8')) + else: + print('cannot hash param: %s, %s' % (param, type(param))) + + 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 + hash_value(mode) + hash_value(params.model) + hash_value(params.provider) + hash_value(params.scheduler) + hash_value(params.prompt) + hash_value(params.negative_prompt) + hash_value(params.cfg) + hash_value(params.steps) + hash_value(params.seed) + hash_value(size.width) + hash_value(size.height) - for param in params: - if param is None: - continue - elif isinstance(param, float): - sha.update(bytearray(pack('!f', param))) - elif isinstance(param, int): - sha.update(bytearray(pack('!I', param))) - elif isinstance(param, str): - sha.update(param.encode('utf-8')) - else: - print('cannot hash param: %s, %s' % (param, type(param))) + for param in extras: + hash_value(sha, param) output_file = '%s_%s_%s_%s.png' % (mode, params.seed, sha.hexdigest(), now) output_full = safer_join(output_path, output_file) @@ -140,7 +150,7 @@ def make_output_path(mode: str, params: BaseParams, size: Size, extras: Tuple[Un return OutputPath(output_full, output_file) -def url_from_rule(rule): +def url_from_rule(rule) -> str: options = {} for arg in rule.arguments: options[arg] = ":%s" % (arg) @@ -195,7 +205,8 @@ def pipeline_from_request() -> Tuple[BaseParams, Size]: 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)) - params = BaseParams(model, provider, scheduler, prompt, negative_prompt, cfg, steps, seed) + params = BaseParams(model, provider, scheduler, prompt, + negative_prompt, cfg, steps, seed) size = Size(width, height) return (params, size) @@ -220,6 +231,7 @@ def load_params(): check_paths() +set_model_path(model_path) load_models() load_params() @@ -301,61 +313,33 @@ def img2img(): print("img2img output: %s" % (output.path)) input_image.thumbnail((size.width, size.height)) - executor.submit_stored(output.file, run_img2img_pipeline, params, output, strength, input_image) + executor.submit_stored(output.file, run_img2img_pipeline, + params, output, strength, input_image) return jsonify({ - 'output': output_file, - 'params': { - 'model': model, - 'provider': provider, - 'scheduler': scheduler.__name__, - 'seed': seed, - 'prompt': prompt, - 'cfg': cfg, - 'negativePrompt': negative_prompt, - 'steps': steps, - 'height': height, - 'width': width, - } + 'output': output.file, + 'params': params.tojson(), + 'size': size.tojson(), }) @app.route('/api/txt2img', methods=['POST']) def txt2img(): - (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( 'txt2img', - seed, ( - model, - provider, - scheduler.__name__, - prompt, - negative_prompt, - cfg, - steps, - height, - width)) - print("txt2img output: %s" % (output_full)) + params, + size) + print("txt2img output: %s" % (output.file)) - executor.submit_stored(output_file, run_txt2img_pipeline, model, - provider, scheduler, prompt, negative_prompt, cfg, steps, seed, output_full, height, width) + executor.submit_stored( + output.file, run_txt2img_pipeline, params, size, output) return jsonify({ - 'output': output_file, - 'params': { - 'model': model, - 'provider': provider, - 'scheduler': scheduler.__name__, - 'seed': seed, - 'prompt': prompt, - 'cfg': cfg, - 'negativePrompt': negative_prompt, - 'steps': steps, - 'height': height, - 'width': width, - } + 'output': output.file, + 'params': params.tojson(), + 'size': size.tojson(), }) @@ -367,8 +351,7 @@ def inpaint(): mask_file = request.files.get('mask') mask_image = Image.open(BytesIO(mask_file.read())).convert('RGB') - (model, provider, scheduler, prompt, negative_prompt, cfg, steps, height, - width, seed) = pipeline_from_request() + params, size = pipeline_from_request() left = get_and_clamp_int(request.args, 'left', 0, config_params.get('width').get('max'), 0) @@ -378,70 +361,45 @@ def inpaint(): config_params.get('height').get('max'), 0) bottom = get_and_clamp_int( request.args, 'bottom', 0, config_params.get('height').get('max'), 0) + expand = Border(left, right, top, bottom) mask_filter = get_from_map(request.args, 'filter', mask_filters, 'none') noise_source = get_from_map( request.args, 'noise', noise_sources, 'histogram') - (output_file, output_full) = make_output_path( - 'inpaint', seed, ( - model, - provider, - scheduler.__name__, - prompt, - negative_prompt, - cfg, - steps, - height, - width, + output = make_output_path( + 'inpaint', + params, + size, + extras=( left, right, top, bottom, mask_filter.__name__, noise_source.__name__, - )) - print("inpaint output: %s" % output_full) + ) + ) + print("inpaint output: %s" % output.file) - source_image.thumbnail((width, height)) - mask_image.thumbnail((width, height)) + source_image.thumbnail((size.width, size.height)) + mask_image.thumbnail((size.width, size.height)) executor.submit_stored( - output_file, + output.file, run_inpaint_pipeline, - model, - provider, - scheduler, - prompt, - negative_prompt, - cfg, - steps, - seed, - output_full, - height, - width, + params, + size, + output, source_image, mask_image, - left, - right, - top, - bottom, + expand, noise_source, mask_filter) return jsonify({ - 'output': output_file, - 'params': { - 'model': model, - 'provider': provider, - 'scheduler': scheduler.__name__, - 'seed': seed, - 'prompt': prompt, - 'cfg': cfg, - 'negativePrompt': negative_prompt, - 'steps': steps, - 'height': height, - 'width': width, - } + 'output': output.file, + 'params': params.tojson(), + 'size': size.tojson(), }) diff --git a/api/onnx_web/utils.py b/api/onnx_web/utils.py index bf6ce51e..38125e1e 100644 --- a/api/onnx_web/utils.py +++ b/api/onnx_web/utils.py @@ -26,6 +26,7 @@ def safer_join(base, tail): return path.join(base, safer_path) +# TODO: .path is only used in one place, can probably just be a str class OutputPath: def __init__(self, path, file): self.path = path @@ -42,7 +43,18 @@ class BaseParams: self.cfg = cfg self.steps = steps self.seed = seed - self.output = output + + def tojson(self) -> Dict[str, Any]: + return { + 'model': self.model, + 'provider': self.provider, + 'scheduler': self.scheduler.__name__, + 'seed': self.seed, + 'prompt': self.prompt, + 'cfg': self.cfg, + 'negativePrompt': self.negative_prompt, + 'steps': self.steps, + } class Border: @@ -57,3 +69,9 @@ class Size: def __init__(self, width, height): self.width = width self.height = height + + def tojson(self) -> Dict[str, Any]: + return { + 'height': self.height, + 'width': self.width, + }