From 02336fdb366a764a8f8e00601f913962ed3ae41a Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Fri, 16 Feb 2024 08:49:42 -0500 Subject: [PATCH] refactor txt2img/img2img api --- CHANGELOG.md | 9 ++- cli/simple-img2img.py | 9 ++- cli/simple-txt2img.py | 8 ++- modules/api/api.py | 129 ++---------------------------------- modules/api/generate.py | 143 ++++++++++++++++++++++++++++++++++++++++ modules/masking.py | 1 - 6 files changed, 165 insertions(+), 134 deletions(-) create mode 100644 modules/api/generate.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 02d931243..8c528e270 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ ## Update for 2024-02-16 -- **improvements**: +- **Improvements**: - **IP Adapter** major refactor - support for **multiple input images** per each ip adapter - support for **multiple concurrent ip adapters** @@ -58,10 +58,13 @@ - add use only for hires pass option - add `--theme` cli param to force theme on startup - add `--allow-paths` cli param to add additional paths that are allowed to be accessed via web, thanks @OuticNZ -- **wiki**: +- **Wiki**: - added benchmark notes for IPEX, OpenVINO and Olive - added ZLUDA wiki page -- **fixes**: +- **Internal** + - update dependencies + - refactor txt2img/img2img api +- **Fixes**: - handle extensions that install conflicting versions of packages `onnxruntime`, `opencv2-python` - installer refresh package cache on any install diff --git a/cli/simple-img2img.py b/cli/simple-img2img.py index 679c788db..7cbfa14b7 100755 --- a/cli/simple-img2img.py +++ b/cli/simple-img2img.py @@ -17,7 +17,6 @@ logging.basicConfig(level = logging.INFO, format = '%(asctime)s %(levelname)s: % log = logging.getLogger(__name__) urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) -filename='/tmp/simple-img2img.jpg' options = { "save_images": False, "send_images": True, @@ -74,8 +73,11 @@ def generate(args): # pylint: disable=redefined-outer-name b64 = data['images'][i].split(',',1)[0] info = data['info'] image = Image.open(io.BytesIO(base64.b64decode(b64))) - image.save(filename) - log.info(f'received image: size={image.size} file={filename} time={t1-t0:.2f} info="{info}"') + log.info(f'received image: size={image.size} time={t1-t0:.2f} info="{info}"') + if args.output: + image.save(args.output) + log.info(f'image saved: size={image.size} filename={args.output}') + else: log.warning(f'no images received: {data}') @@ -89,6 +91,7 @@ if __name__ == "__main__": parser.add_argument('--steps', required=False, default=20, help='number of steps') parser.add_argument('--seed', required=False, default=-1, help='initial seed') parser.add_argument('--sampler', required=False, default='Euler a', help='sampler name') + parser.add_argument('--output', required=False, default=None, help='output image file') parser.add_argument('--model', required=False, help='model name') args = parser.parse_args() log.info(f'img2img: {args}') diff --git a/cli/simple-txt2img.py b/cli/simple-txt2img.py index c2a5ee001..a9696c600 100755 --- a/cli/simple-txt2img.py +++ b/cli/simple-txt2img.py @@ -17,7 +17,6 @@ logging.basicConfig(level = logging.INFO, format = '%(asctime)s %(levelname)s: % log = logging.getLogger(__name__) urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) -filename='/tmp/simple-txt2img.jpg' options = { "save_images": False, "send_images": True, @@ -57,8 +56,10 @@ def generate(args): # pylint: disable=redefined-outer-name b64 = data['images'][i].split(',',1)[0] image = Image.open(io.BytesIO(base64.b64decode(b64))) info = data['info'] - image.save(filename) - log.info(f'received image: size={image.size} file={filename} time={t1-t0:.2f} info="{info}"') + log.info(f'image received: size={image.size} time={t1-t0:.2f} info="{info}"') + if args.output: + image.save(args.output) + log.info(f'image saved: size={image.size} filename={args.output}') else: log.warning(f'no images received: {data}') @@ -72,6 +73,7 @@ if __name__ == "__main__": parser.add_argument('--steps', required=False, default=20, help='number of steps') parser.add_argument('--seed', required=False, default=-1, help='initial seed') parser.add_argument('--sampler', required=False, default='Euler a', help='sampler name') + parser.add_argument('--output', required=False, default=None, help='output image file') parser.add_argument('--model', required=False, help='model name') args = parser.parse_args() log.info(f'txt2img: {args}') diff --git a/modules/api/api.py b/modules/api/api.py index 3f1aedf95..761151a05 100644 --- a/modules/api/api.py +++ b/modules/api/api.py @@ -4,9 +4,8 @@ from secrets import compare_digest from fastapi import FastAPI, APIRouter, Depends, Request from fastapi.security import HTTPBasic, HTTPBasicCredentials from fastapi.exceptions import HTTPException -from modules import errors, shared, scripts, ui, postprocessing -from modules.api import models, endpoints, script, train, helpers, server, nvml -from modules.processing import StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img, process_images +from modules import errors, shared, postprocessing +from modules.api import models, endpoints, script, train, helpers, server, nvml, generate errors.install() @@ -28,6 +27,7 @@ class Api: self.router = APIRouter() self.app = app self.queue_lock = queue_lock + self.generate = generate.APIGenerate(queue_lock) # server api self.add_api_route("/sdapi/v1/motd", server.get_motd, methods=["GET"], response_model=str) @@ -47,8 +47,8 @@ class Api: # core api using locking - self.add_api_route("/sdapi/v1/txt2img", self.post_text2img, methods=["POST"], response_model=models.ResTxt2Img) - self.add_api_route("/sdapi/v1/img2img", self.post_img2img, methods=["POST"], response_model=models.ResImg2Img) + self.add_api_route("/sdapi/v1/txt2img", self.generate.post_text2img, methods=["POST"], response_model=models.ResTxt2Img) + self.add_api_route("/sdapi/v1/img2img", self.generate.post_img2img, methods=["POST"], response_model=models.ResImg2Img) self.add_api_route("/sdapi/v1/extra-single-image", self.extras_single_image_api, methods=["POST"], response_model=models.ResProcessImage) self.add_api_route("/sdapi/v1/extra-batch-images", self.extras_batch_images_api, methods=["POST"], response_model=models.ResProcessBatch) @@ -84,9 +84,6 @@ class Api: self.add_api_route("/sdapi/v1/train/embedding", train.post_train_embedding, methods=["POST"], response_model=models.ResTrain) self.add_api_route("/sdapi/v1/train/hypernetwork", train.post_train_hypernetwork, methods=["POST"], response_model=models.ResTrain) - self.default_script_arg_txt2img = [] - self.default_script_arg_img2img = [] - def add_api_route(self, path: str, endpoint, **kwargs): if (shared.cmd_opts.auth or shared.cmd_opts.auth_file) and shared.cmd_opts.api_only: return self.app.add_api_route(path, endpoint, dependencies=[Depends(self.auth)], **kwargs) @@ -133,122 +130,6 @@ class Api: } del request.ip_adapter - def sanitize_args(self, args: list): - for idx in range(0, len(args)): - if isinstance(args[idx], str) and len(args[idx]) >= 1000: - args[idx] = f"" - - def sanitize_img_gen_request(self, request): - if hasattr(request, "alwayson_scripts") and request.alwayson_scripts: - for script_name in request.alwayson_scripts.keys(): - script_obj = request.alwayson_scripts[script_name] - - if script_obj and "args" in script_obj and script_obj["args"]: - self.sanitize_args(script_obj["args"]) - - if hasattr(request, "script_args") and request.script_args: - self.sanitize_args(request.script_args) - - def post_text2img(self, txt2imgreq: models.ReqTxt2Img): - self.prepare_img_gen_request(txt2imgreq) - - script_runner = scripts.scripts_txt2img - if not script_runner.scripts: - script_runner.initialize_scripts(False) - ui.create_ui(None) - if not self.default_script_arg_txt2img: - self.default_script_arg_txt2img = script.init_default_script_args(script_runner) - selectable_scripts, selectable_script_idx = script.get_selectable_script(txt2imgreq.script_name, script_runner) - populate = txt2imgreq.copy(update={ # Override __init__ params - "sampler_name": helpers.validate_sampler_name(txt2imgreq.sampler_name or txt2imgreq.sampler_index), - "do_not_save_samples": not txt2imgreq.save_images, - "do_not_save_grid": not txt2imgreq.save_images, - }) - if populate.sampler_name: - populate.sampler_index = None # prevent a warning later on - args = vars(populate) - args.pop('script_name', None) - args.pop('script_args', None) # will refeed them to the pipeline directly after initializing them - args.pop('face', None) - args.pop('ip_adapter', None) - args.pop('alwayson_scripts', None) - send_images = args.pop('send_images', True) - args.pop('save_images', None) - - with self.queue_lock: - p = StableDiffusionProcessingTxt2Img(sd_model=shared.sd_model, **args) - p.scripts = script_runner - p.outpath_grids = shared.opts.outdir_grids or shared.opts.outdir_txt2img_grids - p.outpath_samples = shared.opts.outdir_samples or shared.opts.outdir_txt2img_samples - shared.state.begin('api-txt2img', api=True) - script_args = script.init_script_args(p, txt2imgreq, self.default_script_arg_txt2img, selectable_scripts, selectable_script_idx, script_runner) - if selectable_scripts is not None: - processed = scripts.scripts_txt2img.run(p, *script_args) # Need to pass args as list here - else: - p.script_args = tuple(script_args) # Need to pass args as tuple here - processed = process_images(p) - shared.state.end(api=False) - - b64images = list(map(helpers.encode_pil_to_base64, processed.images)) if send_images else [] - self.sanitize_img_gen_request(txt2imgreq) - return models.ResTxt2Img(images=b64images, parameters=vars(txt2imgreq), info=processed.js()) - - def post_img2img(self, img2imgreq: models.ReqImg2Img): - self.prepare_img_gen_request(img2imgreq) - - init_images = img2imgreq.init_images - if init_images is None: - raise HTTPException(status_code=404, detail="Init image not found") - mask = img2imgreq.mask - if mask: - mask = helpers.decode_base64_to_image(mask) - script_runner = scripts.scripts_img2img - if not script_runner.scripts: - script_runner.initialize_scripts(True) - ui.create_ui(None) - if not self.default_script_arg_img2img: - self.default_script_arg_img2img = script.init_default_script_args(script_runner) - selectable_scripts, selectable_script_idx = script.get_selectable_script(img2imgreq.script_name, script_runner) - populate = img2imgreq.copy(update={ # Override __init__ params - "sampler_name": helpers.validate_sampler_name(img2imgreq.sampler_name or img2imgreq.sampler_index), - "do_not_save_samples": not img2imgreq.save_images, - "do_not_save_grid": not img2imgreq.save_images, - "mask": mask, - }) - if populate.sampler_name: - populate.sampler_index = None # prevent a warning later on - args = vars(populate) - args.pop('include_init_images', None) # this is meant to be done by "exclude": True in model, but it's for a reason that I cannot determine. - args.pop('script_name', None) - args.pop('script_args', None) # will refeed them to the pipeline directly after initializing them - args.pop('alwayson_scripts', None) - args.pop('face_id', None) - args.pop('ip_adapter', None) - send_images = args.pop('send_images', True) - args.pop('save_images', None) - - with self.queue_lock: - p = StableDiffusionProcessingImg2Img(sd_model=shared.sd_model, **args) - p.init_images = [helpers.decode_base64_to_image(x) for x in init_images] - p.scripts = script_runner - p.outpath_grids = shared.opts.outdir_img2img_grids - p.outpath_samples = shared.opts.outdir_img2img_samples - shared.state.begin('api-img2img', api=True) - script_args = script.init_script_args(p, img2imgreq, self.default_script_arg_img2img, selectable_scripts, selectable_script_idx, script_runner) - if selectable_scripts is not None: - processed = scripts.scripts_img2img.run(p, *script_args) # Need to pass args as list here - else: - p.script_args = tuple(script_args) # Need to pass args as tuple here - processed = process_images(p) - shared.state.end(api=False) - - b64images = list(map(helpers.encode_pil_to_base64, processed.images)) if send_images else [] - if not img2imgreq.include_init_images: - img2imgreq.init_images = None - img2imgreq.mask = None - self.sanitize_img_gen_request(img2imgreq) - return models.ResImg2Img(images=b64images, parameters=vars(img2imgreq), info=processed.js()) - def set_upscalers(self, req: dict): reqDict = vars(req) reqDict['extras_upscaler_1'] = reqDict.pop('upscaler_1', None) diff --git a/modules/api/generate.py b/modules/api/generate.py new file mode 100644 index 000000000..4f674d716 --- /dev/null +++ b/modules/api/generate.py @@ -0,0 +1,143 @@ +from threading import Lock +from fastapi.exceptions import HTTPException +from modules import errors, shared, scripts, ui +from modules.api import models, script, helpers +from modules.processing import StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img, process_images + + +errors.install() + + +class APIGenerate(): + def __init__(self, queue_lock: Lock): + self.queue_lock = queue_lock + self.default_script_arg_txt2img = [] + self.default_script_arg_img2img = [] + + def sanitize_args(self, args: dict): + args = vars(args) + args.pop('include_init_images', None) # this is meant to be done by "exclude": True in model + args.pop('script_name', None) + args.pop('script_args', None) # will refeed them to the pipeline directly after initializing them + args.pop('alwayson_scripts', None) + args.pop('face', None) + args.pop('face_id', None) + args.pop('ip_adapter', None) + args.pop('save_images', None) + return args + + def sanitize_b64(self, request): + def sanitize_str(args: list): + for idx in range(0, len(args)): + if isinstance(args[idx], str) and len(args[idx]) >= 1000: + args[idx] = f"" + + if hasattr(request, "alwayson_scripts") and request.alwayson_scripts: + for script_name in request.alwayson_scripts.keys(): + script_obj = request.alwayson_scripts[script_name] + if script_obj and "args" in script_obj and script_obj["args"]: + sanitize_str(script_obj["args"]) + if hasattr(request, "script_args") and request.script_args: + sanitize_str(request.script_args) + + def prepare_face_module(self, request): + if hasattr(request, "face") and request.face and not request.script_name and (not request.alwayson_scripts or "face" not in request.alwayson_scripts.keys()): + request.script_name = "face" + request.script_args = [ + request.face.mode, + request.face.source_images, + request.face.ip_model, + request.face.ip_override_sampler, + request.face.ip_cache_model, + request.face.ip_strength, + request.face.ip_structure, + request.face.id_strength, + request.face.id_conditioning, + request.face.id_cache, + request.face.pm_trigger, + request.face.pm_strength, + request.face.pm_start, + request.face.fs_cache + ] + del request.face + + def post_text2img(self, txt2imgreq: models.ReqTxt2Img): + self.prepare_face_module(txt2imgreq) + script_runner = scripts.scripts_txt2img + if not script_runner.scripts: + script_runner.initialize_scripts(False) + ui.create_ui(None) + if not self.default_script_arg_txt2img: + self.default_script_arg_txt2img = script.init_default_script_args(script_runner) + selectable_scripts, selectable_script_idx = script.get_selectable_script(txt2imgreq.script_name, script_runner) + populate = txt2imgreq.copy(update={ # Override __init__ params + "sampler_name": helpers.validate_sampler_name(txt2imgreq.sampler_name or txt2imgreq.sampler_index), + "do_not_save_samples": not txt2imgreq.save_images, + "do_not_save_grid": not txt2imgreq.save_images, + }) + if populate.sampler_name: + populate.sampler_index = None # prevent a warning later on + args = self.sanitize_args(populate) + send_images = args.pop('send_images', True) + with self.queue_lock: + p = StableDiffusionProcessingTxt2Img(sd_model=shared.sd_model, **args) + p.scripts = script_runner + p.outpath_grids = shared.opts.outdir_grids or shared.opts.outdir_txt2img_grids + p.outpath_samples = shared.opts.outdir_samples or shared.opts.outdir_txt2img_samples + shared.state.begin('api-txt2img', api=True) + script_args = script.init_script_args(p, txt2imgreq, self.default_script_arg_txt2img, selectable_scripts, selectable_script_idx, script_runner) + if selectable_scripts is not None: + processed = scripts.scripts_txt2img.run(p, *script_args) # Need to pass args as list here + else: + p.script_args = tuple(script_args) # Need to pass args as tuple here + processed = process_images(p) + shared.state.end(api=False) + b64images = list(map(helpers.encode_pil_to_base64, processed.images)) if send_images else [] + self.sanitize_b64(txt2imgreq) + return models.ResTxt2Img(images=b64images, parameters=vars(txt2imgreq), info=processed.js()) + + def post_img2img(self, img2imgreq: models.ReqImg2Img): + self.prepare_face_module(img2imgreq) + init_images = img2imgreq.init_images + if init_images is None: + raise HTTPException(status_code=404, detail="Init image not found") + mask = img2imgreq.mask + if mask: + mask = helpers.decode_base64_to_image(mask) + script_runner = scripts.scripts_img2img + if not script_runner.scripts: + script_runner.initialize_scripts(True) + ui.create_ui(None) + if not self.default_script_arg_img2img: + self.default_script_arg_img2img = script.init_default_script_args(script_runner) + selectable_scripts, selectable_script_idx = script.get_selectable_script(img2imgreq.script_name, script_runner) + populate = img2imgreq.copy(update={ # Override __init__ params + "sampler_name": helpers.validate_sampler_name(img2imgreq.sampler_name or img2imgreq.sampler_index), + "do_not_save_samples": not img2imgreq.save_images, + "do_not_save_grid": not img2imgreq.save_images, + "mask": mask, + }) + if populate.sampler_name: + populate.sampler_index = None # prevent a warning later on + args = self.sanitize_args(populate) + send_images = args.pop('send_images', True) + with self.queue_lock: + p = StableDiffusionProcessingImg2Img(sd_model=shared.sd_model, **args) + p.init_images = [helpers.decode_base64_to_image(x) for x in init_images] + p.scripts = script_runner + p.outpath_grids = shared.opts.outdir_img2img_grids + p.outpath_samples = shared.opts.outdir_img2img_samples + shared.state.begin('api-img2img', api=True) + script_args = script.init_script_args(p, img2imgreq, self.default_script_arg_img2img, selectable_scripts, selectable_script_idx, script_runner) + if selectable_scripts is not None: + processed = scripts.scripts_img2img.run(p, *script_args) # Need to pass args as list here + else: + p.script_args = tuple(script_args) # Need to pass args as tuple here + processed = process_images(p) + shared.state.end(api=False) + b64images = list(map(helpers.encode_pil_to_base64, processed.images)) if send_images else [] + if not img2imgreq.include_init_images: + img2imgreq.init_images = None + img2imgreq.mask = None + self.sanitize_b64(img2imgreq) + return models.ResImg2Img(images=b64images, parameters=vars(img2imgreq), info=processed.js()) diff --git a/modules/masking.py b/modules/masking.py index ed0fdb7ee..2bf304ba4 100644 --- a/modules/masking.py +++ b/modules/masking.py @@ -354,7 +354,6 @@ def outpaint(input_image: Image.Image, outpaint_type: str = 'Edge'): mask = cv2.erode(mask, kernel, iterations=max(sigmaX, sigmaY) // 3) # increase overlap area mask = cv2.GaussianBlur(mask, (0, 0), sigmaX=sigmaX, sigmaY=sigmaY) # blur mask mask = Image.fromarray(mask) - mask.save('/tmp/mask2.png') if outpaint_type == 'Edge': bordered = cv2.copyMakeBorder(cropped, y1, h0-y2, x1, w0-x2, cv2.BORDER_REPLICATE)