refactor txt2img/img2img api

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  

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}')

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}')

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"<str {len(args[idx])}>"
-
-    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)

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"<str {len(args[idx])}>"
+
+        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())

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)