api.py 36 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806
  1. import base64
  2. import io
  3. import os
  4. import time
  5. import datetime
  6. import uvicorn
  7. import ipaddress
  8. import requests
  9. import gradio as gr
  10. from threading import Lock
  11. from io import BytesIO
  12. from fastapi import APIRouter, Depends, FastAPI, Request, Response
  13. from fastapi.security import HTTPBasic, HTTPBasicCredentials
  14. from fastapi.exceptions import HTTPException
  15. from fastapi.responses import JSONResponse
  16. from fastapi.encoders import jsonable_encoder
  17. from secrets import compare_digest
  18. import modules.shared as shared
  19. from modules import sd_samplers, deepbooru, sd_hijack, images, scripts, ui, postprocessing, errors, restart, shared_items, script_callbacks, generation_parameters_copypaste, sd_models
  20. from modules.api import models
  21. from modules.shared import opts
  22. from modules.processing import StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img, process_images
  23. from modules.textual_inversion.textual_inversion import create_embedding, train_embedding
  24. from modules.textual_inversion.preprocess import preprocess
  25. from modules.hypernetworks.hypernetwork import create_hypernetwork, train_hypernetwork
  26. from PIL import PngImagePlugin, Image
  27. from modules.sd_models_config import find_checkpoint_config_near_filename
  28. from modules.realesrgan_model import get_realesrgan_models
  29. from modules import devices
  30. from typing import Any
  31. import piexif
  32. import piexif.helper
  33. from contextlib import closing
  34. def script_name_to_index(name, scripts):
  35. try:
  36. return [script.title().lower() for script in scripts].index(name.lower())
  37. except Exception as e:
  38. raise HTTPException(status_code=422, detail=f"Script '{name}' not found") from e
  39. def validate_sampler_name(name):
  40. config = sd_samplers.all_samplers_map.get(name, None)
  41. if config is None:
  42. raise HTTPException(status_code=404, detail="Sampler not found")
  43. return name
  44. def setUpscalers(req: dict):
  45. reqDict = vars(req)
  46. reqDict['extras_upscaler_1'] = reqDict.pop('upscaler_1', None)
  47. reqDict['extras_upscaler_2'] = reqDict.pop('upscaler_2', None)
  48. return reqDict
  49. def verify_url(url):
  50. """Returns True if the url refers to a global resource."""
  51. import socket
  52. from urllib.parse import urlparse
  53. try:
  54. parsed_url = urlparse(url)
  55. domain_name = parsed_url.netloc
  56. host = socket.gethostbyname_ex(domain_name)
  57. for ip in host[2]:
  58. ip_addr = ipaddress.ip_address(ip)
  59. if not ip_addr.is_global:
  60. return False
  61. except Exception:
  62. return False
  63. return True
  64. def decode_base64_to_image(encoding):
  65. if encoding.startswith("http://") or encoding.startswith("https://"):
  66. if not opts.api_enable_requests:
  67. raise HTTPException(status_code=500, detail="Requests not allowed")
  68. if opts.api_forbid_local_requests and not verify_url(encoding):
  69. raise HTTPException(status_code=500, detail="Request to local resource not allowed")
  70. headers = {'user-agent': opts.api_useragent} if opts.api_useragent else {}
  71. response = requests.get(encoding, timeout=30, headers=headers)
  72. try:
  73. image = Image.open(BytesIO(response.content))
  74. return image
  75. except Exception as e:
  76. raise HTTPException(status_code=500, detail="Invalid image url") from e
  77. if encoding.startswith("data:image/"):
  78. encoding = encoding.split(";")[1].split(",")[1]
  79. try:
  80. image = Image.open(BytesIO(base64.b64decode(encoding)))
  81. return image
  82. except Exception as e:
  83. raise HTTPException(status_code=500, detail="Invalid encoded image") from e
  84. def encode_pil_to_base64(image):
  85. with io.BytesIO() as output_bytes:
  86. if isinstance(image, str):
  87. return image
  88. if opts.samples_format.lower() == 'png':
  89. use_metadata = False
  90. metadata = PngImagePlugin.PngInfo()
  91. for key, value in image.info.items():
  92. if isinstance(key, str) and isinstance(value, str):
  93. metadata.add_text(key, value)
  94. use_metadata = True
  95. image.save(output_bytes, format="PNG", pnginfo=(metadata if use_metadata else None), quality=opts.jpeg_quality)
  96. elif opts.samples_format.lower() in ("jpg", "jpeg", "webp"):
  97. if image.mode == "RGBA":
  98. image = image.convert("RGB")
  99. parameters = image.info.get('parameters', None)
  100. exif_bytes = piexif.dump({
  101. "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(parameters or "", encoding="unicode") }
  102. })
  103. if opts.samples_format.lower() in ("jpg", "jpeg"):
  104. image.save(output_bytes, format="JPEG", exif = exif_bytes, quality=opts.jpeg_quality)
  105. else:
  106. image.save(output_bytes, format="WEBP", exif = exif_bytes, quality=opts.jpeg_quality)
  107. else:
  108. raise HTTPException(status_code=500, detail="Invalid image format")
  109. bytes_data = output_bytes.getvalue()
  110. return base64.b64encode(bytes_data)
  111. def api_middleware(app: FastAPI):
  112. rich_available = False
  113. try:
  114. if os.environ.get('WEBUI_RICH_EXCEPTIONS', None) is not None:
  115. import anyio # importing just so it can be placed on silent list
  116. import starlette # importing just so it can be placed on silent list
  117. from rich.console import Console
  118. console = Console()
  119. rich_available = True
  120. except Exception:
  121. pass
  122. @app.middleware("http")
  123. async def log_and_time(req: Request, call_next):
  124. ts = time.time()
  125. res: Response = await call_next(req)
  126. duration = str(round(time.time() - ts, 4))
  127. res.headers["X-Process-Time"] = duration
  128. endpoint = req.scope.get('path', 'err')
  129. if shared.cmd_opts.api_log and endpoint.startswith('/sdapi'):
  130. print('API {t} {code} {prot}/{ver} {method} {endpoint} {cli} {duration}'.format(
  131. t=datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"),
  132. code=res.status_code,
  133. ver=req.scope.get('http_version', '0.0'),
  134. cli=req.scope.get('client', ('0:0.0.0', 0))[0],
  135. prot=req.scope.get('scheme', 'err'),
  136. method=req.scope.get('method', 'err'),
  137. endpoint=endpoint,
  138. duration=duration,
  139. ))
  140. return res
  141. def handle_exception(request: Request, e: Exception):
  142. err = {
  143. "error": type(e).__name__,
  144. "detail": vars(e).get('detail', ''),
  145. "body": vars(e).get('body', ''),
  146. "errors": str(e),
  147. }
  148. if not isinstance(e, HTTPException): # do not print backtrace on known httpexceptions
  149. message = f"API error: {request.method}: {request.url} {err}"
  150. if rich_available:
  151. print(message)
  152. console.print_exception(show_locals=True, max_frames=2, extra_lines=1, suppress=[anyio, starlette], word_wrap=False, width=min([console.width, 200]))
  153. else:
  154. errors.report(message, exc_info=True)
  155. return JSONResponse(status_code=vars(e).get('status_code', 500), content=jsonable_encoder(err))
  156. @app.middleware("http")
  157. async def exception_handling(request: Request, call_next):
  158. try:
  159. return await call_next(request)
  160. except Exception as e:
  161. return handle_exception(request, e)
  162. @app.exception_handler(Exception)
  163. async def fastapi_exception_handler(request: Request, e: Exception):
  164. return handle_exception(request, e)
  165. @app.exception_handler(HTTPException)
  166. async def http_exception_handler(request: Request, e: HTTPException):
  167. return handle_exception(request, e)
  168. class Api:
  169. def __init__(self, app: FastAPI, queue_lock: Lock):
  170. if shared.cmd_opts.api_auth:
  171. self.credentials = {}
  172. for auth in shared.cmd_opts.api_auth.split(","):
  173. user, password = auth.split(":")
  174. self.credentials[user] = password
  175. self.router = APIRouter()
  176. self.app = app
  177. self.queue_lock = queue_lock
  178. api_middleware(self.app)
  179. self.add_api_route("/sdapi/v1/txt2img", self.text2imgapi, methods=["POST"], response_model=models.TextToImageResponse)
  180. self.add_api_route("/sdapi/v1/img2img", self.img2imgapi, methods=["POST"], response_model=models.ImageToImageResponse)
  181. self.add_api_route("/sdapi/v1/extra-single-image", self.extras_single_image_api, methods=["POST"], response_model=models.ExtrasSingleImageResponse)
  182. self.add_api_route("/sdapi/v1/extra-batch-images", self.extras_batch_images_api, methods=["POST"], response_model=models.ExtrasBatchImagesResponse)
  183. self.add_api_route("/sdapi/v1/png-info", self.pnginfoapi, methods=["POST"], response_model=models.PNGInfoResponse)
  184. self.add_api_route("/sdapi/v1/progress", self.progressapi, methods=["GET"], response_model=models.ProgressResponse)
  185. self.add_api_route("/sdapi/v1/interrogate", self.interrogateapi, methods=["POST"])
  186. self.add_api_route("/sdapi/v1/interrupt", self.interruptapi, methods=["POST"])
  187. self.add_api_route("/sdapi/v1/skip", self.skip, methods=["POST"])
  188. self.add_api_route("/sdapi/v1/options", self.get_config, methods=["GET"], response_model=models.OptionsModel)
  189. self.add_api_route("/sdapi/v1/options", self.set_config, methods=["POST"])
  190. self.add_api_route("/sdapi/v1/cmd-flags", self.get_cmd_flags, methods=["GET"], response_model=models.FlagsModel)
  191. self.add_api_route("/sdapi/v1/samplers", self.get_samplers, methods=["GET"], response_model=list[models.SamplerItem])
  192. self.add_api_route("/sdapi/v1/upscalers", self.get_upscalers, methods=["GET"], response_model=list[models.UpscalerItem])
  193. self.add_api_route("/sdapi/v1/latent-upscale-modes", self.get_latent_upscale_modes, methods=["GET"], response_model=list[models.LatentUpscalerModeItem])
  194. self.add_api_route("/sdapi/v1/sd-models", self.get_sd_models, methods=["GET"], response_model=list[models.SDModelItem])
  195. self.add_api_route("/sdapi/v1/sd-vae", self.get_sd_vaes, methods=["GET"], response_model=list[models.SDVaeItem])
  196. self.add_api_route("/sdapi/v1/hypernetworks", self.get_hypernetworks, methods=["GET"], response_model=list[models.HypernetworkItem])
  197. self.add_api_route("/sdapi/v1/face-restorers", self.get_face_restorers, methods=["GET"], response_model=list[models.FaceRestorerItem])
  198. self.add_api_route("/sdapi/v1/realesrgan-models", self.get_realesrgan_models, methods=["GET"], response_model=list[models.RealesrganItem])
  199. self.add_api_route("/sdapi/v1/prompt-styles", self.get_prompt_styles, methods=["GET"], response_model=list[models.PromptStyleItem])
  200. self.add_api_route("/sdapi/v1/embeddings", self.get_embeddings, methods=["GET"], response_model=models.EmbeddingsResponse)
  201. self.add_api_route("/sdapi/v1/refresh-checkpoints", self.refresh_checkpoints, methods=["POST"])
  202. self.add_api_route("/sdapi/v1/refresh-vae", self.refresh_vae, methods=["POST"])
  203. self.add_api_route("/sdapi/v1/create/embedding", self.create_embedding, methods=["POST"], response_model=models.CreateResponse)
  204. self.add_api_route("/sdapi/v1/create/hypernetwork", self.create_hypernetwork, methods=["POST"], response_model=models.CreateResponse)
  205. self.add_api_route("/sdapi/v1/preprocess", self.preprocess, methods=["POST"], response_model=models.PreprocessResponse)
  206. self.add_api_route("/sdapi/v1/train/embedding", self.train_embedding, methods=["POST"], response_model=models.TrainResponse)
  207. self.add_api_route("/sdapi/v1/train/hypernetwork", self.train_hypernetwork, methods=["POST"], response_model=models.TrainResponse)
  208. self.add_api_route("/sdapi/v1/memory", self.get_memory, methods=["GET"], response_model=models.MemoryResponse)
  209. self.add_api_route("/sdapi/v1/unload-checkpoint", self.unloadapi, methods=["POST"])
  210. self.add_api_route("/sdapi/v1/reload-checkpoint", self.reloadapi, methods=["POST"])
  211. self.add_api_route("/sdapi/v1/scripts", self.get_scripts_list, methods=["GET"], response_model=models.ScriptsList)
  212. self.add_api_route("/sdapi/v1/script-info", self.get_script_info, methods=["GET"], response_model=list[models.ScriptInfo])
  213. self.add_api_route("/sdapi/v1/extensions", self.get_extensions_list, methods=["GET"], response_model=list[models.ExtensionItem])
  214. if shared.cmd_opts.api_server_stop:
  215. self.add_api_route("/sdapi/v1/server-kill", self.kill_webui, methods=["POST"])
  216. self.add_api_route("/sdapi/v1/server-restart", self.restart_webui, methods=["POST"])
  217. self.add_api_route("/sdapi/v1/server-stop", self.stop_webui, methods=["POST"])
  218. self.default_script_arg_txt2img = []
  219. self.default_script_arg_img2img = []
  220. def add_api_route(self, path: str, endpoint, **kwargs):
  221. if shared.cmd_opts.api_auth:
  222. return self.app.add_api_route(path, endpoint, dependencies=[Depends(self.auth)], **kwargs)
  223. return self.app.add_api_route(path, endpoint, **kwargs)
  224. def auth(self, credentials: HTTPBasicCredentials = Depends(HTTPBasic())):
  225. if credentials.username in self.credentials:
  226. if compare_digest(credentials.password, self.credentials[credentials.username]):
  227. return True
  228. raise HTTPException(status_code=401, detail="Incorrect username or password", headers={"WWW-Authenticate": "Basic"})
  229. def get_selectable_script(self, script_name, script_runner):
  230. if script_name is None or script_name == "":
  231. return None, None
  232. script_idx = script_name_to_index(script_name, script_runner.selectable_scripts)
  233. script = script_runner.selectable_scripts[script_idx]
  234. return script, script_idx
  235. def get_scripts_list(self):
  236. t2ilist = [script.name for script in scripts.scripts_txt2img.scripts if script.name is not None]
  237. i2ilist = [script.name for script in scripts.scripts_img2img.scripts if script.name is not None]
  238. return models.ScriptsList(txt2img=t2ilist, img2img=i2ilist)
  239. def get_script_info(self):
  240. res = []
  241. for script_list in [scripts.scripts_txt2img.scripts, scripts.scripts_img2img.scripts]:
  242. res += [script.api_info for script in script_list if script.api_info is not None]
  243. return res
  244. def get_script(self, script_name, script_runner):
  245. if script_name is None or script_name == "":
  246. return None, None
  247. script_idx = script_name_to_index(script_name, script_runner.scripts)
  248. return script_runner.scripts[script_idx]
  249. def init_default_script_args(self, script_runner):
  250. #find max idx from the scripts in runner and generate a none array to init script_args
  251. last_arg_index = 1
  252. for script in script_runner.scripts:
  253. if last_arg_index < script.args_to:
  254. last_arg_index = script.args_to
  255. # None everywhere except position 0 to initialize script args
  256. script_args = [None]*last_arg_index
  257. script_args[0] = 0
  258. # get default values
  259. with gr.Blocks(): # will throw errors calling ui function without this
  260. for script in script_runner.scripts:
  261. if script.ui(script.is_img2img):
  262. ui_default_values = []
  263. for elem in script.ui(script.is_img2img):
  264. ui_default_values.append(elem.value)
  265. script_args[script.args_from:script.args_to] = ui_default_values
  266. return script_args
  267. def init_script_args(self, request, default_script_args, selectable_scripts, selectable_idx, script_runner):
  268. script_args = default_script_args.copy()
  269. # 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()
  270. if selectable_scripts:
  271. script_args[selectable_scripts.args_from:selectable_scripts.args_to] = request.script_args
  272. script_args[0] = selectable_idx + 1
  273. # Now check for always on scripts
  274. if request.alwayson_scripts:
  275. for alwayson_script_name in request.alwayson_scripts.keys():
  276. alwayson_script = self.get_script(alwayson_script_name, script_runner)
  277. if alwayson_script is None:
  278. raise HTTPException(status_code=422, detail=f"always on script {alwayson_script_name} not found")
  279. # Selectable script in always on script param check
  280. if alwayson_script.alwayson is False:
  281. raise HTTPException(status_code=422, detail="Cannot have a selectable script in the always on scripts params")
  282. # always on script with no arg should always run so you don't really need to add them to the requests
  283. if "args" in request.alwayson_scripts[alwayson_script_name]:
  284. # min between arg length in scriptrunner and arg length in the request
  285. for idx in range(0, min((alwayson_script.args_to - alwayson_script.args_from), len(request.alwayson_scripts[alwayson_script_name]["args"]))):
  286. script_args[alwayson_script.args_from + idx] = request.alwayson_scripts[alwayson_script_name]["args"][idx]
  287. return script_args
  288. def text2imgapi(self, txt2imgreq: models.StableDiffusionTxt2ImgProcessingAPI):
  289. script_runner = scripts.scripts_txt2img
  290. if not script_runner.scripts:
  291. script_runner.initialize_scripts(False)
  292. ui.create_ui()
  293. if not self.default_script_arg_txt2img:
  294. self.default_script_arg_txt2img = self.init_default_script_args(script_runner)
  295. selectable_scripts, selectable_script_idx = self.get_selectable_script(txt2imgreq.script_name, script_runner)
  296. populate = txt2imgreq.copy(update={ # Override __init__ params
  297. "sampler_name": validate_sampler_name(txt2imgreq.sampler_name or txt2imgreq.sampler_index),
  298. "do_not_save_samples": not txt2imgreq.save_images,
  299. "do_not_save_grid": not txt2imgreq.save_images,
  300. })
  301. if populate.sampler_name:
  302. populate.sampler_index = None # prevent a warning later on
  303. args = vars(populate)
  304. args.pop('script_name', None)
  305. args.pop('script_args', None) # will refeed them to the pipeline directly after initializing them
  306. args.pop('alwayson_scripts', None)
  307. script_args = self.init_script_args(txt2imgreq, self.default_script_arg_txt2img, selectable_scripts, selectable_script_idx, script_runner)
  308. send_images = args.pop('send_images', True)
  309. args.pop('save_images', None)
  310. with self.queue_lock:
  311. with closing(StableDiffusionProcessingTxt2Img(sd_model=shared.sd_model, **args)) as p:
  312. p.is_api = True
  313. p.scripts = script_runner
  314. p.outpath_grids = opts.outdir_txt2img_grids
  315. p.outpath_samples = opts.outdir_txt2img_samples
  316. try:
  317. shared.state.begin(job="scripts_txt2img")
  318. if selectable_scripts is not None:
  319. p.script_args = script_args
  320. processed = scripts.scripts_txt2img.run(p, *p.script_args) # Need to pass args as list here
  321. else:
  322. p.script_args = tuple(script_args) # Need to pass args as tuple here
  323. processed = process_images(p)
  324. finally:
  325. shared.state.end()
  326. shared.total_tqdm.clear()
  327. b64images = list(map(encode_pil_to_base64, processed.images)) if send_images else []
  328. return models.TextToImageResponse(images=b64images, parameters=vars(txt2imgreq), info=processed.js())
  329. def img2imgapi(self, img2imgreq: models.StableDiffusionImg2ImgProcessingAPI):
  330. init_images = img2imgreq.init_images
  331. if init_images is None:
  332. raise HTTPException(status_code=404, detail="Init image not found")
  333. mask = img2imgreq.mask
  334. if mask:
  335. mask = decode_base64_to_image(mask)
  336. script_runner = scripts.scripts_img2img
  337. if not script_runner.scripts:
  338. script_runner.initialize_scripts(True)
  339. ui.create_ui()
  340. if not self.default_script_arg_img2img:
  341. self.default_script_arg_img2img = self.init_default_script_args(script_runner)
  342. selectable_scripts, selectable_script_idx = self.get_selectable_script(img2imgreq.script_name, script_runner)
  343. populate = img2imgreq.copy(update={ # Override __init__ params
  344. "sampler_name": validate_sampler_name(img2imgreq.sampler_name or img2imgreq.sampler_index),
  345. "do_not_save_samples": not img2imgreq.save_images,
  346. "do_not_save_grid": not img2imgreq.save_images,
  347. "mask": mask,
  348. })
  349. if populate.sampler_name:
  350. populate.sampler_index = None # prevent a warning later on
  351. args = vars(populate)
  352. 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.
  353. args.pop('script_name', None)
  354. args.pop('script_args', None) # will refeed them to the pipeline directly after initializing them
  355. args.pop('alwayson_scripts', None)
  356. script_args = self.init_script_args(img2imgreq, self.default_script_arg_img2img, selectable_scripts, selectable_script_idx, script_runner)
  357. send_images = args.pop('send_images', True)
  358. args.pop('save_images', None)
  359. with self.queue_lock:
  360. with closing(StableDiffusionProcessingImg2Img(sd_model=shared.sd_model, **args)) as p:
  361. p.init_images = [decode_base64_to_image(x) for x in init_images]
  362. p.is_api = True
  363. p.scripts = script_runner
  364. p.outpath_grids = opts.outdir_img2img_grids
  365. p.outpath_samples = opts.outdir_img2img_samples
  366. try:
  367. shared.state.begin(job="scripts_img2img")
  368. if selectable_scripts is not None:
  369. p.script_args = script_args
  370. processed = scripts.scripts_img2img.run(p, *p.script_args) # Need to pass args as list here
  371. else:
  372. p.script_args = tuple(script_args) # Need to pass args as tuple here
  373. processed = process_images(p)
  374. finally:
  375. shared.state.end()
  376. shared.total_tqdm.clear()
  377. b64images = list(map(encode_pil_to_base64, processed.images)) if send_images else []
  378. if not img2imgreq.include_init_images:
  379. img2imgreq.init_images = None
  380. img2imgreq.mask = None
  381. return models.ImageToImageResponse(images=b64images, parameters=vars(img2imgreq), info=processed.js())
  382. def extras_single_image_api(self, req: models.ExtrasSingleImageRequest):
  383. reqDict = setUpscalers(req)
  384. reqDict['image'] = decode_base64_to_image(reqDict['image'])
  385. with self.queue_lock:
  386. result = postprocessing.run_extras(extras_mode=0, image_folder="", input_dir="", output_dir="", save_output=False, **reqDict)
  387. return models.ExtrasSingleImageResponse(image=encode_pil_to_base64(result[0][0]), html_info=result[1])
  388. def extras_batch_images_api(self, req: models.ExtrasBatchImagesRequest):
  389. reqDict = setUpscalers(req)
  390. image_list = reqDict.pop('imageList', [])
  391. image_folder = [decode_base64_to_image(x.data) for x in image_list]
  392. with self.queue_lock:
  393. result = postprocessing.run_extras(extras_mode=1, image_folder=image_folder, image="", input_dir="", output_dir="", save_output=False, **reqDict)
  394. return models.ExtrasBatchImagesResponse(images=list(map(encode_pil_to_base64, result[0])), html_info=result[1])
  395. def pnginfoapi(self, req: models.PNGInfoRequest):
  396. image = decode_base64_to_image(req.image.strip())
  397. if image is None:
  398. return models.PNGInfoResponse(info="")
  399. geninfo, items = images.read_info_from_image(image)
  400. if geninfo is None:
  401. geninfo = ""
  402. params = generation_parameters_copypaste.parse_generation_parameters(geninfo)
  403. script_callbacks.infotext_pasted_callback(geninfo, params)
  404. return models.PNGInfoResponse(info=geninfo, items=items, parameters=params)
  405. def progressapi(self, req: models.ProgressRequest = Depends()):
  406. # copy from check_progress_call of ui.py
  407. if shared.state.job_count == 0:
  408. return models.ProgressResponse(progress=0, eta_relative=0, state=shared.state.dict(), textinfo=shared.state.textinfo)
  409. # avoid dividing zero
  410. progress = 0.01
  411. if shared.state.job_count > 0:
  412. progress += shared.state.job_no / shared.state.job_count
  413. if shared.state.sampling_steps > 0:
  414. progress += 1 / shared.state.job_count * shared.state.sampling_step / shared.state.sampling_steps
  415. time_since_start = time.time() - shared.state.time_start
  416. eta = (time_since_start/progress)
  417. eta_relative = eta-time_since_start
  418. progress = min(progress, 1)
  419. shared.state.set_current_image()
  420. current_image = None
  421. if shared.state.current_image and not req.skip_current_image:
  422. current_image = encode_pil_to_base64(shared.state.current_image)
  423. return models.ProgressResponse(progress=progress, eta_relative=eta_relative, state=shared.state.dict(), current_image=current_image, textinfo=shared.state.textinfo)
  424. def interrogateapi(self, interrogatereq: models.InterrogateRequest):
  425. image_b64 = interrogatereq.image
  426. if image_b64 is None:
  427. raise HTTPException(status_code=404, detail="Image not found")
  428. img = decode_base64_to_image(image_b64)
  429. img = img.convert('RGB')
  430. # Override object param
  431. with self.queue_lock:
  432. if interrogatereq.model == "clip":
  433. processed = shared.interrogator.interrogate(img)
  434. elif interrogatereq.model == "deepdanbooru":
  435. processed = deepbooru.model.tag(img)
  436. else:
  437. raise HTTPException(status_code=404, detail="Model not found")
  438. return models.InterrogateResponse(caption=processed)
  439. def interruptapi(self):
  440. shared.state.interrupt()
  441. return {}
  442. def unloadapi(self):
  443. sd_models.unload_model_weights()
  444. return {}
  445. def reloadapi(self):
  446. sd_models.send_model_to_device(shared.sd_model)
  447. return {}
  448. def skip(self):
  449. shared.state.skip()
  450. def get_config(self):
  451. options = {}
  452. for key in shared.opts.data.keys():
  453. metadata = shared.opts.data_labels.get(key)
  454. if(metadata is not None):
  455. options.update({key: shared.opts.data.get(key, shared.opts.data_labels.get(key).default)})
  456. else:
  457. options.update({key: shared.opts.data.get(key, None)})
  458. return options
  459. def set_config(self, req: dict[str, Any]):
  460. checkpoint_name = req.get("sd_model_checkpoint", None)
  461. if checkpoint_name is not None and checkpoint_name not in sd_models.checkpoint_aliases:
  462. raise RuntimeError(f"model {checkpoint_name!r} not found")
  463. for k, v in req.items():
  464. shared.opts.set(k, v, is_api=True)
  465. shared.opts.save(shared.config_filename)
  466. return
  467. def get_cmd_flags(self):
  468. return vars(shared.cmd_opts)
  469. def get_samplers(self):
  470. return [{"name": sampler[0], "aliases":sampler[2], "options":sampler[3]} for sampler in sd_samplers.all_samplers]
  471. def get_upscalers(self):
  472. return [
  473. {
  474. "name": upscaler.name,
  475. "model_name": upscaler.scaler.model_name,
  476. "model_path": upscaler.data_path,
  477. "model_url": None,
  478. "scale": upscaler.scale,
  479. }
  480. for upscaler in shared.sd_upscalers
  481. ]
  482. def get_latent_upscale_modes(self):
  483. return [
  484. {
  485. "name": upscale_mode,
  486. }
  487. for upscale_mode in [*(shared.latent_upscale_modes or {})]
  488. ]
  489. def get_sd_models(self):
  490. import modules.sd_models as sd_models
  491. return [{"title": x.title, "model_name": x.model_name, "hash": x.shorthash, "sha256": x.sha256, "filename": x.filename, "config": find_checkpoint_config_near_filename(x)} for x in sd_models.checkpoints_list.values()]
  492. def get_sd_vaes(self):
  493. import modules.sd_vae as sd_vae
  494. return [{"model_name": x, "filename": sd_vae.vae_dict[x]} for x in sd_vae.vae_dict.keys()]
  495. def get_hypernetworks(self):
  496. return [{"name": name, "path": shared.hypernetworks[name]} for name in shared.hypernetworks]
  497. def get_face_restorers(self):
  498. return [{"name":x.name(), "cmd_dir": getattr(x, "cmd_dir", None)} for x in shared.face_restorers]
  499. def get_realesrgan_models(self):
  500. return [{"name":x.name,"path":x.data_path, "scale":x.scale} for x in get_realesrgan_models(None)]
  501. def get_prompt_styles(self):
  502. styleList = []
  503. for k in shared.prompt_styles.styles:
  504. style = shared.prompt_styles.styles[k]
  505. styleList.append({"name":style[0], "prompt": style[1], "negative_prompt": style[2]})
  506. return styleList
  507. def get_embeddings(self):
  508. db = sd_hijack.model_hijack.embedding_db
  509. def convert_embedding(embedding):
  510. return {
  511. "step": embedding.step,
  512. "sd_checkpoint": embedding.sd_checkpoint,
  513. "sd_checkpoint_name": embedding.sd_checkpoint_name,
  514. "shape": embedding.shape,
  515. "vectors": embedding.vectors,
  516. }
  517. def convert_embeddings(embeddings):
  518. return {embedding.name: convert_embedding(embedding) for embedding in embeddings.values()}
  519. return {
  520. "loaded": convert_embeddings(db.word_embeddings),
  521. "skipped": convert_embeddings(db.skipped_embeddings),
  522. }
  523. def refresh_checkpoints(self):
  524. with self.queue_lock:
  525. shared.refresh_checkpoints()
  526. def refresh_vae(self):
  527. with self.queue_lock:
  528. shared_items.refresh_vae_list()
  529. def create_embedding(self, args: dict):
  530. try:
  531. shared.state.begin(job="create_embedding")
  532. filename = create_embedding(**args) # create empty embedding
  533. sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings() # reload embeddings so new one can be immediately used
  534. return models.CreateResponse(info=f"create embedding filename: {filename}")
  535. except AssertionError as e:
  536. return models.TrainResponse(info=f"create embedding error: {e}")
  537. finally:
  538. shared.state.end()
  539. def create_hypernetwork(self, args: dict):
  540. try:
  541. shared.state.begin(job="create_hypernetwork")
  542. filename = create_hypernetwork(**args) # create empty embedding
  543. return models.CreateResponse(info=f"create hypernetwork filename: {filename}")
  544. except AssertionError as e:
  545. return models.TrainResponse(info=f"create hypernetwork error: {e}")
  546. finally:
  547. shared.state.end()
  548. def preprocess(self, args: dict):
  549. try:
  550. shared.state.begin(job="preprocess")
  551. preprocess(**args) # quick operation unless blip/booru interrogation is enabled
  552. shared.state.end()
  553. return models.PreprocessResponse(info='preprocess complete')
  554. except KeyError as e:
  555. return models.PreprocessResponse(info=f"preprocess error: invalid token: {e}")
  556. except Exception as e:
  557. return models.PreprocessResponse(info=f"preprocess error: {e}")
  558. finally:
  559. shared.state.end()
  560. def train_embedding(self, args: dict):
  561. try:
  562. shared.state.begin(job="train_embedding")
  563. apply_optimizations = shared.opts.training_xattention_optimizations
  564. error = None
  565. filename = ''
  566. if not apply_optimizations:
  567. sd_hijack.undo_optimizations()
  568. try:
  569. embedding, filename = train_embedding(**args) # can take a long time to complete
  570. except Exception as e:
  571. error = e
  572. finally:
  573. if not apply_optimizations:
  574. sd_hijack.apply_optimizations()
  575. return models.TrainResponse(info=f"train embedding complete: filename: {filename} error: {error}")
  576. except Exception as msg:
  577. return models.TrainResponse(info=f"train embedding error: {msg}")
  578. finally:
  579. shared.state.end()
  580. def train_hypernetwork(self, args: dict):
  581. try:
  582. shared.state.begin(job="train_hypernetwork")
  583. shared.loaded_hypernetworks = []
  584. apply_optimizations = shared.opts.training_xattention_optimizations
  585. error = None
  586. filename = ''
  587. if not apply_optimizations:
  588. sd_hijack.undo_optimizations()
  589. try:
  590. hypernetwork, filename = train_hypernetwork(**args)
  591. except Exception as e:
  592. error = e
  593. finally:
  594. shared.sd_model.cond_stage_model.to(devices.device)
  595. shared.sd_model.first_stage_model.to(devices.device)
  596. if not apply_optimizations:
  597. sd_hijack.apply_optimizations()
  598. shared.state.end()
  599. return models.TrainResponse(info=f"train embedding complete: filename: {filename} error: {error}")
  600. except Exception as exc:
  601. return models.TrainResponse(info=f"train embedding error: {exc}")
  602. finally:
  603. shared.state.end()
  604. def get_memory(self):
  605. try:
  606. import os
  607. import psutil
  608. process = psutil.Process(os.getpid())
  609. res = process.memory_info() # only rss is cross-platform guaranteed so we dont rely on other values
  610. ram_total = 100 * res.rss / process.memory_percent() # and total memory is calculated as actual value is not cross-platform safe
  611. ram = { 'free': ram_total - res.rss, 'used': res.rss, 'total': ram_total }
  612. except Exception as err:
  613. ram = { 'error': f'{err}' }
  614. try:
  615. import torch
  616. if torch.cuda.is_available():
  617. s = torch.cuda.mem_get_info()
  618. system = { 'free': s[0], 'used': s[1] - s[0], 'total': s[1] }
  619. s = dict(torch.cuda.memory_stats(shared.device))
  620. allocated = { 'current': s['allocated_bytes.all.current'], 'peak': s['allocated_bytes.all.peak'] }
  621. reserved = { 'current': s['reserved_bytes.all.current'], 'peak': s['reserved_bytes.all.peak'] }
  622. active = { 'current': s['active_bytes.all.current'], 'peak': s['active_bytes.all.peak'] }
  623. inactive = { 'current': s['inactive_split_bytes.all.current'], 'peak': s['inactive_split_bytes.all.peak'] }
  624. warnings = { 'retries': s['num_alloc_retries'], 'oom': s['num_ooms'] }
  625. cuda = {
  626. 'system': system,
  627. 'active': active,
  628. 'allocated': allocated,
  629. 'reserved': reserved,
  630. 'inactive': inactive,
  631. 'events': warnings,
  632. }
  633. else:
  634. cuda = {'error': 'unavailable'}
  635. except Exception as err:
  636. cuda = {'error': f'{err}'}
  637. return models.MemoryResponse(ram=ram, cuda=cuda)
  638. def get_extensions_list(self):
  639. from modules import extensions
  640. extensions.list_extensions()
  641. ext_list = []
  642. for ext in extensions.extensions:
  643. ext: extensions.Extension
  644. ext.read_info_from_repo()
  645. if ext.remote is not None:
  646. ext_list.append({
  647. "name": ext.name,
  648. "remote": ext.remote,
  649. "branch": ext.branch,
  650. "commit_hash":ext.commit_hash,
  651. "commit_date":ext.commit_date,
  652. "version":ext.version,
  653. "enabled":ext.enabled
  654. })
  655. return ext_list
  656. def launch(self, server_name, port, root_path):
  657. self.app.include_router(self.router)
  658. uvicorn.run(self.app, host=server_name, port=port, timeout_keep_alive=shared.cmd_opts.timeout_keep_alive, root_path=root_path)
  659. def kill_webui(self):
  660. restart.stop_program()
  661. def restart_webui(self):
  662. if restart.is_restartable():
  663. restart.restart_program()
  664. return Response(status_code=501)
  665. def stop_webui(request):
  666. shared.state.server_command = "stop"
  667. return Response("Stopping.")