api.py 36 KB

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