Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-18 16:54:33 +02:00
Code Issues Packages Projects Releases Wiki Activity

refactor api

Browse Source
This commit is contained in:
Vladimir Mandic
2024-01-28 10:12:24 -05:00
parent dc70246c78
commit 7e447222a1
13 changed files with 778 additions and 785 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CHANGELOG.md
+5 -6
View File
@@ -8,20 +8,16 @@ BLOCKERS:
OPTIONAL:
- pending `diffusers==0.26.0`
- wuerstchen v3 [pr](https://github.com/huggingface/diffusers/pull/6487)
- animatediff image2video [pr](https://github.com/huggingface/diffusers/pull/6509)
- tiledvae [pr](https://github.com/huggingface/diffusers/pull/1441)
- style aligned [pr](https://github.com/huggingface/diffusers/pull/6489)
- mixture tiling [pr](https://github.com/huggingface/diffusers/tree/main/examples/community#stable-diffusion-mixture-tiling)
- depth anything [repo](https://depth-anything.github.io/)
- instaflow [pr](https://github.com/huggingface/diffusers/pull/6057)[repo](https://github.com/gnobitab/RectifiedFlow)
- control api
- photomaker api
- interrogate api
- remb api
- face api
- masking api
- preprocess api
## Update for 2023-01-27
## Update for 2023-01-28
Another big release, highlights being:
- A lot more functionality in the **Control** module:
@@ -277,6 +273,9 @@ As of this release, default backend is set to **diffusers** as its more feature
- major internal ui module refactoring
this may cause compatibility issues if an extension is doing a direct import from `ui.py`
in which case, report it so we can add a compatibility layer
- major public api refactoring
this may cause compatibility issues if an extension is doing a direct import from `api.py` or `models.py`
in which case, report it so we can add a compatibility layer
## Update for 2023-12-29
installer.py
+1 -1
View File
@@ -48,7 +48,7 @@ args = Dot({
git_commit = "unknown"
submodules_commit = {
'sd-webui-controlnet': 'ecd33eb',
'stable-diffusion-webui-images-browser': '27fe4a7',
# 'stable-diffusion-webui-images-browser': '27fe4a7',
}
# setup console and file logging
modules/api/api.py
+102 -552
View File
@@ -1,101 +1,17 @@
import io
import time
import base64
from io import BytesIO
from typing import List, Dict, Any, Optional
from typing import List, Optional
from threading import Lock
from secrets import compare_digest
from fastapi import FastAPI, APIRouter, Depends, Request
from fastapi.security import HTTPBasic, HTTPBasicCredentials
from fastapi.exceptions import HTTPException
from PIL import PngImagePlugin,Image
import requests
import piexif
import piexif.helper
import gradio as gr
from modules import errors, shared, sd_samplers, deepbooru, sd_hijack, images, scripts, ui, postprocessing, script_callbacks, generation_parameters_copypaste
from modules.sd_vae import vae_dict
from modules.api import models
from modules import errors, shared, sd_samplers, scripts, ui, postprocessing
from modules.api import models, endpoints, script, train, helpers, server
from modules.processing import StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img, process_images
from modules.textual_inversion.textual_inversion import create_embedding, train_embedding
from modules.textual_inversion.preprocess import preprocess
from modules.hypernetworks.hypernetwork import create_hypernetwork, train_hypernetwork
from modules.sd_models import checkpoints_list, unload_model_weights, reload_model_weights
from modules.sd_models_config import find_checkpoint_config_near_filename
from modules import devices
errors.install()
def upscaler_to_index(name: str):
try:
return [x.name.lower() for x in shared.sd_upscalers].index(name.lower())
except Exception as e:
raise HTTPException(status_code=400, detail=f"Invalid upscaler, needs to be one of these: {' , '.join([x.name for x in shared.sd_upscalers])}") from e
def script_name_to_index(name, scripts_list):
try:
return [script.title().lower() for script in scripts_list].index(name.lower())
except Exception as e:
raise HTTPException(status_code=422, detail=f"Script '{name}' not found") from e
def validate_sampler_name(name):
config = sd_samplers.all_samplers_map.get(name, None)
if config is None:
raise HTTPException(status_code=404, detail="Sampler not found")
return name
def setUpscalers(req: dict):
reqDict = vars(req)
reqDict['extras_upscaler_1'] = reqDict.pop('upscaler_1', None)
reqDict['extras_upscaler_2'] = reqDict.pop('upscaler_2', None)
return reqDict
def decode_base64_to_image(encoding):
if encoding.startswith("data:image/"):
encoding = encoding.split(";")[1].split(",")[1]
try:
image = Image.open(BytesIO(base64.b64decode(encoding)))
return image
except Exception as e:
shared.log.warning(f'API cannot decode image: {e}')
raise HTTPException(status_code=500, detail="Invalid encoded image") from e
def save_image(image, fn, ext):
# actual save
parameters = image.info.get('parameters', None)
image_format = Image.registered_extensions()[f'.{ext}']
if image_format == 'PNG':
pnginfo_data = PngImagePlugin.PngInfo()
for k, v in image.info.items():
pnginfo_data.add_text(k, str(v))
image.save(fn, format=image_format, quality=shared.opts.jpeg_quality, pnginfo=pnginfo_data)
elif image_format == 'JPEG':
if image.mode == 'RGBA':
shared.log.warning('Saving RGBA image as JPEG: Alpha channel will be lost')
image = image.convert("RGB")
elif image.mode == 'I;16':
image = image.point(lambda p: p * 0.0038910505836576).convert("L")
exif_bytes = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(parameters or "", encoding="unicode") } })
image.save(fn, format=image_format, quality=shared.opts.jpeg_quality, exif=exif_bytes)
elif image_format == 'WEBP':
if image.mode == 'I;16':
image = image.point(lambda p: p * 0.0038910505836576).convert("RGB")
exif_bytes = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(parameters or "", encoding="unicode") } })
image.save(fn, format=image_format, quality=shared.opts.jpeg_quality, lossless=shared.opts.webp_lossless, exif=exif_bytes)
else:
# shared.log.warning(f'Unrecognized image format: {extension} attempting save as {image_format}')
image.save(fn, format=image_format, quality=shared.opts.jpeg_quality)
def encode_pil_to_base64(image):
with io.BytesIO() as output_bytes:
save_image(image, output_bytes, shared.opts.samples_format)
bytes_data = output_bytes.getvalue()
return base64.b64encode(bytes_data)
class Api:
def __init__(self, app: FastAPI, queue_lock: Lock):
self.credentials = {}
@@ -112,44 +28,58 @@ class Api:
self.router = APIRouter()
self.app = app
self.queue_lock = queue_lock
self.add_api_route("/sdapi/v1/txt2img", self.text2imgapi, methods=["POST"], response_model=models.TextToImageResponse)
self.add_api_route("/sdapi/v1/img2img", self.img2imgapi, methods=["POST"], response_model=models.ImageToImageResponse)
self.add_api_route("/sdapi/v1/extra-single-image", self.extras_single_image_api, methods=["POST"], response_model=models.ExtrasSingleImageResponse)
self.add_api_route("/sdapi/v1/extra-batch-images", self.extras_batch_images_api, methods=["POST"], response_model=models.ExtrasBatchImagesResponse)
self.add_api_route("/sdapi/v1/png-info", self.pnginfoapi, methods=["POST"], response_model=models.PNGInfoResponse)
self.add_api_route("/sdapi/v1/progress", self.progressapi, methods=["GET"], response_model=models.ProgressResponse)
self.add_api_route("/sdapi/v1/interrogate", self.interrogateapi, methods=["POST"])
self.add_api_route("/sdapi/v1/interrupt", self.interruptapi, methods=["POST"])
self.add_api_route("/sdapi/v1/skip", self.skip, methods=["POST"])
self.add_api_route("/sdapi/v1/options", self.get_config, methods=["GET"], response_model=models.OptionsModel)
self.add_api_route("/sdapi/v1/options", self.set_config, methods=["POST"])
self.add_api_route("/sdapi/v1/cmd-flags", self.get_cmd_flags, methods=["GET"], response_model=models.FlagsModel)
self.add_api_route("/sdapi/v1/samplers", self.get_samplers, methods=["GET"], response_model=List[models.SamplerItem])
self.add_api_route("/sdapi/v1/upscalers", self.get_upscalers, methods=["GET"], response_model=List[models.UpscalerItem])
self.add_api_route("/sdapi/v1/sd-models", self.get_sd_models, methods=["GET"], response_model=List[models.SDModelItem])
self.add_api_route("/sdapi/v1/hypernetworks", self.get_hypernetworks, methods=["GET"], response_model=List[models.HypernetworkItem])
self.add_api_route("/sdapi/v1/face-restorers", self.get_face_restorers, methods=["GET"], response_model=List[models.FaceRestorerItem])
self.add_api_route("/sdapi/v1/prompt-styles", self.get_prompt_styles, methods=["GET"], response_model=List[models.StyleItem])
self.add_api_route("/sdapi/v1/embeddings", self.get_embeddings, methods=["GET"], response_model=models.EmbeddingsResponse)
self.add_api_route("/sdapi/v1/refresh-checkpoints", self.refresh_checkpoints, methods=["POST"])
self.add_api_route("/sdapi/v1/sd-vae", self.get_sd_vaes, methods=["GET"], response_model=List[models.SDVaeItem])
self.add_api_route("/sdapi/v1/refresh-vae", self.refresh_vaes, methods=["POST"])
self.add_api_route("/sdapi/v1/create/embedding", self.create_embedding, methods=["POST"], response_model=models.CreateResponse)
self.add_api_route("/sdapi/v1/create/hypernetwork", self.create_hypernetwork, methods=["POST"], response_model=models.CreateResponse)
self.add_api_route("/sdapi/v1/preprocess", self.preprocess, methods=["POST"], response_model=models.PreprocessResponse)
self.add_api_route("/sdapi/v1/train/embedding", self.train_embedding, methods=["POST"], response_model=models.TrainResponse)
self.add_api_route("/sdapi/v1/train/hypernetwork", self.train_hypernetwork, methods=["POST"], response_model=models.TrainResponse)
self.add_api_route("/sdapi/v1/shutdown", self.shutdown, methods=["POST"])
self.add_api_route("/sdapi/v1/memory", self.get_memory, methods=["GET"], response_model=models.MemoryResponse)
self.add_api_route("/sdapi/v1/unload-checkpoint", self.unloadapi, methods=["POST"])
self.add_api_route("/sdapi/v1/reload-checkpoint", self.reloadapi, methods=["POST"])
self.add_api_route("/sdapi/v1/scripts", self.get_scripts_list, methods=["GET"], response_model=models.ScriptsList)
self.add_api_route("/sdapi/v1/script-info", self.get_script_info, methods=["GET"], response_model=List[models.ScriptInfo])
self.add_api_route("/sdapi/v1/extensions", self.get_extensions_list, methods=["GET"], response_model=List[models.ExtensionItem])
self.add_api_route("/sdapi/v1/log", self.get_log_buffer, methods=["GET"], response_model=List)
self.add_api_route("/sdapi/v1/start", self.session_start, methods=["GET"])
self.add_api_route("/sdapi/v1/motd", self.get_motd, methods=["GET"], response_model=str)
self.add_api_route("/sdapi/v1/extra-networks", self.get_extra_networks, methods=["GET"], response_model=List[models.ExtraNetworkItem])
# server api
self.add_api_route("/sdapi/v1/motd", server.get_motd, methods=["GET"], response_model=str)
self.add_api_route("/sdapi/v1/log", server.get_log_buffer, methods=["GET"], response_model=List[str])
self.add_api_route("/sdapi/v1/start", self.get_session_start, methods=["GET"])
self.add_api_route("/sdapi/v1/progress", server.get_progress, methods=["GET"], response_model=models.ResProgress)
self.add_api_route("/sdapi/v1/interrupt", server.post_interrupt, methods=["POST"])
self.add_api_route("/sdapi/v1/skip", server.post_skip, methods=["POST"])
self.add_api_route("/sdapi/v1/shutdown", server.post_shutdown, methods=["POST"])
self.add_api_route("/sdapi/v1/memory", server.get_memory, methods=["GET"], response_model=models.ResMemory)
self.add_api_route("/sdapi/v1/options", server.get_config, methods=["GET"], response_model=models.OptionsModel)
self.add_api_route("/sdapi/v1/options", server.set_config, methods=["POST"])
self.add_api_route("/sdapi/v1/cmd-flags", server.get_cmd_flags, methods=["GET"], response_model=models.FlagsModel)
# 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/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)
# api dealing with optional scripts
self.add_api_route("/sdapi/v1/scripts", script.get_scripts_list, methods=["GET"], response_model=models.ResScripts)
self.add_api_route("/sdapi/v1/script-info", script.get_script_info, methods=["GET"], response_model=List[models.ItemScript])
# enumerator api
self.add_api_route("/sdapi/v1/interrogate", endpoints.get_interrogate, methods=["GET"], response_model=List[str])
self.add_api_route("/sdapi/v1/samplers", endpoints.get_samplers, methods=["GET"], response_model=List[models.ItemSampler])
self.add_api_route("/sdapi/v1/upscalers", endpoints.get_upscalers, methods=["GET"], response_model=List[models.ItemUpscaler])
self.add_api_route("/sdapi/v1/sd-models", endpoints.get_sd_models, methods=["GET"], response_model=List[models.ItemModel])
self.add_api_route("/sdapi/v1/hypernetworks", endpoints.get_hypernetworks, methods=["GET"], response_model=List[models.ItemHypernetwork])
self.add_api_route("/sdapi/v1/face-restorers", endpoints.get_face_restorers, methods=["GET"], response_model=List[models.ItemFaceRestorer])
self.add_api_route("/sdapi/v1/prompt-styles", endpoints.get_prompt_styles, methods=["GET"], response_model=List[models.ItemStyle])
self.add_api_route("/sdapi/v1/embeddings", endpoints.get_embeddings, methods=["GET"], response_model=models.ResEmbeddings)
self.add_api_route("/sdapi/v1/sd-vae", endpoints.get_sd_vaes, methods=["GET"], response_model=List[models.ItemVae])
self.add_api_route("/sdapi/v1/extensions", endpoints.get_extensions_list, methods=["GET"], response_model=List[models.ItemExtension])
self.add_api_route("/sdapi/v1/extra-networks", endpoints.get_extra_networks, methods=["GET"], response_model=List[models.ItemExtraNetwork])
# functional api
self.add_api_route("/sdapi/v1/png-info", endpoints.post_pnginfo, methods=["POST"], response_model=models.ResImageInfo)
self.add_api_route("/sdapi/v1/interrogate", endpoints.post_interrogate, methods=["POST"])
self.add_api_route("/sdapi/v1/refresh-checkpoints", endpoints.post_refresh_checkpoints, methods=["POST"])
self.add_api_route("/sdapi/v1/unload-checkpoint", endpoints.post_unload_checkpoint, methods=["POST"])
self.add_api_route("/sdapi/v1/reload-checkpoint", endpoints.post_reload_checkpoint, methods=["POST"])
self.add_api_route("/sdapi/v1/refresh-vae", endpoints.post_refresh_vae, methods=["POST"])
# train api
self.add_api_route("/sdapi/v1/create/embedding", train.post_create_embedding, methods=["POST"], response_model=models.ResCreate)
self.add_api_route("/sdapi/v1/create/hypernetwork", train.post_create_hypernetwork, methods=["POST"], response_model=models.ResCreate)
self.add_api_route("/sdapi/v1/preprocess", train.post_preprocess, methods=["POST"], response_model=models.ResPreprocess)
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 = []
@@ -165,103 +95,13 @@ class Api:
return True
raise HTTPException(status_code=401, detail="Unauthorized", headers={"WWW-Authenticate": "Basic"})
def get_log_buffer(self, req: models.LogRequest = Depends()):
lines = shared.log.buffer[:req.lines] if req.lines > 0 else shared.log.buffer.copy()
if req.clear:
shared.log.buffer.clear()
return lines
def session_start(self, req: Request, agent: Optional[str] = None):
def get_session_start(self, req: Request, agent: Optional[str] = None):
token = req.cookies.get("access-token") or req.cookies.get("access-token-unsecure")
user = self.app.tokens.get(token) if hasattr(self.app, 'tokens') else None
shared.log.info(f'Browser session: user={user} client={req.client.host} agent={agent}')
return {}
def get_motd(self):
from installer import get_version
motd = ''
ver = get_version()
if ver.get('updated', None) is not None:
motd = f"version <b>{ver['hash']} {ver['updated']}</b> <span style='color: var(--primary-500)'>{ver['url'].split('/')[-1]}</span><br>"
if shared.opts.motd:
res = requests.get('https://vladmandic.github.io/automatic/motd', timeout=10)
if res.status_code == 200:
msg = (res.text or '').strip()
shared.log.info(f'MOTD: {msg if len(msg) > 0 else "N/A"}')
motd += res.text
return motd
def get_selectable_script(self, script_name, script_runner):
if script_name is None or script_name == "":
return None, None
script_idx = script_name_to_index(script_name, script_runner.selectable_scripts)
script = script_runner.selectable_scripts[script_idx]
return script, script_idx
def get_scripts_list(self):
t2ilist = [script.name for script in scripts.scripts_txt2img.scripts if script.name is not None]
i2ilist = [script.name for script in scripts.scripts_img2img.scripts if script.name is not None]
control = [script.name for script in scripts.scripts_control.scripts if script.name is not None]
return models.ScriptsList(txt2img = t2ilist, img2img = i2ilist, control = control)
def get_script_info(self, script_name: Optional[str] = None):
res = []
for script_list in [scripts.scripts_txt2img.scripts, scripts.scripts_img2img.scripts, scripts.scripts_control.scripts]:
for script in script_list:
if script.api_info is not None and (script_name is None or script_name == script.api_info.name):
res.append(script.api_info)
return res
def get_script(self, script_name, script_runner):
if script_name is None or script_name == "":
return None, None
script_idx = script_name_to_index(script_name, script_runner.scripts)
return script_runner.scripts[script_idx]
def init_default_script_args(self, script_runner):
#find max idx from the scripts in runner and generate a none array to init script_args
last_arg_index = 1
for script in script_runner.scripts:
if last_arg_index < script.args_to:
last_arg_index = script.args_to
# None everywhere except position 0 to initialize script args
script_args = [None]*last_arg_index
script_args[0] = 0
# get default values
if gr is None:
return script_args
with gr.Blocks(): # will throw errors calling ui function without this
for script in script_runner.scripts:
if script.ui(script.is_img2img):
ui_default_values = []
for elem in script.ui(script.is_img2img):
ui_default_values.append(elem.value)
script_args[script.args_from:script.args_to] = ui_default_values
return script_args
def init_script_args(self, p, request, default_script_args, selectable_scripts, selectable_script_idx, script_runner):
script_args = default_script_args.copy()
# position 0 in script_arg is the idx+1 of the selectable script that is going to be run when using scripts.scripts_*2img.run()
if selectable_scripts:
script_args[selectable_scripts.args_from:selectable_scripts.args_to] = request.script_args
script_args[0] = selectable_script_idx + 1
# Now check for always on scripts
if request.alwayson_scripts and (len(request.alwayson_scripts) > 0):
for alwayson_script_name in request.alwayson_scripts.keys():
alwayson_script = self.get_script(alwayson_script_name, script_runner)
if alwayson_script is None:
raise HTTPException(status_code=422, detail=f"Always on script not found: {alwayson_script_name}")
if not alwayson_script.alwayson:
raise HTTPException(status_code=422, detail=f"Selectable script cannot be in always on params: {alwayson_script_name}")
if "args" in request.alwayson_scripts[alwayson_script_name]:
# min between arg length in scriptrunner and arg length in the request
for idx in range(0, min((alwayson_script.args_to - alwayson_script.args_from), len(request.alwayson_scripts[alwayson_script_name]["args"]))):
script_args[alwayson_script.args_from + idx] = request.alwayson_scripts[alwayson_script_name]["args"][idx]
p.per_script_args[alwayson_script.title()] = request.alwayson_scripts[alwayson_script_name]["args"]
return script_args
def prepare_img_gen_request(self, request, img_gen_type: str): # pylint: disable=unused-argument
def prepare_img_gen_request(self, request):
if hasattr(request, "face_id") and request.face_id and not request.script_name and (not request.alwayson_scripts or "FaceID" not in request.alwayson_scripts.keys()):
request.script_name = "FaceID"
request.script_args = [
@@ -289,7 +129,7 @@ class Api:
if isinstance(args[idx], str) and len(args[idx]) >= 1000:
args[idx] = f"<str {len(args[idx])}>"
def sanitize_img_gen_request(self, request, img_gen_type: str):
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]
@@ -300,18 +140,24 @@ class Api:
if hasattr(request, "script_args") and request.script_args:
self.sanitize_args(request.script_args)
def text2imgapi(self, txt2imgreq: models.StableDiffusionTxt2ImgProcessingAPI):
self.prepare_img_gen_request(txt2imgreq, "txt2img")
def validate_sampler_name(self, name):
config = sd_samplers.all_samplers_map.get(name, None)
if config is None:
raise HTTPException(status_code=404, detail="Sampler not found")
return name
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 = self.init_default_script_args(script_runner)
selectable_scripts, selectable_script_idx = self.get_selectable_script(txt2imgreq.script_name, script_runner)
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": validate_sampler_name(txt2imgreq.sampler_name or txt2imgreq.sampler_index),
"sampler_name": self.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,
})
@@ -332,7 +178,7 @@ class Api:
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 = self.init_script_args(p, txt2imgreq, self.default_script_arg_txt2img, selectable_scripts, selectable_script_idx, script_runner)
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:
@@ -340,28 +186,28 @@ class Api:
processed = process_images(p)
shared.state.end(api=False)
b64images = list(map(encode_pil_to_base64, processed.images)) if send_images else []
self.sanitize_img_gen_request(txt2imgreq, "txt2img")
return models.TextToImageResponse(images=b64images, parameters=vars(txt2imgreq), info=processed.js())
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 img2imgapi(self, img2imgreq: models.StableDiffusionImg2ImgProcessingAPI):
self.prepare_img_gen_request(img2imgreq, "img2img")
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 = decode_base64_to_image(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 = self.init_default_script_args(script_runner)
selectable_scripts, selectable_script_idx = self.get_selectable_script(img2imgreq.script_name, script_runner)
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": validate_sampler_name(img2imgreq.sampler_name or img2imgreq.sampler_index),
"sampler_name": self.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,
@@ -380,12 +226,12 @@ class Api:
with self.queue_lock:
p = StableDiffusionProcessingImg2Img(sd_model=shared.sd_model, **args)
p.init_images = [decode_base64_to_image(x) for x in init_images]
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 = self.init_script_args(p, img2imgreq, self.default_script_arg_img2img, selectable_scripts, selectable_script_idx, script_runner)
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:
@@ -393,329 +239,33 @@ class Api:
processed = process_images(p)
shared.state.end(api=False)
b64images = list(map(encode_pil_to_base64, processed.images)) if send_images else []
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, "img2img")
return models.ImageToImageResponse(images=b64images, parameters=vars(img2imgreq), info=processed.js())
self.sanitize_img_gen_request(img2imgreq)
return models.ResImg2Img(images=b64images, parameters=vars(img2imgreq), info=processed.js())
def extras_single_image_api(self, req: models.ExtrasSingleImageRequest):
reqDict = setUpscalers(req)
reqDict['image'] = decode_base64_to_image(reqDict['image'])
def set_upscalers(self, req: dict):
reqDict = vars(req)
reqDict['extras_upscaler_1'] = reqDict.pop('upscaler_1', None)
reqDict['extras_upscaler_2'] = reqDict.pop('upscaler_2', None)
return reqDict
def extras_single_image_api(self, req: models.ReqProcessImage):
reqDict = self.set_upscalers(req)
reqDict['image'] = helpers.decode_base64_to_image(reqDict['image'])
with self.queue_lock:
result = postprocessing.run_extras(extras_mode=0, image_folder="", input_dir="", output_dir="", save_output=False, **reqDict)
return models.ExtrasSingleImageResponse(image=encode_pil_to_base64(result[0][0]), html_info=result[1])
return models.ResProcessImage(image=helpers.encode_pil_to_base64(result[0][0]), html_info=result[1])
def extras_batch_images_api(self, req: models.ExtrasBatchImagesRequest):
reqDict = setUpscalers(req)
def extras_batch_images_api(self, req: models.ReqProcessBatch):
reqDict = self.set_upscalers(req)
image_list = reqDict.pop('imageList', [])
image_folder = [decode_base64_to_image(x.data) for x in image_list]
image_folder = [helpers.decode_base64_to_image(x.data) for x in image_list]
with self.queue_lock:
result = postprocessing.run_extras(extras_mode=1, image_folder=image_folder, image="", input_dir="", output_dir="", save_output=False, **reqDict)
return models.ExtrasBatchImagesResponse(images=list(map(encode_pil_to_base64, result[0])), html_info=result[1])
def pnginfoapi(self, req: models.PNGInfoRequest):
if not req.image.strip():
return models.PNGInfoResponse(info="")
image = decode_base64_to_image(req.image.strip())
if image is None:
return models.PNGInfoResponse(info="")
geninfo, items = images.read_info_from_image(image)
if geninfo is None:
geninfo = ""
if items and items['parameters']:
del items['parameters']
params = generation_parameters_copypaste.parse_generation_parameters(geninfo)
script_callbacks.infotext_pasted_callback(geninfo, params)
return models.PNGInfoResponse(info=geninfo, items=items, parameters=params)
def progressapi(self, req: models.ProgressRequest = Depends()):
if shared.state.job_count == 0:
return models.ProgressResponse(progress=0, eta_relative=0, state=shared.state.dict(), textinfo=shared.state.textinfo)
shared.state.do_set_current_image()
current_image = None
if shared.state.current_image and not req.skip_current_image:
current_image = encode_pil_to_base64(shared.state.current_image)
batch_x = max(shared.state.job_no, 0)
batch_y = max(shared.state.job_count, 1)
step_x = max(shared.state.sampling_step, 0)
step_y = max(shared.state.sampling_steps, 1)
current = step_y * batch_x + step_x
total = step_y * batch_y
progress = current / total if current > 0 and total > 0 else 0
time_since_start = time.time() - shared.state.time_start
eta_relative = (time_since_start / progress) - time_since_start if progress > 0 else 0
res = models.ProgressResponse(progress=progress, eta_relative=eta_relative, state=shared.state.dict(), current_image=current_image, textinfo=shared.state.textinfo)
return res
def interrogateapi(self, interrogatereq: models.InterrogateRequest):
image_b64 = interrogatereq.image
if image_b64 is None:
raise HTTPException(status_code=404, detail="Image not found")
img = decode_base64_to_image(image_b64)
img = img.convert('RGB')
with self.queue_lock:
if interrogatereq.model == "clip":
processed = shared.interrogator.interrogate(img)
elif interrogatereq.model == "deepdanbooru":
processed = deepbooru.model.tag(img)
else:
raise HTTPException(status_code=404, detail="Model not found")
return models.InterrogateResponse(caption=processed)
def interruptapi(self):
shared.state.interrupt()
return {}
def unloadapi(self):
unload_model_weights(op='model')
unload_model_weights(op='refiner')
return {}
def reloadapi(self):
reload_model_weights()
return {}
def skip(self):
shared.state.skip()
def get_config(self):
options = {}
for k in shared.opts.data.keys():
if shared.opts.data_labels.get(k) is not None:
options.update({k: shared.opts.data.get(k, shared.opts.data_labels.get(k).default)})
else:
options.update({k: shared.opts.data.get(k, None)})
if 'sd_lyco' in options:
del options['sd_lyco']
if 'sd_lora' in options:
del options['sd_lora']
return options
def set_config(self, req: Dict[str, Any]):
updated = []
for k, v in req.items():
updated.append({ k: shared.opts.set(k, v) })
shared.opts.save(shared.config_filename)
return { "updated": updated }
def get_cmd_flags(self):
return vars(shared.cmd_opts)
def get_samplers(self):
return [{"name": sampler[0], "aliases":sampler[2], "options":sampler[3]} for sampler in sd_samplers.all_samplers]
def get_sd_vaes(self):
return [{"model_name": x, "filename": vae_dict[x]} for x in vae_dict.keys()]
def get_upscalers(self):
return [{"name": upscaler.name, "model_name": upscaler.scaler.model_name, "model_path": upscaler.data_path, "model_url": None, "scale": upscaler.scale} for upscaler in shared.sd_upscalers]
def get_sd_models(self):
return [{"title": x.title, "model_name": x.name, "filename": x.filename, "type": x.type, "hash": x.shorthash, "sha256": x.sha256, "config": find_checkpoint_config_near_filename(x)} for x in checkpoints_list.values()]
def get_hypernetworks(self):
return [{"name": name, "path": shared.hypernetworks[name]} for name in shared.hypernetworks]
def get_face_restorers(self):
return [{"name":x.name(), "cmd_dir": getattr(x, "cmd_dir", None)} for x in shared.face_restorers]
def get_prompt_styles(self):
return [{ 'name': v.name, 'prompt': v.prompt, 'negative_prompt': v.negative_prompt, 'extra': v.extra, 'filename': v.filename, 'preview': v.preview} for v in shared.prompt_styles.styles.values()]
def get_embeddings(self):
db = sd_hijack.model_hijack.embedding_db
def convert_embedding(embedding):
return {"step": embedding.step, "sd_checkpoint": embedding.sd_checkpoint, "sd_checkpoint_name": embedding.sd_checkpoint_name, "shape": embedding.shape, "vectors": embedding.vectors}
def convert_embeddings(embeddings):
return {embedding.name: convert_embedding(embedding) for embedding in embeddings.values()}
return {"loaded": convert_embeddings(db.word_embeddings), "skipped": convert_embeddings(db.skipped_embeddings)}
def get_extra_networks(self, page: Optional[str] = None, name: Optional[str] = None, filename: Optional[str] = None, title: Optional[str] = None, fullname: Optional[str] = None, hash: Optional[str] = None): # pylint: disable=redefined-builtin
res = []
for pg in shared.extra_networks:
if page is not None and pg.name != page.lower():
continue
for item in pg.items:
if name is not None and item.get('name', '') != name:
continue
if title is not None and item.get('title', '') != title:
continue
if filename is not None and item.get('filename', '') != filename:
continue
if fullname is not None and item.get('fullname', '') != fullname:
continue
if hash is not None and (item.get('shorthash', None) or item.get('hash')) != hash:
continue
res.append({
'name': item.get('name', ''),
'type': pg.name,
'title': item.get('title', None),
'fullname': item.get('fullname', None),
'filename': item.get('filename', None),
'hash': item.get('shorthash', None) or item.get('hash'),
"preview": item.get('preview', None),
})
return res
def refresh_checkpoints(self):
return shared.refresh_checkpoints()
def refresh_vaes(self):
return shared.refresh_vaes()
def create_embedding(self, args: dict):
try:
shared.state.begin('api-embedding')
filename = create_embedding(**args) # create empty embedding
sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings() # reload embeddings so new one can be immediately used
shared.state.end()
return models.CreateResponse(info = f"create embedding filename: {filename}")
except AssertionError as e:
shared.state.end()
return models.TrainResponse(info = f"create embedding error: {e}")
def create_hypernetwork(self, args: dict):
try:
shared.state.begin('api-hypernetwork')
filename = create_hypernetwork(**args) # create empty embedding # pylint: disable=E1111
shared.state.end()
return models.CreateResponse(info = f"create hypernetwork filename: {filename}")
except AssertionError as e:
shared.state.end()
return models.TrainResponse(info = f"create hypernetwork error: {e}")
def preprocess(self, args: dict):
try:
shared.state.begin('api-preprocess')
preprocess(**args) # quick operation unless blip/booru interrogation is enabled
shared.state.end()
return models.PreprocessResponse(info = 'preprocess complete')
except KeyError as e:
shared.state.end()
return models.PreprocessResponse(info = f"preprocess error: invalid token: {e}")
except AssertionError as e:
shared.state.end()
return models.PreprocessResponse(info = f"preprocess error: {e}")
except FileNotFoundError as e:
shared.state.end()
return models.PreprocessResponse(info = f'preprocess error: {e}')
def train_embedding(self, args: dict):
try:
shared.state.begin('api-embedding')
apply_optimizations = False
error = None
filename = ''
if not apply_optimizations:
sd_hijack.undo_optimizations()
try:
_embedding, filename = train_embedding(**args) # can take a long time to complete
except Exception as e:
error = e
finally:
if not apply_optimizations:
sd_hijack.apply_optimizations()
shared.state.end()
return models.TrainResponse(info = f"train embedding complete: filename: {filename} error: {error}")
except AssertionError as msg:
shared.state.end()
return models.TrainResponse(info = f"train embedding error: {msg}")
def train_hypernetwork(self, args: dict):
try:
shared.state.begin('api-hypernetwork')
shared.loaded_hypernetworks = []
apply_optimizations = False
error = None
filename = ''
if not apply_optimizations:
sd_hijack.undo_optimizations()
try:
_hypernetwork, filename = train_hypernetwork(**args)
except Exception as e:
error = e
finally:
shared.sd_model.cond_stage_model.to(devices.device)
shared.sd_model.first_stage_model.to(devices.device)
if not apply_optimizations:
sd_hijack.apply_optimizations()
shared.state.end()
return models.TrainResponse(info=f"train embedding complete: filename: {filename} error: {error}")
except AssertionError:
shared.state.end()
return models.TrainResponse(info=f"train embedding error: {error}")
def shutdown(self):
shared.log.info('Shutdown request received')
import sys
sys.exit(0)
def get_memory(self):
try:
import os
import psutil
process = psutil.Process(os.getpid())
res = process.memory_info() # only rss is cross-platform guaranteed so we dont rely on other values
ram_total = 100 * res.rss / process.memory_percent() # and total memory is calculated as actual value is not cross-platform safe
ram = { 'free': ram_total - res.rss, 'used': res.rss, 'total': ram_total }
except Exception as err:
ram = { 'error': f'{err}' }
try:
import torch
if torch.cuda.is_available():
s = torch.cuda.mem_get_info()
system = { 'free': s[0], 'used': s[1] - s[0], 'total': s[1] }
s = dict(torch.cuda.memory_stats(shared.device))
allocated = { 'current': s['allocated_bytes.all.current'], 'peak': s['allocated_bytes.all.peak'] }
reserved = { 'current': s['reserved_bytes.all.current'], 'peak': s['reserved_bytes.all.peak'] }
active = { 'current': s['active_bytes.all.current'], 'peak': s['active_bytes.all.peak'] }
inactive = { 'current': s['inactive_split_bytes.all.current'], 'peak': s['inactive_split_bytes.all.peak'] }
warnings = { 'retries': s['num_alloc_retries'], 'oom': s['num_ooms'] }
cuda = {
'system': system,
'active': active,
'allocated': allocated,
'reserved': reserved,
'inactive': inactive,
'events': warnings,
}
else:
cuda = { 'error': 'unavailable' }
except Exception as err:
cuda = { 'error': f'{err}' }
return models.MemoryResponse(ram = ram, cuda = cuda)
def get_extensions_list(self):
from modules import extensions
extensions.list_extensions()
ext_list = []
for ext in extensions.extensions:
ext: extensions.Extension
ext.read_info()
if ext.remote is not None:
ext_list.append({
"name": ext.name,
"remote": ext.remote,
"branch": ext.branch,
"commit_hash":ext.commit_hash,
"commit_date":ext.commit_date,
"version":ext.version,
"enabled":ext.enabled
})
return ext_list
return models.ResProcessBatch(images=list(map(helpers.encode_pil_to_base64, result[0])), html_info=result[1])
def launch(self):
config = {
@@ -727,9 +277,9 @@ class Api:
"http": "auto", # auto, h11, httptools
}
from modules.server import UvicornServer
server = UvicornServer(self.app, **config)
http_server = UvicornServer(self.app, **config)
# from modules.server import HypercornServer
# server = HypercornServer(self.app, **config)
server.start()
http_server.start()
shared.log.info(f'API server: Uvicorn options={config}')
return server
return http_server
modules/api/endpoints.py
+147
View File
@@ -0,0 +1,147 @@
from typing import Optional
from fastapi.exceptions import HTTPException
from modules import shared
from modules.api import models, helpers
def get_samplers():
from modules import sd_samplers
return [{"name": sampler[0], "aliases":sampler[2], "options":sampler[3]} for sampler in sd_samplers.all_samplers]
def get_sd_vaes():
from modules.sd_vae import vae_dict
return [{"model_name": x, "filename": vae_dict[x]} for x in vae_dict.keys()]
def get_upscalers():
return [{"name": upscaler.name, "model_name": upscaler.scaler.model_name, "model_path": upscaler.data_path, "model_url": None, "scale": upscaler.scale} for upscaler in shared.sd_upscalers]
def get_sd_models():
from modules import sd_models, sd_models_config
return [{"title": x.title, "model_name": x.name, "filename": x.filename, "type": x.type, "hash": x.shorthash, "sha256": x.sha256, "config": sd_models_config.find_checkpoint_config_near_filename(x)} for x in sd_models.checkpoints_list.values()]
def get_hypernetworks():
return [{"name": name, "path": shared.hypernetworks[name]} for name in shared.hypernetworks]
def get_face_restorers():
return [{"name":x.name(), "cmd_dir": getattr(x, "cmd_dir", None)} for x in shared.face_restorers]
def get_prompt_styles():
return [{ 'name': v.name, 'prompt': v.prompt, 'negative_prompt': v.negative_prompt, 'extra': v.extra, 'filename': v.filename, 'preview': v.preview} for v in shared.prompt_styles.styles.values()]
def get_embeddings():
from modules import sd_hijack
db = sd_hijack.model_hijack.embedding_db
def convert_embedding(embedding):
return {"step": embedding.step, "sd_checkpoint": embedding.sd_checkpoint, "sd_checkpoint_name": embedding.sd_checkpoint_name, "shape": embedding.shape, "vectors": embedding.vectors}
def convert_embeddings(embeddings):
return {embedding.name: convert_embedding(embedding) for embedding in embeddings.values()}
return {"loaded": convert_embeddings(db.word_embeddings), "skipped": convert_embeddings(db.skipped_embeddings)}
def get_extra_networks(page: Optional[str] = None, name: Optional[str] = None, filename: Optional[str] = None, title: Optional[str] = None, fullname: Optional[str] = None, hash: Optional[str] = None): # pylint: disable=redefined-builtin
res = []
for pg in shared.extra_networks:
if page is not None and pg.name != page.lower():
continue
for item in pg.items:
if name is not None and item.get('name', '') != name:
continue
if title is not None and item.get('title', '') != title:
continue
if filename is not None and item.get('filename', '') != filename:
continue
if fullname is not None and item.get('fullname', '') != fullname:
continue
if hash is not None and (item.get('shorthash', None) or item.get('hash')) != hash:
continue
res.append({
'name': item.get('name', ''),
'type': pg.name,
'title': item.get('title', None),
'fullname': item.get('fullname', None),
'filename': item.get('filename', None),
'hash': item.get('shorthash', None) or item.get('hash'),
"preview": item.get('preview', None),
})
return res
def get_interrogate():
from modules.ui_interrogate import get_models
return ['clip', 'deepdanbooru'] + get_models()
def post_interrogate(req: models.InterrogateRequest):
if req.image is None or len(req.image) < 64:
raise HTTPException(status_code=404, detail="Image not found")
image = helpers.decode_base64_to_image(req.image)
image = image.convert('RGB')
if req.model == "clip":
caption = shared.interrogator.interrogate(image)
return models.InterrogateResponse(caption)
elif req.model == "deepdanbooru":
from mobules import deepbooru
caption = deepbooru.model.tag(image)
return models.InterrogateResponse(caption)
else:
from modules.ui_interrogate import interrogate_image, analyze_image, get_models
if req.model not in get_models():
raise HTTPException(status_code=404, detail="Model not found")
caption = interrogate_image(image, model=req.model, mode=req.mode)
if not req.analyze:
return models.InterrogateResponse(caption)
else:
medium, artist, movement, trending, flavor = analyze_image(image, model=req.model)
return models.InterrogateResponse(caption, medium, artist, movement, trending, flavor)
def post_unload_checkpoint():
from modules import sd_models
sd_models.unload_model_weights(op='model')
sd_models.unload_model_weights(op='refiner')
return {}
def post_reload_checkpoint():
from modules import sd_models
sd_models.reload_model_weights()
return {}
def post_refresh_checkpoints():
return shared.refresh_checkpoints()
def post_refresh_vae():
return shared.refresh_vaes()
def get_extensions_list():
from modules import extensions
extensions.list_extensions()
ext_list = []
for ext in extensions.extensions:
ext: extensions.Extension
ext.read_info()
if ext.remote is not None:
ext_list.append({
"name": ext.name,
"remote": ext.remote,
"branch": ext.branch,
"commit_hash":ext.commit_hash,
"commit_date":ext.commit_date,
"version":ext.version,
"enabled":ext.enabled
})
return ext_list
def post_pnginfo(req: models.PNGInfoRequest):
from modules import images, script_callbacks, generation_parameters_copypaste
if not req.image.strip():
return models.PNGInfoResponse(info="")
image = helpers.decode_base64_to_image(req.image.strip())
if image is None:
return models.PNGInfoResponse(info="")
geninfo, items = images.read_info_from_image(image)
if geninfo is None:
geninfo = ""
if items and items['parameters']:
del items['parameters']
params = generation_parameters_copypaste.parse_generation_parameters(geninfo)
script_callbacks.infotext_pasted_callback(geninfo, params)
return models.PNGInfoResponse(info=geninfo, items=items, parameters=params)
modules/api/helpers.py
+57
View File
@@ -0,0 +1,57 @@
import io
import base64
from PIL import Image, PngImagePlugin
import piexif
import piexif.helper
from fastapi.exceptions import HTTPException
from modules import shared, images
def decode_base64_to_image(encoding):
if encoding.startswith("data:image/"):
encoding = encoding.split(";")[1].split(",")[1]
try:
image = Image.open(io.BytesIO(base64.b64decode(encoding)))
return image
except Exception as e:
shared.log.warning(f'API cannot decode image: {e}')
raise HTTPException(status_code=500, detail="Invalid encoded image") from e
def encode_pil_to_base64(image):
with io.BytesIO() as output_bytes:
images.save_image(image, output_bytes, shared.opts.samples_format)
bytes_data = output_bytes.getvalue()
return base64.b64encode(bytes_data)
def upscaler_to_index(name: str):
try:
return [x.name.lower() for x in shared.sd_upscalers].index(name.lower())
except Exception as e:
raise HTTPException(status_code=400, detail=f"Invalid upscaler, needs to be one of these: {' , '.join([x.name for x in shared.sd_upscalers])}") from e
def save_image(image, fn, ext):
# actual save
parameters = image.info.get('parameters', None)
image_format = Image.registered_extensions()[f'.{ext}']
if image_format == 'PNG':
pnginfo_data = PngImagePlugin.PngInfo()
for k, v in image.info.items():
pnginfo_data.add_text(k, str(v))
image.save(fn, format=image_format, quality=shared.opts.jpeg_quality, pnginfo=pnginfo_data)
elif image_format == 'JPEG':
if image.mode == 'RGBA':
shared.log.warning('Saving RGBA image as JPEG: Alpha channel will be lost')
image = image.convert("RGB")
elif image.mode == 'I;16':
image = image.point(lambda p: p * 0.0038910505836576).convert("L")
exif_bytes = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(parameters or "", encoding="unicode") } })
image.save(fn, format=image_format, quality=shared.opts.jpeg_quality, exif=exif_bytes)
elif image_format == 'WEBP':
if image.mode == 'I;16':
image = image.point(lambda p: p * 0.0038910505836576).convert("RGB")
exif_bytes = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(parameters or "", encoding="unicode") } })
image.save(fn, format=image_format, quality=shared.opts.jpeg_quality, lossless=shared.opts.webp_lossless, exif=exif_bytes)
else:
# shared.log.warning(f'Unrecognized image format: {extension} attempting save as {image_format}')
image.save(fn, format=image_format, quality=shared.opts.jpeg_quality)
modules/middleware.py → modules/api/middleware.py
View File
modules/api/models.py
+136 -134
View File
@@ -91,14 +91,81 @@ class PydanticModelGenerator:
DynamicModel.__config__.allow_mutation = True
return DynamicModel
### item classes
class IPAdapterItem(BaseModel):
class ItemSampler(BaseModel):
name: str = Field(title="Name")
aliases: List[str] = Field(title="Aliases")
options: Dict[str, str] = Field(title="Options")
class ItemVae(BaseModel):
model_name: str = Field(title="Model Name")
filename: str = Field(title="Filename")
class ItemUpscaler(BaseModel):
name: str = Field(title="Name")
model_name: Optional[str] = Field(title="Model Name")
model_path: Optional[str] = Field(title="Path")
model_url: Optional[str] = Field(title="URL")
scale: Optional[float] = Field(title="Scale")
class ItemModel(BaseModel):
title: str = Field(title="Title")
model_name: str = Field(title="Model Name")
filename: str = Field(title="Filename")
type: str = Field(title="Model type")
sha256: Optional[str] = Field(title="SHA256 hash")
hash: Optional[str] = Field(title="Short hash")
config: Optional[str] = Field(title="Config file")
class ItemHypernetwork(BaseModel):
name: str = Field(title="Name")
path: Optional[str] = Field(title="Path")
class ItemFaceRestorer(BaseModel):
name: str = Field(title="Name")
cmd_dir: Optional[str] = Field(title="Path")
class ItemGAN(BaseModel):
name: str = Field(title="Name")
path: Optional[str] = Field(title="Path")
scale: Optional[int] = Field(title="Scale")
class ItemStyle(BaseModel):
name: str = Field(title="Name")
prompt: Optional[str] = Field(title="Prompt")
negative_prompt: Optional[str] = Field(title="Negative Prompt")
extra: Optional[str] = Field(title="Extra")
filename: Optional[str] = Field(title="Filename")
preview: Optional[str] = Field(title="Preview")
class ItemExtraNetwork(BaseModel):
name: str = Field(title="Name")
type: str = Field(title="Type")
title: Optional[str] = Field(title="Title")
fullname: Optional[str] = Field(title="Fullname")
filename: Optional[str] = Field(title="Filename")
hash: Optional[str] = Field(title="Hash")
preview: Optional[str] = Field(title="Preview image URL")
class ItemArtist(BaseModel):
name: str = Field(title="Name")
score: float = Field(title="Score")
category: str = Field(title="Category")
class ItemEmbedding(BaseModel):
step: Optional[int] = Field(title="Step", description="The number of steps that were used to train this embedding, if available")
sd_checkpoint: Optional[str] = Field(title="SD Checkpoint", description="The hash of the checkpoint this embedding was trained on, if available")
sd_checkpoint_name: Optional[str] = Field(title="SD Checkpoint Name", description="The name of the checkpoint this embedding was trained on, if available. Note that this is the name that was used by the trainer; for a stable identifier, use `sd_checkpoint` instead")
shape: int = Field(title="Shape", description="The length of each individual vector in the embedding")
vectors: int = Field(title="Vectors", description="The number of vectors in the embedding")
class ItemIPAdapter(BaseModel):
adapter: str = Field(title="Adapter", default="Base", description="Adapter to use")
image: str = Field(title="Image", default="", description="Adapter image, must be a base64 string containing the image's data.")
scale: float = Field(title="Scale", default=0.5, gt=0, le=1, description="Scale of the adapter image, must be between 0 and 1.")
class FaceIDItem(BaseModel):
class ItemFaceID(BaseModel):
mode: list[str] = Field(title="Mode", default=["FaceID"], description="The mode to use (available values: FaceID, FaceSwap).")
model: str = Field(title="Model", default="FaceID Base", description="The FaceID model to use.")
image: str = Field(title="Image", default="", description="Source face image, must be a base64 string containing the image's data.")
@@ -109,8 +176,32 @@ class FaceIDItem(BaseModel):
tokens: int = Field(title="Tokens", default=4, ge=1, le=16, description="Amount of tokens to use, must be between 1 and 16.")
cache_model: bool = Field(title="Cache", default=True, description="Should the model be cached?")
class ScriptArg(BaseModel):
label: str = Field(default=None, title="Label", description="Name of the argument in UI")
value: Optional[Any] = Field(default=None, title="Value", description="Default value of the argument")
minimum: Optional[Any] = Field(default=None, title="Minimum", description="Minimum allowed value for the argumentin UI")
maximum: Optional[Any] = Field(default=None, title="Minimum", description="Maximum allowed value for the argumentin UI")
step: Optional[Any] = Field(default=None, title="Minimum", description="Step for changing value of the argumentin UI")
choices: Optional[Any] = Field(default=None, title="Choices", description="Possible values for the argument")
StableDiffusionTxt2ImgProcessingAPI = PydanticModelGenerator(
class ItemScript(BaseModel):
name: str = Field(default=None, title="Name", description="Script name")
is_alwayson: bool = Field(default=None, title="IsAlwayson", description="Flag specifying whether this script is an alwayson script")
is_img2img: bool = Field(default=None, title="IsImg2img", description="Flag specifying whether this script is an img2img script")
args: List[ScriptArg] = Field(title="Arguments", description="List of script's arguments")
class ItemExtension(BaseModel):
name: str = Field(title="Name", description="Extension name")
remote: str = Field(title="Remote", description="Extension Repository URL")
branch: str = Field(title="Branch", description="Extension Repository Branch")
commit_hash: str = Field(title="Commit Hash", description="Extension Repository Commit Hash")
version: str = Field(title="Version", description="Extension Version")
commit_date: str = Field(title="Commit Date", description="Extension Repository Commit Date")
enabled: bool = Field(title="Enabled", description="Flag specifying whether this extension is enabled")
### request/response classes
ReqTxt2Img = PydanticModelGenerator(
"StableDiffusionProcessingTxt2Img",
StableDiffusionProcessingTxt2Img,
[
@@ -120,12 +211,17 @@ StableDiffusionTxt2ImgProcessingAPI = PydanticModelGenerator(
{"key": "send_images", "type": bool, "default": True},
{"key": "save_images", "type": bool, "default": False},
{"key": "alwayson_scripts", "type": dict, "default": {}},
{"key": "ip_adapter", "type": Optional[IPAdapterItem], "default": None, "exclude": True},
{"key": "face_id", "type": Optional[FaceIDItem], "default": None, "exclude": True},
{"key": "ip_adapter", "type": Optional[ItemIPAdapter], "default": None, "exclude": True},
{"key": "face_id", "type": Optional[ItemFaceID], "default": None, "exclude": True},
]
).generate_model()
StableDiffusionImg2ImgProcessingAPI = PydanticModelGenerator(
class ResTxt2Img(BaseModel):
images: List[str] = Field(default=None, title="Image", description="The generated image in base64 format.")
parameters: dict
info: str
ReqImg2Img = PydanticModelGenerator(
"StableDiffusionProcessingImg2Img",
StableDiffusionProcessingImg2Img,
[
@@ -139,22 +235,21 @@ StableDiffusionImg2ImgProcessingAPI = PydanticModelGenerator(
{"key": "send_images", "type": bool, "default": True},
{"key": "save_images", "type": bool, "default": False},
{"key": "alwayson_scripts", "type": dict, "default": {}},
{"key": "ip_adapter", "type": Optional[IPAdapterItem], "default": None, "exclude": True},
{"key": "face_id", "type": Optional[FaceIDItem], "default": None, "exclude": True},
{"key": "ip_adapter", "type": Optional[ItemIPAdapter], "default": None, "exclude": True},
{"key": "face_id", "type": Optional[ItemFaceID], "default": None, "exclude": True},
]
).generate_model()
class TextToImageResponse(BaseModel):
class ResImg2Img(BaseModel):
images: List[str] = Field(default=None, title="Image", description="The generated image in base64 format.")
parameters: dict
info: str
class ImageToImageResponse(BaseModel):
images: List[str] = Field(default=None, title="Image", description="The generated image in base64 format.")
parameters: dict
info: str
class FileData(BaseModel):
data: str = Field(title="File data", description="Base64 representation of the file")
name: str = Field(title="File name")
class ExtrasBaseRequest(BaseModel):
class ReqProcess(BaseModel):
resize_mode: float = Field(default=0, title="Resize Mode", description="Sets the resize mode: 0 to upscale by upscaling_resize amount, 1 to upscale up to upscaling_resize_h x upscaling_resize_w.")
show_extras_results: bool = Field(default=True, title="Show results", description="Should the backend return the generated image?")
gfpgan_visibility: float = Field(default=0, title="GFPGAN Visibility", ge=0, le=1, allow_inf_nan=False, description="Sets the visibility of GFPGAN, values should be between 0 and 1.")
@@ -169,61 +264,62 @@ class ExtrasBaseRequest(BaseModel):
extras_upscaler_2_visibility: float = Field(default=0, title="Secondary upscaler visibility", ge=0, le=1, allow_inf_nan=False, description="Sets the visibility of secondary upscaler, values should be between 0 and 1.")
upscale_first: bool = Field(default=False, title="Upscale first", description="Should the upscaler run before restoring faces?")
class ExtraBaseResponse(BaseModel):
class ResProcess(BaseModel):
html_info: str = Field(title="HTML info", description="A series of HTML tags containing the process info.")
class ExtrasSingleImageRequest(ExtrasBaseRequest):
class ReqProcessImage(ReqProcess):
image: str = Field(default="", title="Image", description="Image to work on, must be a Base64 string containing the image's data.")
class ExtrasSingleImageResponse(ExtraBaseResponse):
class ResProcessImage(ResProcess):
image: str = Field(default=None, title="Image", description="The generated image in base64 format.")
class FileData(BaseModel):
data: str = Field(title="File data", description="Base64 representation of the file")
name: str = Field(title="File name")
class ExtrasBatchImagesRequest(ExtrasBaseRequest):
class ReqProcessBatch(ReqProcess):
imageList: List[FileData] = Field(title="Images", description="List of images to work on. Must be Base64 strings")
class ExtrasBatchImagesResponse(ExtraBaseResponse):
class ResProcessBatch(ResProcess):
images: List[str] = Field(title="Images", description="The generated images in base64 format.")
class PNGInfoRequest(BaseModel):
class ReqImageInfo(BaseModel):
image: str = Field(title="Image", description="The base64 encoded PNG image")
class PNGInfoResponse(BaseModel):
class ResImageInfo(BaseModel):
info: str = Field(title="Image info", description="A string with the parameters used to generate the image")
items: dict = Field(title="Items", description="A dictionary containing all the other fields the image had")
parameters: dict = Field(title="Parameters", description="A dictionary with parsed generation info fields")
class LogRequest(BaseModel):
class ReqLog(BaseModel):
lines: int = Field(default=100, title="Lines", description="How many lines to return")
clear: bool = Field(default=False, title="Clear", description="Should the log be cleared after returning the lines?")
class ProgressRequest(BaseModel):
class ReqProgress(BaseModel):
skip_current_image: bool = Field(default=False, title="Skip current image", description="Skip current image serialization")
class ProgressResponse(BaseModel):
class ResProgress(BaseModel):
progress: float = Field(title="Progress", description="The progress with a range of 0 to 1")
eta_relative: float = Field(title="ETA in secs")
state: dict = Field(title="State", description="The current state snapshot")
current_image: str = Field(default=None, title="Current image", description="The current image in base64 format. opts.show_progress_every_n_steps is required for this to work.")
textinfo: str = Field(default=None, title="Info text", description="Info text used by WebUI.")
class InterrogateRequest(BaseModel):
class ReqInterrogate(BaseModel):
image: str = Field(default="", title="Image", description="Image to work on, must be a Base64 string containing the image's data.")
model: str = Field(default="clip", title="Model", description="The interrogate model used.")
class InterrogateResponse(BaseModel):
caption: str = Field(default=None, title="Caption", description="The generated caption for the image.")
class ResInterrogate(BaseModel):
caption: Optional[str] = Field(default=None, title="Caption", description="The generated caption for the image.")
medium: Optional[str] = Field(default=None, title="Medium", description="Image medium.")
artist: Optional[str] = Field(default=None, title="Medium", description="Image artist.")
movement: Optional[str] = Field(default=None, title="Medium", description="Image movement.")
trending: Optional[str] = Field(default=None, title="Medium", description="Image trending.")
flavor: Optional[str] = Field(default=None, title="Medium", description="Image flavor.")
class TrainResponse(BaseModel):
class ResTrain(BaseModel):
info: str = Field(title="Train info", description="Response string from train embedding or hypernetwork task.")
class CreateResponse(BaseModel):
class ResCreate(BaseModel):
info: str = Field(title="Create info", description="Response string from create embedding or hypernetwork task.")
class PreprocessResponse(BaseModel):
class ResPreprocess(BaseModel):
info: str = Field(title="Preprocess info", description="Response string from preprocessing task.")
fields = {}
@@ -251,109 +347,15 @@ for key in _options:
FlagsModel = create_model("Flags", **flags)
class SamplerItem(BaseModel):
name: str = Field(title="Name")
aliases: List[str] = Field(title="Aliases")
options: Dict[str, str] = Field(title="Options")
class ResEmbeddings(BaseModel):
loaded: Dict[str, ItemEmbedding] = Field(title="Loaded", description="Embeddings loaded for the current model")
skipped: Dict[str, ItemEmbedding] = Field(title="Skipped", description="Embeddings skipped for the current model (likely due to architecture incompatibility)")
class SDVaeItem(BaseModel):
model_name: str = Field(title="Model Name")
filename: str = Field(title="Filename")
class UpscalerItem(BaseModel):
name: str = Field(title="Name")
model_name: Optional[str] = Field(title="Model Name")
model_path: Optional[str] = Field(title="Path")
model_url: Optional[str] = Field(title="URL")
scale: Optional[float] = Field(title="Scale")
class SDModelItem(BaseModel):
title: str = Field(title="Title")
model_name: str = Field(title="Model Name")
filename: str = Field(title="Filename")
type: str = Field(title="Model type")
sha256: Optional[str] = Field(title="SHA256 hash")
hash: Optional[str] = Field(title="Short hash")
config: Optional[str] = Field(title="Config file")
class HypernetworkItem(BaseModel):
name: str = Field(title="Name")
path: Optional[str] = Field(title="Path")
class FaceRestorerItem(BaseModel):
name: str = Field(title="Name")
cmd_dir: Optional[str] = Field(title="Path")
class RealesrganItem(BaseModel):
name: str = Field(title="Name")
path: Optional[str] = Field(title="Path")
scale: Optional[int] = Field(title="Scale")
class StyleItem(BaseModel):
name: str = Field(title="Name")
prompt: Optional[str] = Field(title="Prompt")
negative_prompt: Optional[str] = Field(title="Negative Prompt")
extra: Optional[str] = Field(title="Extra")
filename: Optional[str] = Field(title="Filename")
preview: Optional[str] = Field(title="Preview")
class ExtraNetworkItem(BaseModel):
name: str = Field(title="Name")
type: str = Field(title="Type")
title: Optional[str] = Field(title="Title")
fullname: Optional[str] = Field(title="Fullname")
filename: Optional[str] = Field(title="Filename")
hash: Optional[str] = Field(title="Hash")
preview: Optional[str] = Field(title="Preview image URL")
# description: Optional[str] = Field(title="Description")
# info: Optional[str] = Field(title="Information")
# metadata: Optional[Any] = Field(title="Metadata")
# local: Optional[str] = Field(title="Local")
class ArtistItem(BaseModel):
name: str = Field(title="Name")
score: float = Field(title="Score")
category: str = Field(title="Category")
class EmbeddingItem(BaseModel):
step: Optional[int] = Field(title="Step", description="The number of steps that were used to train this embedding, if available")
sd_checkpoint: Optional[str] = Field(title="SD Checkpoint", description="The hash of the checkpoint this embedding was trained on, if available")
sd_checkpoint_name: Optional[str] = Field(title="SD Checkpoint Name", description="The name of the checkpoint this embedding was trained on, if available. Note that this is the name that was used by the trainer; for a stable identifier, use `sd_checkpoint` instead")
shape: int = Field(title="Shape", description="The length of each individual vector in the embedding")
vectors: int = Field(title="Vectors", description="The number of vectors in the embedding")
class EmbeddingsResponse(BaseModel):
loaded: Dict[str, EmbeddingItem] = Field(title="Loaded", description="Embeddings loaded for the current model")
skipped: Dict[str, EmbeddingItem] = Field(title="Skipped", description="Embeddings skipped for the current model (likely due to architecture incompatibility)")
class MemoryResponse(BaseModel):
class ResMemory(BaseModel):
ram: dict = Field(title="RAM", description="System memory stats")
cuda: dict = Field(title="CUDA", description="nVidia CUDA memory stats")
class ScriptsList(BaseModel):
class ResScripts(BaseModel):
txt2img: list = Field(default=None, title="Txt2img", description="Titles of scripts (txt2img)")
img2img: list = Field(default=None, title="Img2img", description="Titles of scripts (img2img)")
control: list = Field(default=None, title="Control", description="Titles of scripts (control)")
class ScriptArg(BaseModel):
label: str = Field(default=None, title="Label", description="Name of the argument in UI")
value: Optional[Any] = Field(default=None, title="Value", description="Default value of the argument")
minimum: Optional[Any] = Field(default=None, title="Minimum", description="Minimum allowed value for the argumentin UI")
maximum: Optional[Any] = Field(default=None, title="Minimum", description="Maximum allowed value for the argumentin UI")
step: Optional[Any] = Field(default=None, title="Minimum", description="Step for changing value of the argumentin UI")
choices: Optional[Any] = Field(default=None, title="Choices", description="Possible values for the argument")
class ScriptInfo(BaseModel):
name: str = Field(default=None, title="Name", description="Script name")
is_alwayson: bool = Field(default=None, title="IsAlwayson", description="Flag specifying whether this script is an alwayson script")
is_img2img: bool = Field(default=None, title="IsImg2img", description="Flag specifying whether this script is an img2img script")
args: List[ScriptArg] = Field(title="Arguments", description="List of script's arguments")
class ExtensionItem(BaseModel):
name: str = Field(title="Name", description="Extension name")
remote: str = Field(title="Remote", description="Extension Repository URL")
branch: str = Field(title="Branch", description="Extension Repository Branch")
commit_hash: str = Field(title="Commit Hash", description="Extension Repository Commit Hash")
version: str = Field(title="Version", description="Extension Version")
commit_date: str = Field(title="Commit Date", description="Extension Repository Commit Date")
enabled: bool = Field(title="Enabled", description="Flag specifying whether this extension is enabled")
modules/api/script.py
+82
View File
@@ -0,0 +1,82 @@
from typing import Optional
from fastapi.exceptions import HTTPException
import gradio as gr
from modules.api import models
from modules import scripts
def script_name_to_index(name, scripts_list):
try:
return [script.title().lower() for script in scripts_list].index(name.lower())
except Exception as e:
raise HTTPException(status_code=422, detail=f"Script '{name}' not found") from e
def get_selectable_script(script_name, script_runner):
if script_name is None or script_name == "":
return None, None
script_idx = script_name_to_index(script_name, script_runner.selectable_scripts)
script = script_runner.selectable_scripts[script_idx]
return script, script_idx
def get_scripts_list():
t2ilist = [script.name for script in scripts.scripts_txt2img.scripts if script.name is not None]
i2ilist = [script.name for script in scripts.scripts_img2img.scripts if script.name is not None]
control = [script.name for script in scripts.scripts_control.scripts if script.name is not None]
return models.ScriptsList(txt2img = t2ilist, img2img = i2ilist, control = control)
def get_script_info(script_name: Optional[str] = None):
res = []
for script_list in [scripts.scripts_txt2img.scripts, scripts.scripts_img2img.scripts, scripts.scripts_control.scripts]:
for script in script_list:
if script.api_info is not None and (script_name is None or script_name == script.api_info.name):
res.append(script.api_info)
return res
def get_script(script_name, script_runner):
if script_name is None or script_name == "":
return None, None
script_idx = script_name_to_index(script_name, script_runner.scripts)
return script_runner.scripts[script_idx]
def init_default_script_args(script_runner):
#find max idx from the scripts in runner and generate a none array to init script_args
last_arg_index = 1
for script in script_runner.scripts:
if last_arg_index < script.args_to:
last_arg_index = script.args_to
# None everywhere except position 0 to initialize script args
script_args = [None]*last_arg_index
script_args[0] = 0
# get default values
if gr is None:
return script_args
with gr.Blocks(): # will throw errors calling ui function without this
for script in script_runner.scripts:
if script.ui(script.is_img2img):
ui_default_values = []
for elem in script.ui(script.is_img2img):
ui_default_values.append(elem.value)
script_args[script.args_from:script.args_to] = ui_default_values
return script_args
def init_script_args(p, request, default_script_args, selectable_scripts, selectable_script_idx, script_runner):
script_args = default_script_args.copy()
# position 0 in script_arg is the idx+1 of the selectable script that is going to be run when using scripts.scripts_*2img.run()
if selectable_scripts:
script_args[selectable_scripts.args_from:selectable_scripts.args_to] = request.script_args
script_args[0] = selectable_script_idx + 1
# Now check for always on scripts
if request.alwayson_scripts and (len(request.alwayson_scripts) > 0):
for alwayson_script_name in request.alwayson_scripts.keys():
alwayson_script = get_script(alwayson_script_name, script_runner)
if alwayson_script is None:
raise HTTPException(status_code=422, detail=f"Always on script not found: {alwayson_script_name}")
if not alwayson_script.alwayson:
raise HTTPException(status_code=422, detail=f"Selectable script cannot be in always on params: {alwayson_script_name}")
if "args" in request.alwayson_scripts[alwayson_script_name]:
# min between arg length in scriptrunner and arg length in the request
for idx in range(0, min((alwayson_script.args_to - alwayson_script.args_from), len(request.alwayson_scripts[alwayson_script_name]["args"]))):
script_args[alwayson_script.args_from + idx] = request.alwayson_scripts[alwayson_script_name]["args"][idx]
p.per_script_args[alwayson_script.title()] = request.alwayson_scripts[alwayson_script_name]["args"]
return script_args
modules/api/server.py
+116
View File
@@ -0,0 +1,116 @@
from typing import Any, Dict
from fastapi import Depends
from modules import shared
from modules.api import models, helpers
def post_shutdown():
shared.log.info('Shutdown request received')
import sys
sys.exit(0)
def get_motd():
import requests
from installer import get_version
motd = ''
ver = get_version()
if ver.get('updated', None) is not None:
motd = f"version <b>{ver['hash']} {ver['updated']}</b> <span style='color: var(--primary-500)'>{ver['url'].split('/')[-1]}</span><br>"
if shared.opts.motd:
res = requests.get('https://vladmandic.github.io/automatic/motd', timeout=10)
if res.status_code == 200:
msg = (res.text or '').strip()
shared.log.info(f'MOTD: {msg if len(msg) > 0 else "N/A"}')
motd += res.text
return motd
def get_log_buffer(req: models.LogRequest = Depends()):
lines = shared.log.buffer[:req.lines] if req.lines > 0 else shared.log.buffer.copy()
if req.clear:
shared.log.buffer.clear()
return lines
def get_config():
options = {}
for k in shared.opts.data.keys():
if shared.opts.data_labels.get(k) is not None:
options.update({k: shared.opts.data.get(k, shared.opts.data_labels.get(k).default)})
else:
options.update({k: shared.opts.data.get(k, None)})
if 'sd_lyco' in options:
del options['sd_lyco']
if 'sd_lora' in options:
del options['sd_lora']
return options
def set_config(req: Dict[str, Any]):
updated = []
for k, v in req.items():
updated.append({ k: shared.opts.set(k, v) })
shared.opts.save(shared.config_filename)
return { "updated": updated }
def get_cmd_flags():
return vars(shared.cmd_opts)
def get_progress(req: models.ProgressRequest = Depends()):
import time
if shared.state.job_count == 0:
return models.ProgressResponse(progress=0, eta_relative=0, state=shared.state.dict(), textinfo=shared.state.textinfo)
shared.state.do_set_current_image()
current_image = None
if shared.state.current_image and not req.skip_current_image:
current_image = helpers.encode_pil_to_base64(shared.state.current_image)
batch_x = max(shared.state.job_no, 0)
batch_y = max(shared.state.job_count, 1)
step_x = max(shared.state.sampling_step, 0)
step_y = max(shared.state.sampling_steps, 1)
current = step_y * batch_x + step_x
total = step_y * batch_y
progress = current / total if current > 0 and total > 0 else 0
time_since_start = time.time() - shared.state.time_start
eta_relative = (time_since_start / progress) - time_since_start if progress > 0 else 0
res = models.ProgressResponse(progress=progress, eta_relative=eta_relative, state=shared.state.dict(), current_image=current_image, textinfo=shared.state.textinfo)
return res
def post_interrupt():
shared.state.interrupt()
return {}
def post_skip():
shared.state.skip()
def get_memory():
try:
import os
import psutil
process = psutil.Process(os.getpid())
res = process.memory_info() # only rss is cross-platform guaranteed so we dont rely on other values
ram_total = 100 * res.rss / process.memory_percent() # and total memory is calculated as actual value is not cross-platform safe
ram = { 'free': ram_total - res.rss, 'used': res.rss, 'total': ram_total }
except Exception as err:
ram = { 'error': f'{err}' }
try:
import torch
if torch.cuda.is_available():
s = torch.cuda.mem_get_info()
system = { 'free': s[0], 'used': s[1] - s[0], 'total': s[1] }
s = dict(torch.cuda.memory_stats(shared.device))
allocated = { 'current': s['allocated_bytes.all.current'], 'peak': s['allocated_bytes.all.peak'] }
reserved = { 'current': s['reserved_bytes.all.current'], 'peak': s['reserved_bytes.all.peak'] }
active = { 'current': s['active_bytes.all.current'], 'peak': s['active_bytes.all.peak'] }
inactive = { 'current': s['inactive_split_bytes.all.current'], 'peak': s['inactive_split_bytes.all.peak'] }
warnings = { 'retries': s['num_alloc_retries'], 'oom': s['num_ooms'] }
cuda = {
'system': system,
'active': active,
'allocated': allocated,
'reserved': reserved,
'inactive': inactive,
'events': warnings,
}
else:
cuda = { 'error': 'unavailable' }
except Exception as err:
cuda = { 'error': f'{err}' }
return models.MemoryResponse(ram = ram, cuda = cuda)
modules/api/train.py
+90
View File
@@ -0,0 +1,90 @@
from modules import shared, sd_hijack, devices
from modules.api import models
from modules.textual_inversion.preprocess import preprocess
def post_create_embedding(args: dict):
from modules.textual_inversion.textual_inversion import create_embedding
try:
shared.state.begin('api-embedding')
filename = create_embedding(**args) # create empty embedding
sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings() # reload embeddings so new one can be immediately used
shared.state.end()
return models.CreateResponse(info = f"create embedding filename: {filename}")
except AssertionError as e:
shared.state.end()
return models.TrainResponse(info = f"create embedding error: {e}")
def post_create_hypernetwork(args: dict):
from modules.hypernetworks.hypernetwork import create_hypernetwork
try:
shared.state.begin('api-hypernetwork')
filename = create_hypernetwork(**args) # create empty embedding # pylint: disable=E1111
shared.state.end()
return models.CreateResponse(info = f"create hypernetwork filename: {filename}")
except AssertionError as e:
shared.state.end()
return models.TrainResponse(info = f"create hypernetwork error: {e}")
def post_preprocess(args: dict):
try:
shared.state.begin('api-preprocess')
preprocess(**args) # quick operation unless blip/booru interrogation is enabled
shared.state.end()
return models.PreprocessResponse(info = 'preprocess complete')
except KeyError as e:
shared.state.end()
return models.PreprocessResponse(info = f"preprocess error: invalid token: {e}")
except AssertionError as e:
shared.state.end()
return models.PreprocessResponse(info = f"preprocess error: {e}")
except FileNotFoundError as e:
shared.state.end()
return models.PreprocessResponse(info = f'preprocess error: {e}')
def post_train_embedding(args: dict):
from modules.textual_inversion.textual_inversion import train_embedding
try:
shared.state.begin('api-embedding')
apply_optimizations = False
error = None
filename = ''
if not apply_optimizations:
sd_hijack.undo_optimizations()
try:
_embedding, filename = train_embedding(**args) # can take a long time to complete
except Exception as e:
error = e
finally:
if not apply_optimizations:
sd_hijack.apply_optimizations()
shared.state.end()
return models.TrainResponse(info = f"train embedding complete: filename: {filename} error: {error}")
except AssertionError as msg:
shared.state.end()
return models.TrainResponse(info = f"train embedding error: {msg}")
def post_train_hypernetwork(args: dict):
from modules.hypernetworks.hypernetwork import train_hypernetwork
try:
shared.state.begin('api-hypernetwork')
shared.loaded_hypernetworks = []
apply_optimizations = False
error = None
filename = ''
if not apply_optimizations:
sd_hijack.undo_optimizations()
try:
_hypernetwork, filename = train_hypernetwork(**args)
except Exception as e:
error = e
finally:
shared.sd_model.cond_stage_model.to(devices.device)
shared.sd_model.first_stage_model.to(devices.device)
if not apply_optimizations:
sd_hijack.apply_optimizations()
shared.state.end()
return models.TrainResponse(info=f"train embedding complete: filename: {filename} error: {error}")
except AssertionError:
shared.state.end()
return models.TrainResponse(info=f"train embedding error: {error}")
modules/shared.py
+1
View File
@@ -27,6 +27,7 @@ from installer import log as central_logger # pylint: disable=E0611
errors.install([gr])
demo: gr.Blocks = None
api = None
log = central_logger
progress_print_out = sys.stdout
parser = cmd_args.parser
modules/ui_interrogate.py
+36 -87
View File
@@ -1,11 +1,7 @@
import os
import base64
from io import BytesIO
import gradio as gr
import torch
from PIL import Image
from pydantic import BaseModel, Field # pylint: disable=no-name-in-module
from fastapi.exceptions import HTTPException
import modules.generation_parameters_copypaste as parameters_copypaste
from modules import devices, lowvram, shared, paths, ui_common
@@ -29,7 +25,12 @@ class BatchWriter:
self.file.close()
def load(clip_model_name):
def get_models():
import open_clip
return ['/'.join(x) for x in open_clip.list_pretrained()]
def load_interrogator(clip_model_name):
from clip_interrogator import Config, Interrogator
global ci # pylint: disable=global-statement
if ci is None:
@@ -54,23 +55,6 @@ def unload():
devices.torch_gc()
def image_analysis(image, clip_model_name):
load(clip_model_name)
image = image.convert('RGB')
image_features = ci.image_to_features(image)
top_mediums = ci.mediums.rank(image_features, 5)
top_artists = ci.artists.rank(image_features, 5)
top_movements = ci.movements.rank(image_features, 5)
top_trendings = ci.trendings.rank(image_features, 5)
top_flavors = ci.flavors.rank(image_features, 5)
medium_ranks = dict(zip(top_mediums, ci.similarities(image_features, top_mediums)))
artist_ranks = dict(zip(top_artists, ci.similarities(image_features, top_artists)))
movement_ranks = dict(zip(top_movements, ci.similarities(image_features, top_movements)))
trending_ranks = dict(zip(top_trendings, ci.similarities(image_features, top_trendings)))
flavor_ranks = dict(zip(top_flavors, ci.similarities(image_features, top_flavors)))
return medium_ranks, artist_ranks, movement_ranks, trending_ranks, flavor_ranks
def interrogate(image, mode, caption=None):
shared.log.info(f'Interrogate: image={image} mode={mode} config={ci.config}')
if mode == 'best':
@@ -88,14 +72,14 @@ def interrogate(image, mode, caption=None):
return prompt
def image_to_prompt(image, mode, clip_model_name):
def interrogate_image(image, model, mode):
shared.state.begin()
shared.state.job = 'interrogate'
try:
if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
lowvram.send_everything_to_cpu()
devices.torch_gc()
load(clip_model_name)
load_interrogator(model)
image = image.convert('RGB')
shared.log.info(f'Interrogate: image={image} mode={mode} config={ci.config}')
prompt = interrogate(image, mode)
@@ -106,12 +90,7 @@ def image_to_prompt(image, mode, clip_model_name):
return prompt
def get_models():
import open_clip
return ['/'.join(x) for x in open_clip.list_pretrained()]
def batch_process(batch_files, batch_folder, batch_str, mode, clip_model, write):
def interrogate_batch(batch_files, batch_folder, batch_str, model, mode, write):
files = []
if batch_files is not None:
files += [f.name for f in batch_files]
@@ -122,7 +101,6 @@ def batch_process(batch_files, batch_folder, batch_str, mode, clip_model, write)
if len(files) == 0:
shared.log.error('Interrogate batch no images')
return ''
shared.log.info(f'Interrogate batch: images={len(files)} mode={mode} config={ci.config}')
shared.state.begin()
shared.state.job = 'batch interrogate'
prompts = []
@@ -130,7 +108,8 @@ def batch_process(batch_files, batch_folder, batch_str, mode, clip_model, write)
if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
lowvram.send_everything_to_cpu()
devices.torch_gc()
load(clip_model)
load_interrogator(model)
shared.log.info(f'Interrogate batch: images={len(files)} mode={mode} config={ci.config}')
captions = []
# first pass: generate captions
for file in files:
@@ -168,6 +147,23 @@ def batch_process(batch_files, batch_folder, batch_str, mode, clip_model, write)
return '\n\n'.join(prompts)
def analyze_image(image, model):
load_interrogator(model)
image = image.convert('RGB')
image_features = ci.image_to_features(image)
top_mediums = ci.mediums.rank(image_features, 5)
top_artists = ci.artists.rank(image_features, 5)
top_movements = ci.movements.rank(image_features, 5)
top_trendings = ci.trendings.rank(image_features, 5)
top_flavors = ci.flavors.rank(image_features, 5)
medium_ranks = dict(zip(top_mediums, ci.similarities(image_features, top_mediums)))
artist_ranks = dict(zip(top_artists, ci.similarities(image_features, top_artists)))
movement_ranks = dict(zip(top_movements, ci.similarities(image_features, top_movements)))
trending_ranks = dict(zip(top_trendings, ci.similarities(image_features, top_trendings)))
flavor_ranks = dict(zip(top_flavors, ci.similarities(image_features, top_flavors)))
return medium_ranks, artist_ranks, movement_ranks, trending_ranks, flavor_ranks
def create_ui():
global low_vram # pylint: disable=global-statement
low_vram = shared.cmd_opts.lowvram or shared.cmd_opts.medvram
@@ -190,9 +186,9 @@ def create_ui():
trending = gr.Label(label="Trending", num_top_classes=5)
flavor = gr.Label(label="Flavor", num_top_classes=5)
with gr.Row():
interrogate_btn = gr.Button("Interrogate", variant='primary')
analyze_btn = gr.Button("Analyze", variant='primary')
unload_btn = gr.Button("Unload")
btn_interrogate_img = gr.Button("Interrogate", variant='primary')
btn_analyze_img = gr.Button("Analyze", variant='primary')
btn_unload = gr.Button("Unload")
with gr.Row():
buttons = parameters_copypaste.create_buttons(["txt2img", "img2img", "extras", "control"])
for tabname, button in buttons.items():
@@ -209,7 +205,7 @@ def create_ui():
with gr.Row():
write = gr.Checkbox(label='Write prompts to files', value=False)
with gr.Row():
batch_btn = gr.Button("Interrogate", variant='primary')
btn_interrogate_batch = gr.Button("Interrogate", variant='primary')
with gr.Column():
with gr.Row():
# clip_model = gr.Dropdown(get_models(), value='ViT-L-14/openai', label='CLIP Model')
@@ -217,54 +213,7 @@ def create_ui():
ui_common.create_refresh_button(clip_model, get_models, lambda: {"choices": get_models()}, 'refresh_interrogate_models')
with gr.Row():
mode = gr.Radio(['best', 'fast', 'classic', 'caption', 'negative'], label='Mode', value='best')
interrogate_btn.click(image_to_prompt, inputs=[image, mode, clip_model], outputs=prompt)
analyze_btn.click(image_analysis, inputs=[image, clip_model], outputs=[medium, artist, movement, trending, flavor])
unload_btn.click(unload)
batch_btn.click(batch_process, inputs=[batch_files, batch_folder, batch_str, mode, clip_model, write], outputs=[batch])
def decode_base64_to_image(encoding):
if encoding.startswith("data:image/"):
encoding = encoding.split(";")[1].split(",")[1]
try:
image = Image.open(BytesIO(base64.b64decode(encoding)))
return image
except Exception as e:
raise HTTPException(status_code=500, detail="Invalid encoded image") from e
# TODO redesign interrogator api
def mount_interrogator_api(_: gr.Blocks, app):
class InterrogatorAnalyzeRequest(BaseModel):
image: str = Field(default="", title="Image", description="Image to work on, must be a Base64 string containing the image's data.")
clip_model_name: str = Field(default="ViT-L-14/openai", title="Model", description="The interrogate model used. See the models endpoint for a list of available models.")
class InterrogatorPromptRequest(InterrogatorAnalyzeRequest):
mode: str = Field(default="fast", title="Mode", description="The mode used to generate the prompt. Can be one of: best, fast, classic, negative.")
@app.get("/interrogator/models")
async def api_get_models():
import open_clip
return ["/".join(x) for x in open_clip.list_pretrained()]
@app.post("/interrogator/prompt")
async def api_get_prompt(analyzereq: InterrogatorPromptRequest):
image_b64 = analyzereq.image
if image_b64 is None:
raise HTTPException(status_code=404, detail="Image not found")
img = decode_base64_to_image(image_b64)
prompt = image_to_prompt(img, analyzereq.mode, analyzereq.clip_model_name)
return {"prompt": prompt}
@app.post("/interrogator/analyze")
async def api_analyze(analyzereq: InterrogatorAnalyzeRequest):
image_b64 = analyzereq.image
if image_b64 is None:
raise HTTPException(status_code=404, detail="Image not found")
img = decode_base64_to_image(image_b64)
(medium_ranks, artist_ranks, movement_ranks, trending_ranks, flavor_ranks) = image_analysis(img, analyzereq.clip_model_name)
return {"medium": medium_ranks, "artist": artist_ranks, "movement": movement_ranks, "trending": trending_ranks, "flavor": flavor_ranks}
# script_callbacks.on_app_started(mount_interrogator_api)
btn_interrogate_img.click(interrogate_image, inputs=[image, clip_model, mode], outputs=prompt)
btn_analyze_img.click(analyze_image, inputs=[image, clip_model], outputs=[medium, artist, movement, trending, flavor])
btn_interrogate_batch.click(interrogate_batch, inputs=[batch_files, batch_folder, batch_str, clip_model, mode, write], outputs=[batch])
btn_unload.click(unload)
webui.py
+5 -5
View File
@@ -31,7 +31,7 @@ import modules.upscaler
import modules.textual_inversion.textual_inversion
import modules.hypernetworks.hypernetwork
import modules.script_callbacks
from modules.middleware import setup_middleware
from modules.api.middleware import setup_middleware
from modules.shared import cmd_opts, opts
@@ -284,7 +284,7 @@ def start_ui():
timer.startup.record("launch")
modules.progress.setup_progress_api(app)
create_api(app)
shared.api = create_api(app)
timer.startup.record("api")
ui_extra_networks.init_api(app)
@@ -347,12 +347,12 @@ def api_only():
from fastapi import FastAPI
app = FastAPI(**fastapi_args)
setup_middleware(app, cmd_opts)
api = create_api(app)
api.wants_restart = False
shared.api = create_api(app)
shared.api.wants_restart = False
modules.script_callbacks.app_started_callback(None, app)
modules.sd_models.write_metadata()
log.info(f"Startup time: {timer.startup.summary()}")
server = api.launch()
server = shared.api.launch()
return server