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
  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 import unload_model_weights, reload_model_weights, checkpoint_aliases
  28. from modules.sd_models_config import find_checkpoint_config_near_filename
  29. from modules.realesrgan_model import get_realesrgan_models
  30. from modules import devices
  31. from typing import Any
  32. import piexif
  33. import piexif.helper
  34. from contextlib import closing
  35. def script_name_to_index(name, scripts):
  36. try:
  37. return [script.title().lower() for script in scripts].index(name.lower())
  38. except Exception as e:
  39. raise HTTPException(status_code=422, detail=f"Script '{name}' not found") from e
  40. def validate_sampler_name(name):
  41. config = sd_samplers.all_samplers_map.get(name, None)
  42. if config is None:
  43. raise HTTPException(status_code=404, detail="Sampler not found")
  44. return name
  45. def setUpscalers(req: dict):
  46. reqDict = vars(req)
  47. reqDict['extras_upscaler_1'] = reqDict.pop('upscaler_1', None)
  48. reqDict['extras_upscaler_2'] = reqDict.pop('upscaler_2', None)
  49. return reqDict
  50. def verify_url(url):
  51. """Returns True if the url refers to a global resource."""
  52. import socket
  53. from urllib.parse import urlparse
  54. try:
  55. parsed_url = urlparse(url)
  56. domain_name = parsed_url.netloc
  57. host = socket.gethostbyname_ex(domain_name)
  58. for ip in host[2]:
  59. ip_addr = ipaddress.ip_address(ip)
  60. if not ip_addr.is_global:
  61. return False
  62. except Exception:
  63. return False
  64. return True
  65. def decode_base64_to_image(encoding):
  66. if encoding.startswith("http://") or encoding.startswith("https://"):
  67. if not opts.api_enable_requests:
  68. raise HTTPException(status_code=500, detail="Requests not allowed")
  69. if opts.api_forbid_local_requests and not verify_url(encoding):
  70. raise HTTPException(status_code=500, detail="Request to local resource not allowed")
  71. headers = {'user-agent': opts.api_useragent} if opts.api_useragent else {}
  72. response = requests.get(encoding, timeout=30, headers=headers)
  73. try:
  74. image = Image.open(BytesIO(response.content))
  75. return image
  76. except Exception as e:
  77. raise HTTPException(status_code=500, detail="Invalid image url") from e
  78. if encoding.startswith("data:image/"):
  79. encoding = encoding.split(";")[1].split(",")[1]
  80. try:
  81. image = Image.open(BytesIO(base64.b64decode(encoding)))
  82. return image
  83. except Exception as e:
  84. raise HTTPException(status_code=500, detail="Invalid encoded image") from e
  85. def encode_pil_to_base64(image):
  86. with io.BytesIO() as output_bytes:
  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/preprocess", self.preprocess, methods=["POST"], response_model=models.PreprocessResponse)
  205. self.add_api_route("/sdapi/v1/train/embedding", self.train_embedding, methods=["POST"], response_model=models.TrainResponse)
  206. self.add_api_route("/sdapi/v1/train/hypernetwork", self.train_hypernetwork, methods=["POST"], response_model=models.TrainResponse)
  207. self.add_api_route("/sdapi/v1/memory", self.get_memory, methods=["GET"], response_model=models.MemoryResponse)
  208. self.add_api_route("/sdapi/v1/unload-checkpoint", self.unloadapi, methods=["POST"])
  209. self.add_api_route("/sdapi/v1/reload-checkpoint", self.reloadapi, methods=["POST"])
  210. self.add_api_route("/sdapi/v1/scripts", self.get_scripts_list, methods=["GET"], response_model=models.ScriptsList)
  211. self.add_api_route("/sdapi/v1/script-info", self.get_script_info, methods=["GET"], response_model=list[models.ScriptInfo])
  212. self.add_api_route("/sdapi/v1/extensions", self.get_extensions_list, methods=["GET"], response_model=list[models.ExtensionItem])
  213. if shared.cmd_opts.api_server_stop:
  214. self.add_api_route("/sdapi/v1/server-kill", self.kill_webui, methods=["POST"])
  215. self.add_api_route("/sdapi/v1/server-restart", self.restart_webui, methods=["POST"])
  216. self.add_api_route("/sdapi/v1/server-stop", self.stop_webui, methods=["POST"])
  217. self.default_script_arg_txt2img = []
  218. self.default_script_arg_img2img = []
  219. def add_api_route(self, path: str, endpoint, **kwargs):
  220. if shared.cmd_opts.api_auth:
  221. return self.app.add_api_route(path, endpoint, dependencies=[Depends(self.auth)], **kwargs)
  222. return self.app.add_api_route(path, endpoint, **kwargs)
  223. def auth(self, credentials: HTTPBasicCredentials = Depends(HTTPBasic())):
  224. if credentials.username in self.credentials:
  225. if compare_digest(credentials.password, self.credentials[credentials.username]):
  226. return True
  227. raise HTTPException(status_code=401, detail="Incorrect username or password", headers={"WWW-Authenticate": "Basic"})
  228. def get_selectable_script(self, script_name, script_runner):
  229. if script_name is None or script_name == "":
  230. return None, None
  231. script_idx = script_name_to_index(script_name, script_runner.selectable_scripts)
  232. script = script_runner.selectable_scripts[script_idx]
  233. return script, script_idx
  234. def get_scripts_list(self):
  235. t2ilist = [script.name for script in scripts.scripts_txt2img.scripts if script.name is not None]
  236. i2ilist = [script.name for script in scripts.scripts_img2img.scripts if script.name is not None]
  237. return models.ScriptsList(txt2img=t2ilist, img2img=i2ilist)
  238. def get_script_info(self):
  239. res = []
  240. for script_list in [scripts.scripts_txt2img.scripts, scripts.scripts_img2img.scripts]:
  241. res += [script.api_info for script in script_list if script.api_info is not None]
  242. return res
  243. def get_script(self, script_name, script_runner):
  244. if script_name is None or script_name == "":
  245. return None, None
  246. script_idx = script_name_to_index(script_name, script_runner.scripts)
  247. return script_runner.scripts[script_idx]
  248. def init_default_script_args(self, script_runner):
  249. #find max idx from the scripts in runner and generate a none array to init script_args
  250. last_arg_index = 1
  251. for script in script_runner.scripts:
  252. if last_arg_index < script.args_to:
  253. last_arg_index = script.args_to
  254. # None everywhere except position 0 to initialize script args
  255. script_args = [None]*last_arg_index
  256. script_args[0] = 0
  257. # get default values
  258. with gr.Blocks(): # will throw errors calling ui function without this
  259. for script in script_runner.scripts:
  260. if script.ui(script.is_img2img):
  261. ui_default_values = []
  262. for elem in script.ui(script.is_img2img):
  263. ui_default_values.append(elem.value)
  264. script_args[script.args_from:script.args_to] = ui_default_values
  265. return script_args
  266. def init_script_args(self, request, default_script_args, selectable_scripts, selectable_idx, script_runner):
  267. script_args = default_script_args.copy()
  268. # 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()
  269. if selectable_scripts:
  270. script_args[selectable_scripts.args_from:selectable_scripts.args_to] = request.script_args
  271. script_args[0] = selectable_idx + 1
  272. # Now check for always on scripts
  273. if request.alwayson_scripts:
  274. for alwayson_script_name in request.alwayson_scripts.keys():
  275. alwayson_script = self.get_script(alwayson_script_name, script_runner)
  276. if alwayson_script is None:
  277. raise HTTPException(status_code=422, detail=f"always on script {alwayson_script_name} not found")
  278. # Selectable script in always on script param check
  279. if alwayson_script.alwayson is False:
  280. raise HTTPException(status_code=422, detail="Cannot have a selectable script in the always on scripts params")
  281. # always on script with no arg should always run so you don't really need to add them to the requests
  282. if "args" in request.alwayson_scripts[alwayson_script_name]:
  283. # min between arg length in scriptrunner and arg length in the request
  284. for idx in range(0, min((alwayson_script.args_to - alwayson_script.args_from), len(request.alwayson_scripts[alwayson_script_name]["args"]))):
  285. script_args[alwayson_script.args_from + idx] = request.alwayson_scripts[alwayson_script_name]["args"][idx]
  286. return script_args
  287. def text2imgapi(self, txt2imgreq: models.StableDiffusionTxt2ImgProcessingAPI):
  288. script_runner = scripts.scripts_txt2img
  289. if not script_runner.scripts:
  290. script_runner.initialize_scripts(False)
  291. ui.create_ui()
  292. if not self.default_script_arg_txt2img:
  293. self.default_script_arg_txt2img = self.init_default_script_args(script_runner)
  294. selectable_scripts, selectable_script_idx = self.get_selectable_script(txt2imgreq.script_name, script_runner)
  295. populate = txt2imgreq.copy(update={ # Override __init__ params
  296. "sampler_name": validate_sampler_name(txt2imgreq.sampler_name or txt2imgreq.sampler_index),
  297. "do_not_save_samples": not txt2imgreq.save_images,
  298. "do_not_save_grid": not txt2imgreq.save_images,
  299. })
  300. if populate.sampler_name:
  301. populate.sampler_index = None # prevent a warning later on
  302. args = vars(populate)
  303. args.pop('script_name', None)
  304. args.pop('script_args', None) # will refeed them to the pipeline directly after initializing them
  305. args.pop('alwayson_scripts', None)
  306. script_args = self.init_script_args(txt2imgreq, self.default_script_arg_txt2img, selectable_scripts, selectable_script_idx, script_runner)
  307. send_images = args.pop('send_images', True)
  308. args.pop('save_images', None)
  309. with self.queue_lock:
  310. with closing(StableDiffusionProcessingTxt2Img(sd_model=shared.sd_model, **args)) as p:
  311. p.is_api = True
  312. p.scripts = script_runner
  313. p.outpath_grids = opts.outdir_txt2img_grids
  314. p.outpath_samples = opts.outdir_txt2img_samples
  315. try:
  316. shared.state.begin(job="scripts_txt2img")
  317. if selectable_scripts is not None:
  318. p.script_args = script_args
  319. processed = scripts.scripts_txt2img.run(p, *p.script_args) # Need to pass args as list here
  320. else:
  321. p.script_args = tuple(script_args) # Need to pass args as tuple here
  322. processed = process_images(p)
  323. finally:
  324. shared.state.end()
  325. shared.total_tqdm.clear()
  326. b64images = list(map(encode_pil_to_base64, processed.images)) if send_images else []
  327. return models.TextToImageResponse(images=b64images, parameters=vars(txt2imgreq), info=processed.js())
  328. def img2imgapi(self, img2imgreq: models.StableDiffusionImg2ImgProcessingAPI):
  329. init_images = img2imgreq.init_images
  330. if init_images is None:
  331. raise HTTPException(status_code=404, detail="Init image not found")
  332. mask = img2imgreq.mask
  333. if mask:
  334. mask = decode_base64_to_image(mask)
  335. script_runner = scripts.scripts_img2img
  336. if not script_runner.scripts:
  337. script_runner.initialize_scripts(True)
  338. ui.create_ui()
  339. if not self.default_script_arg_img2img:
  340. self.default_script_arg_img2img = self.init_default_script_args(script_runner)
  341. selectable_scripts, selectable_script_idx = self.get_selectable_script(img2imgreq.script_name, script_runner)
  342. populate = img2imgreq.copy(update={ # Override __init__ params
  343. "sampler_name": validate_sampler_name(img2imgreq.sampler_name or img2imgreq.sampler_index),
  344. "do_not_save_samples": not img2imgreq.save_images,
  345. "do_not_save_grid": not img2imgreq.save_images,
  346. "mask": mask,
  347. })
  348. if populate.sampler_name:
  349. populate.sampler_index = None # prevent a warning later on
  350. args = vars(populate)
  351. 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.
  352. args.pop('script_name', None)
  353. args.pop('script_args', None) # will refeed them to the pipeline directly after initializing them
  354. args.pop('alwayson_scripts', None)
  355. script_args = self.init_script_args(img2imgreq, self.default_script_arg_img2img, selectable_scripts, selectable_script_idx, script_runner)
  356. send_images = args.pop('send_images', True)
  357. args.pop('save_images', None)
  358. with self.queue_lock:
  359. with closing(StableDiffusionProcessingImg2Img(sd_model=shared.sd_model, **args)) as p:
  360. p.init_images = [decode_base64_to_image(x) for x in init_images]
  361. p.is_api = True
  362. p.scripts = script_runner
  363. p.outpath_grids = opts.outdir_img2img_grids
  364. p.outpath_samples = opts.outdir_img2img_samples
  365. try:
  366. shared.state.begin(job="scripts_img2img")
  367. if selectable_scripts is not None:
  368. p.script_args = script_args
  369. processed = scripts.scripts_img2img.run(p, *p.script_args) # Need to pass args as list here
  370. else:
  371. p.script_args = tuple(script_args) # Need to pass args as tuple here
  372. processed = process_images(p)
  373. finally:
  374. shared.state.end()
  375. shared.total_tqdm.clear()
  376. b64images = list(map(encode_pil_to_base64, processed.images)) if send_images else []
  377. if not img2imgreq.include_init_images:
  378. img2imgreq.init_images = None
  379. img2imgreq.mask = None
  380. return models.ImageToImageResponse(images=b64images, parameters=vars(img2imgreq), info=processed.js())
  381. def extras_single_image_api(self, req: models.ExtrasSingleImageRequest):
  382. reqDict = setUpscalers(req)
  383. reqDict['image'] = decode_base64_to_image(reqDict['image'])
  384. with self.queue_lock:
  385. result = postprocessing.run_extras(extras_mode=0, image_folder="", input_dir="", output_dir="", save_output=False, **reqDict)
  386. return models.ExtrasSingleImageResponse(image=encode_pil_to_base64(result[0][0]), html_info=result[1])
  387. def extras_batch_images_api(self, req: models.ExtrasBatchImagesRequest):
  388. reqDict = setUpscalers(req)
  389. image_list = reqDict.pop('imageList', [])
  390. image_folder = [decode_base64_to_image(x.data) for x in image_list]
  391. with self.queue_lock:
  392. result = postprocessing.run_extras(extras_mode=1, image_folder=image_folder, image="", input_dir="", output_dir="", save_output=False, **reqDict)
  393. return models.ExtrasBatchImagesResponse(images=list(map(encode_pil_to_base64, result[0])), html_info=result[1])
  394. def pnginfoapi(self, req: models.PNGInfoRequest):
  395. image = decode_base64_to_image(req.image.strip())
  396. if image is None:
  397. return models.PNGInfoResponse(info="")
  398. geninfo, items = images.read_info_from_image(image)
  399. if geninfo is None:
  400. geninfo = ""
  401. params = generation_parameters_copypaste.parse_generation_parameters(geninfo)
  402. script_callbacks.infotext_pasted_callback(geninfo, params)
  403. return models.PNGInfoResponse(info=geninfo, items=items, parameters=params)
  404. def progressapi(self, req: models.ProgressRequest = Depends()):
  405. # copy from check_progress_call of ui.py
  406. if shared.state.job_count == 0:
  407. return models.ProgressResponse(progress=0, eta_relative=0, state=shared.state.dict(), textinfo=shared.state.textinfo)
  408. # avoid dividing zero
  409. progress = 0.01
  410. if shared.state.job_count > 0:
  411. progress += shared.state.job_no / shared.state.job_count
  412. if shared.state.sampling_steps > 0:
  413. progress += 1 / shared.state.job_count * shared.state.sampling_step / shared.state.sampling_steps
  414. time_since_start = time.time() - shared.state.time_start
  415. eta = (time_since_start/progress)
  416. eta_relative = eta-time_since_start
  417. progress = min(progress, 1)
  418. shared.state.set_current_image()
  419. current_image = None
  420. if shared.state.current_image and not req.skip_current_image:
  421. current_image = encode_pil_to_base64(shared.state.current_image)
  422. return models.ProgressResponse(progress=progress, eta_relative=eta_relative, state=shared.state.dict(), current_image=current_image, textinfo=shared.state.textinfo)
  423. def interrogateapi(self, interrogatereq: models.InterrogateRequest):
  424. image_b64 = interrogatereq.image
  425. if image_b64 is None:
  426. raise HTTPException(status_code=404, detail="Image not found")
  427. img = decode_base64_to_image(image_b64)
  428. img = img.convert('RGB')
  429. # Override object param
  430. with self.queue_lock:
  431. if interrogatereq.model == "clip":
  432. processed = shared.interrogator.interrogate(img)
  433. elif interrogatereq.model == "deepdanbooru":
  434. processed = deepbooru.model.tag(img)
  435. else:
  436. raise HTTPException(status_code=404, detail="Model not found")
  437. return models.InterrogateResponse(caption=processed)
  438. def interruptapi(self):
  439. shared.state.interrupt()
  440. return {}
  441. def unloadapi(self):
  442. unload_model_weights()
  443. return {}
  444. def reloadapi(self):
  445. reload_model_weights()
  446. return {}
  447. def skip(self):
  448. shared.state.skip()
  449. def get_config(self):
  450. options = {}
  451. for key in shared.opts.data.keys():
  452. metadata = shared.opts.data_labels.get(key)
  453. if(metadata is not None):
  454. options.update({key: shared.opts.data.get(key, shared.opts.data_labels.get(key).default)})
  455. else:
  456. options.update({key: shared.opts.data.get(key, None)})
  457. return options
  458. def set_config(self, req: dict[str, Any]):
  459. checkpoint_name = req.get("sd_model_checkpoint", None)
  460. if checkpoint_name is not None and checkpoint_name not in checkpoint_aliases:
  461. raise RuntimeError(f"model {checkpoint_name!r} not found")
  462. for k, v in req.items():
  463. shared.opts.set(k, v, is_api=True)
  464. shared.opts.save(shared.config_filename)
  465. return
  466. def get_cmd_flags(self):
  467. return vars(shared.cmd_opts)
  468. def get_samplers(self):
  469. return [{"name": sampler[0], "aliases":sampler[2], "options":sampler[3]} for sampler in sd_samplers.all_samplers]
  470. def get_upscalers(self):
  471. return [
  472. {
  473. "name": upscaler.name,
  474. "model_name": upscaler.scaler.model_name,
  475. "model_path": upscaler.data_path,
  476. "model_url": None,
  477. "scale": upscaler.scale,
  478. }
  479. for upscaler in shared.sd_upscalers
  480. ]
  481. def get_latent_upscale_modes(self):
  482. return [
  483. {
  484. "name": upscale_mode,
  485. }
  486. for upscale_mode in [*(shared.latent_upscale_modes or {})]
  487. ]
  488. def get_sd_models(self):
  489. import modules.sd_models as sd_models
  490. 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()]
  491. def get_sd_vaes(self):
  492. import modules.sd_vae as sd_vae
  493. return [{"model_name": x, "filename": sd_vae.vae_dict[x]} for x in sd_vae.vae_dict.keys()]
  494. def get_hypernetworks(self):
  495. return [{"name": name, "path": shared.hypernetworks[name]} for name in shared.hypernetworks]
  496. def get_face_restorers(self):
  497. return [{"name":x.name(), "cmd_dir": getattr(x, "cmd_dir", None)} for x in shared.face_restorers]
  498. def get_realesrgan_models(self):
  499. return [{"name":x.name,"path":x.data_path, "scale":x.scale} for x in get_realesrgan_models(None)]
  500. def get_prompt_styles(self):
  501. styleList = []
  502. for k in shared.prompt_styles.styles:
  503. style = shared.prompt_styles.styles[k]
  504. styleList.append({"name":style[0], "prompt": style[1], "negative_prompt": style[2]})
  505. return styleList
  506. def get_embeddings(self):
  507. db = sd_hijack.model_hijack.embedding_db
  508. def convert_embedding(embedding):
  509. return {
  510. "step": embedding.step,
  511. "sd_checkpoint": embedding.sd_checkpoint,
  512. "sd_checkpoint_name": embedding.sd_checkpoint_name,
  513. "shape": embedding.shape,
  514. "vectors": embedding.vectors,
  515. }
  516. def convert_embeddings(embeddings):
  517. return {embedding.name: convert_embedding(embedding) for embedding in embeddings.values()}
  518. return {
  519. "loaded": convert_embeddings(db.word_embeddings),
  520. "skipped": convert_embeddings(db.skipped_embeddings),
  521. }
  522. def refresh_checkpoints(self):
  523. with self.queue_lock:
  524. shared.refresh_checkpoints()
  525. def refresh_vae(self):
  526. with self.queue_lock:
  527. shared_items.refresh_vae_list()
  528. def create_embedding(self, args: dict):
  529. try:
  530. shared.state.begin(job="create_embedding")
  531. filename = create_embedding(**args) # create empty embedding
  532. sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings() # reload embeddings so new one can be immediately used
  533. return models.CreateResponse(info=f"create embedding filename: {filename}")
  534. except AssertionError as e:
  535. return models.TrainResponse(info=f"create embedding error: {e}")
  536. finally:
  537. shared.state.end()
  538. def create_hypernetwork(self, args: dict):
  539. try:
  540. shared.state.begin(job="create_hypernetwork")
  541. filename = create_hypernetwork(**args) # create empty embedding
  542. return models.CreateResponse(info=f"create hypernetwork filename: {filename}")
  543. except AssertionError as e:
  544. return models.TrainResponse(info=f"create hypernetwork error: {e}")
  545. finally:
  546. shared.state.end()
  547. def preprocess(self, args: dict):
  548. try:
  549. shared.state.begin(job="preprocess")
  550. preprocess(**args) # quick operation unless blip/booru interrogation is enabled
  551. shared.state.end()
  552. return models.PreprocessResponse(info='preprocess complete')
  553. except KeyError as e:
  554. return models.PreprocessResponse(info=f"preprocess error: invalid token: {e}")
  555. except Exception as e:
  556. return models.PreprocessResponse(info=f"preprocess error: {e}")
  557. finally:
  558. shared.state.end()
  559. def train_embedding(self, args: dict):
  560. try:
  561. shared.state.begin(job="train_embedding")
  562. apply_optimizations = shared.opts.training_xattention_optimizations
  563. error = None
  564. filename = ''
  565. if not apply_optimizations:
  566. sd_hijack.undo_optimizations()
  567. try:
  568. embedding, filename = train_embedding(**args) # can take a long time to complete
  569. except Exception as e:
  570. error = e
  571. finally:
  572. if not apply_optimizations:
  573. sd_hijack.apply_optimizations()
  574. return models.TrainResponse(info=f"train embedding complete: filename: {filename} error: {error}")
  575. except Exception as msg:
  576. return models.TrainResponse(info=f"train embedding error: {msg}")
  577. finally:
  578. shared.state.end()
  579. def train_hypernetwork(self, args: dict):
  580. try:
  581. shared.state.begin(job="train_hypernetwork")
  582. shared.loaded_hypernetworks = []
  583. apply_optimizations = shared.opts.training_xattention_optimizations
  584. error = None
  585. filename = ''
  586. if not apply_optimizations:
  587. sd_hijack.undo_optimizations()
  588. try:
  589. hypernetwork, filename = train_hypernetwork(**args)
  590. except Exception as e:
  591. error = e
  592. finally:
  593. shared.sd_model.cond_stage_model.to(devices.device)
  594. shared.sd_model.first_stage_model.to(devices.device)
  595. if not apply_optimizations:
  596. sd_hijack.apply_optimizations()
  597. shared.state.end()
  598. return models.TrainResponse(info=f"train embedding complete: filename: {filename} error: {error}")
  599. except Exception as exc:
  600. return models.TrainResponse(info=f"train embedding error: {exc}")
  601. finally:
  602. shared.state.end()
  603. def get_memory(self):
  604. try:
  605. import os
  606. import psutil
  607. process = psutil.Process(os.getpid())
  608. res = process.memory_info() # only rss is cross-platform guaranteed so we dont rely on other values
  609. ram_total = 100 * res.rss / process.memory_percent() # and total memory is calculated as actual value is not cross-platform safe
  610. ram = { 'free': ram_total - res.rss, 'used': res.rss, 'total': ram_total }
  611. except Exception as err:
  612. ram = { 'error': f'{err}' }
  613. try:
  614. import torch
  615. if torch.cuda.is_available():
  616. s = torch.cuda.mem_get_info()
  617. system = { 'free': s[0], 'used': s[1] - s[0], 'total': s[1] }
  618. s = dict(torch.cuda.memory_stats(shared.device))
  619. allocated = { 'current': s['allocated_bytes.all.current'], 'peak': s['allocated_bytes.all.peak'] }
  620. reserved = { 'current': s['reserved_bytes.all.current'], 'peak': s['reserved_bytes.all.peak'] }
  621. active = { 'current': s['active_bytes.all.current'], 'peak': s['active_bytes.all.peak'] }
  622. inactive = { 'current': s['inactive_split_bytes.all.current'], 'peak': s['inactive_split_bytes.all.peak'] }
  623. warnings = { 'retries': s['num_alloc_retries'], 'oom': s['num_ooms'] }
  624. cuda = {
  625. 'system': system,
  626. 'active': active,
  627. 'allocated': allocated,
  628. 'reserved': reserved,
  629. 'inactive': inactive,
  630. 'events': warnings,
  631. }
  632. else:
  633. cuda = {'error': 'unavailable'}
  634. except Exception as err:
  635. cuda = {'error': f'{err}'}
  636. return models.MemoryResponse(ram=ram, cuda=cuda)
  637. def get_extensions_list(self):
  638. from modules import extensions
  639. extensions.list_extensions()
  640. ext_list = []
  641. for ext in extensions.extensions:
  642. ext: extensions.Extension
  643. ext.read_info_from_repo()
  644. if ext.remote is not None:
  645. ext_list.append({
  646. "name": ext.name,
  647. "remote": ext.remote,
  648. "branch": ext.branch,
  649. "commit_hash":ext.commit_hash,
  650. "commit_date":ext.commit_date,
  651. "version":ext.version,
  652. "enabled":ext.enabled
  653. })
  654. return ext_list
  655. def launch(self, server_name, port, root_path):
  656. self.app.include_router(self.router)
  657. uvicorn.run(self.app, host=server_name, port=port, timeout_keep_alive=shared.cmd_opts.timeout_keep_alive, root_path=root_path)
  658. def kill_webui(self):
  659. restart.stop_program()
  660. def restart_webui(self):
  661. if restart.is_restartable():
  662. restart.restart_program()
  663. return Response(status_code=501)
  664. def stop_webui(request):
  665. shared.state.server_command = "stop"
  666. return Response("Stopping.")