mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 09:14:35 +02:00
refactor txt2img/img2img api
This commit is contained in:
Vladimir Mandic
+5
-124
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user