diff --git a/api/onnx_web/diffusion/run.py b/api/onnx_web/diffusion/run.py index d2694d8b..2a5b289f 100644 --- a/api/onnx_web/diffusion/run.py +++ b/api/onnx_web/diffusion/run.py @@ -71,7 +71,7 @@ def run_txt2img_pipeline( del result run_gc() - logger.info('saved txt2img output: %s', dest) + logger.info('finished txt2img job: %s', dest) def run_img2img_pipeline( @@ -108,7 +108,7 @@ def run_img2img_pipeline( del result run_gc() - logger.info('saved img2img output: %s', dest) + logger.info('finished img2img job: %s', dest) def run_inpaint_pipeline( @@ -154,7 +154,7 @@ def run_inpaint_pipeline( del image run_gc() - logger.info('saved inpaint output: %s', dest) + logger.info('finished inpaint job: %s', dest) def run_upscale_pipeline( @@ -174,4 +174,4 @@ def run_upscale_pipeline( del image run_gc() - logger.info('saved img2img output: %s', dest) + logger.info('finished upscale job: %s', dest) diff --git a/api/onnx_web/output.py b/api/onnx_web/output.py index 11b43f87..692ab516 100644 --- a/api/onnx_web/output.py +++ b/api/onnx_web/output.py @@ -1,10 +1,15 @@ +from hashlib import sha256 from json import dumps +from logging import getLogger from PIL import Image -from typing import Any, Optional +from struct import pack +from time import time +from typing import Any, Optional, Tuple from .params import ( Border, ImageParams, + Param, Size, UpscaleParams, ) @@ -13,6 +18,20 @@ from .utils import ( ServerContext, ) +logger = getLogger(__name__) + +def hash_value(sha, param: 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: + logger.warn('cannot hash param: %s, %s', param, type(param)) + def json_params( output: str, @@ -42,9 +61,38 @@ def json_params( return json +def make_output_name( + mode: str, + params: ImageParams, + size: Size, + extras: Optional[Tuple[Param]] = None +) -> str: + now = int(time()) + sha = sha256() + + hash_value(sha, mode) + hash_value(sha, params.model) + hash_value(sha, params.provider) + hash_value(sha, params.scheduler.__name__) + hash_value(sha, params.prompt) + hash_value(sha, params.negative_prompt) + hash_value(sha, params.cfg) + hash_value(sha, params.steps) + hash_value(sha, params.seed) + hash_value(sha, size.width) + hash_value(sha, size.height) + + if extras is not None: + for param in extras: + hash_value(sha, param) + + return '%s_%s_%s_%s' % (mode, params.seed, sha.hexdigest(), now) + + def save_image(ctx: ServerContext, output: str, image: Image.Image) -> str: path = base_join(ctx.output_path, '%s.%s' % (output, ctx.image_format)) image.save(path, format=ctx.image_format) + logger.debug('saved output image to: %s', path) return path @@ -60,3 +108,5 @@ def save_params( json = json_params(output, params, size, upscale=upscale, border=border) with open(path, 'w') as f: f.write(dumps(json)) + logger.debug('saved image params to: %s', path) + return path diff --git a/api/onnx_web/serve.py b/api/onnx_web/serve.py index 7b0250c6..de35feb4 100644 --- a/api/onnx_web/serve.py +++ b/api/onnx_web/serve.py @@ -436,7 +436,7 @@ def img2img(): params, size, extras=(strength,)) - logger.info("img2img output saved: %s", output) + logger.info("img2img job queued for: %s", output) source_image.thumbnail((size.width, size.height)) executor.submit_stored(output, run_img2img_pipeline, @@ -454,7 +454,7 @@ def txt2img(): 'txt2img', params, size) - logger.info("txt2img output saved: %s", output) + logger.info("txt2img job queued for: %s", output) executor.submit_stored( output, run_txt2img_pipeline, context, params, size, output, upscale) @@ -506,7 +506,7 @@ def inpaint(): fill_color, ) ) - logger.info("inpaint output saved: %s", output) + logger.info("inpaint job queued for: %s", output) source_image.thumbnail((size.width, size.height)) mask_image.thumbnail((size.width, size.height)) diff --git a/api/onnx_web/utils.py b/api/onnx_web/utils.py index cf5b4953..8387d189 100644 --- a/api/onnx_web/utils.py +++ b/api/onnx_web/utils.py @@ -1,17 +1,11 @@ -from hashlib import sha256 from logging import getLogger from os import environ, path -from struct import pack -from time import time -from typing import Any, Dict, List, Optional, Union, Tuple +from typing import Any, Dict, List, Optional, Union import gc import torch from .params import ( - ImageParams, - Param, - Size, SizeChart, ) @@ -64,6 +58,11 @@ class ServerContext: ) +def base_join(base: str, tail: str) -> str: + tail_path = path.relpath(path.normpath(path.join('/', tail)), '/') + return path.join(base, tail_path) + + def is_debug() -> bool: return environ.get('DEBUG') is not None @@ -122,53 +121,7 @@ def get_size(val: Union[int, str, None]) -> SizeChart: raise Exception('invalid size') -def hash_value(sha, param: 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: - logger.warn('cannot hash param: %s, %s', param, type(param)) - - -def make_output_name( - mode: str, - params: ImageParams, - size: Size, - extras: Optional[Tuple[Param]] = None -) -> str: - now = int(time()) - sha = sha256() - - hash_value(sha, mode) - hash_value(sha, params.model) - hash_value(sha, params.provider) - hash_value(sha, params.scheduler.__name__) - hash_value(sha, params.prompt) - hash_value(sha, params.negative_prompt) - hash_value(sha, params.cfg) - hash_value(sha, params.steps) - hash_value(sha, params.seed) - hash_value(sha, size.width) - hash_value(sha, size.height) - - if extras is not None: - for param in extras: - hash_value(sha, param) - - return '%s_%s_%s_%s' % (mode, params.seed, sha.hexdigest(), now) - - -def base_join(base: str, tail: str) -> str: - tail_path = path.relpath(path.normpath(path.join('/', tail)), '/') - return path.join(base, tail_path) - - def run_gc(): logger.debug('running garbage collection') gc.collect() - torch.cuda.empty_cache() \ No newline at end of file + torch.cuda.empty_cache() diff --git a/onnx-web.code-workspace b/onnx-web.code-workspace index 504a5b13..64c0f889 100644 --- a/onnx-web.code-workspace +++ b/onnx-web.code-workspace @@ -60,6 +60,7 @@ "stringcase", "timestep", "timesteps", + "tojson", "uncond", "unet", "untruncated",