api.py 40 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908
  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, infotext_utils, 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. from modules.progress import create_task_id, add_task_to_queue, start_task, finish_task, current_task
  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-embeddings", self.refresh_embeddings, methods=["POST"])
  202. self.add_api_route("/sdapi/v1/refresh-checkpoints", self.refresh_checkpoints, methods=["POST"])
  203. self.add_api_route("/sdapi/v1/refresh-vae", self.refresh_vae, methods=["POST"])
  204. self.add_api_route("/sdapi/v1/create/embedding", self.create_embedding, methods=["POST"], response_model=models.CreateResponse)
  205. self.add_api_route("/sdapi/v1/create/hypernetwork", self.create_hypernetwork, methods=["POST"], response_model=models.CreateResponse)
  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. txt2img_script_runner = scripts.scripts_txt2img
  221. img2img_script_runner = scripts.scripts_img2img
  222. if not txt2img_script_runner.scripts or not img2img_script_runner.scripts:
  223. ui.create_ui()
  224. if not txt2img_script_runner.scripts:
  225. txt2img_script_runner.initialize_scripts(False)
  226. if not self.default_script_arg_txt2img:
  227. self.default_script_arg_txt2img = self.init_default_script_args(txt2img_script_runner)
  228. if not img2img_script_runner.scripts:
  229. img2img_script_runner.initialize_scripts(True)
  230. if not self.default_script_arg_img2img:
  231. self.default_script_arg_img2img = self.init_default_script_args(img2img_script_runner)
  232. def add_api_route(self, path: str, endpoint, **kwargs):
  233. if shared.cmd_opts.api_auth:
  234. return self.app.add_api_route(path, endpoint, dependencies=[Depends(self.auth)], **kwargs)
  235. return self.app.add_api_route(path, endpoint, **kwargs)
  236. def auth(self, credentials: HTTPBasicCredentials = Depends(HTTPBasic())):
  237. if credentials.username in self.credentials:
  238. if compare_digest(credentials.password, self.credentials[credentials.username]):
  239. return True
  240. raise HTTPException(status_code=401, detail="Incorrect username or password", headers={"WWW-Authenticate": "Basic"})
  241. def get_selectable_script(self, script_name, script_runner):
  242. if script_name is None or script_name == "":
  243. return None, None
  244. script_idx = script_name_to_index(script_name, script_runner.selectable_scripts)
  245. script = script_runner.selectable_scripts[script_idx]
  246. return script, script_idx
  247. def get_scripts_list(self):
  248. t2ilist = [script.name for script in scripts.scripts_txt2img.scripts if script.name is not None]
  249. i2ilist = [script.name for script in scripts.scripts_img2img.scripts if script.name is not None]
  250. return models.ScriptsList(txt2img=t2ilist, img2img=i2ilist)
  251. def get_script_info(self):
  252. res = []
  253. for script_list in [scripts.scripts_txt2img.scripts, scripts.scripts_img2img.scripts]:
  254. res += [script.api_info for script in script_list if script.api_info is not None]
  255. return res
  256. def get_script(self, script_name, script_runner):
  257. if script_name is None or script_name == "":
  258. return None, None
  259. script_idx = script_name_to_index(script_name, script_runner.scripts)
  260. return script_runner.scripts[script_idx]
  261. def init_default_script_args(self, script_runner):
  262. #find max idx from the scripts in runner and generate a none array to init script_args
  263. last_arg_index = 1
  264. for script in script_runner.scripts:
  265. if last_arg_index < script.args_to:
  266. last_arg_index = script.args_to
  267. # None everywhere except position 0 to initialize script args
  268. script_args = [None]*last_arg_index
  269. script_args[0] = 0
  270. # get default values
  271. with gr.Blocks(): # will throw errors calling ui function without this
  272. for script in script_runner.scripts:
  273. if script.ui(script.is_img2img):
  274. ui_default_values = []
  275. for elem in script.ui(script.is_img2img):
  276. ui_default_values.append(elem.value)
  277. script_args[script.args_from:script.args_to] = ui_default_values
  278. return script_args
  279. def init_script_args(self, request, default_script_args, selectable_scripts, selectable_idx, script_runner, *, input_script_args=None):
  280. script_args = default_script_args.copy()
  281. if input_script_args is not None:
  282. for index, value in input_script_args.items():
  283. script_args[index] = value
  284. # 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()
  285. if selectable_scripts:
  286. script_args[selectable_scripts.args_from:selectable_scripts.args_to] = request.script_args
  287. script_args[0] = selectable_idx + 1
  288. # Now check for always on scripts
  289. if request.alwayson_scripts:
  290. for alwayson_script_name in request.alwayson_scripts.keys():
  291. alwayson_script = self.get_script(alwayson_script_name, script_runner)
  292. if alwayson_script is None:
  293. raise HTTPException(status_code=422, detail=f"always on script {alwayson_script_name} not found")
  294. # Selectable script in always on script param check
  295. if alwayson_script.alwayson is False:
  296. raise HTTPException(status_code=422, detail="Cannot have a selectable script in the always on scripts params")
  297. # always on script with no arg should always run so you don't really need to add them to the requests
  298. if "args" in request.alwayson_scripts[alwayson_script_name]:
  299. # min between arg length in scriptrunner and arg length in the request
  300. for idx in range(0, min((alwayson_script.args_to - alwayson_script.args_from), len(request.alwayson_scripts[alwayson_script_name]["args"]))):
  301. script_args[alwayson_script.args_from + idx] = request.alwayson_scripts[alwayson_script_name]["args"][idx]
  302. return script_args
  303. def apply_infotext(self, request, tabname, *, script_runner=None, mentioned_script_args=None):
  304. """Processes `infotext` field from the `request`, and sets other fields of the `request` according to what's in infotext.
  305. If request already has a field set, and that field is encountered in infotext too, the value from infotext is ignored.
  306. Additionally, fills `mentioned_script_args` dict with index: value pairs for script arguments read from infotext.
  307. """
  308. if not request.infotext:
  309. return {}
  310. possible_fields = infotext_utils.paste_fields[tabname]["fields"]
  311. set_fields = request.model_dump(exclude_unset=True) if hasattr(request, "request") else request.dict(exclude_unset=True) # pydantic v1/v2 have differenrt names for this
  312. params = infotext_utils.parse_generation_parameters(request.infotext)
  313. def get_field_value(field, params):
  314. value = field.function(params) if field.function else params.get(field.label)
  315. if value is None:
  316. return None
  317. if field.api in request.__fields__:
  318. target_type = request.__fields__[field.api].type_
  319. else:
  320. target_type = type(field.component.value)
  321. if target_type == type(None):
  322. return None
  323. if isinstance(value, dict) and value.get('__type__') == 'generic_update': # this is a gradio.update rather than a value
  324. value = value.get('value')
  325. if value is not None and not isinstance(value, target_type):
  326. value = target_type(value)
  327. return value
  328. for field in possible_fields:
  329. if not field.api:
  330. continue
  331. if field.api in set_fields:
  332. continue
  333. value = get_field_value(field, params)
  334. if value is not None:
  335. setattr(request, field.api, value)
  336. if request.override_settings is None:
  337. request.override_settings = {}
  338. overridden_settings = infotext_utils.get_override_settings(params)
  339. for _, setting_name, value in overridden_settings:
  340. if setting_name not in request.override_settings:
  341. request.override_settings[setting_name] = value
  342. if script_runner is not None and mentioned_script_args is not None:
  343. indexes = {v: i for i, v in enumerate(script_runner.inputs)}
  344. script_fields = ((field, indexes[field.component]) for field in possible_fields if field.component in indexes)
  345. for field, index in script_fields:
  346. value = get_field_value(field, params)
  347. if value is None:
  348. continue
  349. mentioned_script_args[index] = value
  350. return params
  351. def text2imgapi(self, txt2imgreq: models.StableDiffusionTxt2ImgProcessingAPI):
  352. task_id = txt2imgreq.force_task_id or create_task_id("txt2img")
  353. script_runner = scripts.scripts_txt2img
  354. infotext_script_args = {}
  355. self.apply_infotext(txt2imgreq, "txt2img", script_runner=script_runner, mentioned_script_args=infotext_script_args)
  356. selectable_scripts, selectable_script_idx = self.get_selectable_script(txt2imgreq.script_name, script_runner)
  357. populate = txt2imgreq.copy(update={ # Override __init__ params
  358. "sampler_name": validate_sampler_name(txt2imgreq.sampler_name or txt2imgreq.sampler_index),
  359. "do_not_save_samples": not txt2imgreq.save_images,
  360. "do_not_save_grid": not txt2imgreq.save_images,
  361. })
  362. if populate.sampler_name:
  363. populate.sampler_index = None # prevent a warning later on
  364. args = vars(populate)
  365. args.pop('script_name', None)
  366. args.pop('script_args', None) # will refeed them to the pipeline directly after initializing them
  367. args.pop('alwayson_scripts', None)
  368. args.pop('infotext', None)
  369. script_args = self.init_script_args(txt2imgreq, self.default_script_arg_txt2img, selectable_scripts, selectable_script_idx, script_runner, input_script_args=infotext_script_args)
  370. send_images = args.pop('send_images', True)
  371. args.pop('save_images', None)
  372. add_task_to_queue(task_id)
  373. with self.queue_lock:
  374. with closing(StableDiffusionProcessingTxt2Img(sd_model=shared.sd_model, **args)) as p:
  375. p.is_api = True
  376. p.scripts = script_runner
  377. p.outpath_grids = opts.outdir_txt2img_grids
  378. p.outpath_samples = opts.outdir_txt2img_samples
  379. try:
  380. shared.state.begin(job="scripts_txt2img")
  381. start_task(task_id)
  382. if selectable_scripts is not None:
  383. p.script_args = script_args
  384. processed = scripts.scripts_txt2img.run(p, *p.script_args) # Need to pass args as list here
  385. else:
  386. p.script_args = tuple(script_args) # Need to pass args as tuple here
  387. processed = process_images(p)
  388. finish_task(task_id)
  389. finally:
  390. shared.state.end()
  391. shared.total_tqdm.clear()
  392. b64images = list(map(encode_pil_to_base64, processed.images)) if send_images else []
  393. return models.TextToImageResponse(images=b64images, parameters=vars(txt2imgreq), info=processed.js())
  394. def img2imgapi(self, img2imgreq: models.StableDiffusionImg2ImgProcessingAPI):
  395. task_id = img2imgreq.force_task_id or create_task_id("img2img")
  396. init_images = img2imgreq.init_images
  397. if init_images is None:
  398. raise HTTPException(status_code=404, detail="Init image not found")
  399. mask = img2imgreq.mask
  400. if mask:
  401. mask = decode_base64_to_image(mask)
  402. script_runner = scripts.scripts_img2img
  403. infotext_script_args = {}
  404. self.apply_infotext(img2imgreq, "img2img", script_runner=script_runner, mentioned_script_args=infotext_script_args)
  405. selectable_scripts, selectable_script_idx = self.get_selectable_script(img2imgreq.script_name, script_runner)
  406. populate = img2imgreq.copy(update={ # Override __init__ params
  407. "sampler_name": validate_sampler_name(img2imgreq.sampler_name or img2imgreq.sampler_index),
  408. "do_not_save_samples": not img2imgreq.save_images,
  409. "do_not_save_grid": not img2imgreq.save_images,
  410. "mask": mask,
  411. })
  412. if populate.sampler_name:
  413. populate.sampler_index = None # prevent a warning later on
  414. args = vars(populate)
  415. 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.
  416. args.pop('script_name', None)
  417. args.pop('script_args', None) # will refeed them to the pipeline directly after initializing them
  418. args.pop('alwayson_scripts', None)
  419. args.pop('infotext', None)
  420. script_args = self.init_script_args(img2imgreq, self.default_script_arg_img2img, selectable_scripts, selectable_script_idx, script_runner, input_script_args=infotext_script_args)
  421. send_images = args.pop('send_images', True)
  422. args.pop('save_images', None)
  423. add_task_to_queue(task_id)
  424. with self.queue_lock:
  425. with closing(StableDiffusionProcessingImg2Img(sd_model=shared.sd_model, **args)) as p:
  426. p.init_images = [decode_base64_to_image(x) for x in init_images]
  427. p.is_api = True
  428. p.scripts = script_runner
  429. p.outpath_grids = opts.outdir_img2img_grids
  430. p.outpath_samples = opts.outdir_img2img_samples
  431. try:
  432. shared.state.begin(job="scripts_img2img")
  433. start_task(task_id)
  434. if selectable_scripts is not None:
  435. p.script_args = script_args
  436. processed = scripts.scripts_img2img.run(p, *p.script_args) # Need to pass args as list here
  437. else:
  438. p.script_args = tuple(script_args) # Need to pass args as tuple here
  439. processed = process_images(p)
  440. finish_task(task_id)
  441. finally:
  442. shared.state.end()
  443. shared.total_tqdm.clear()
  444. b64images = list(map(encode_pil_to_base64, processed.images)) if send_images else []
  445. if not img2imgreq.include_init_images:
  446. img2imgreq.init_images = None
  447. img2imgreq.mask = None
  448. return models.ImageToImageResponse(images=b64images, parameters=vars(img2imgreq), info=processed.js())
  449. def extras_single_image_api(self, req: models.ExtrasSingleImageRequest):
  450. reqDict = setUpscalers(req)
  451. reqDict['image'] = decode_base64_to_image(reqDict['image'])
  452. with self.queue_lock:
  453. result = postprocessing.run_extras(extras_mode=0, image_folder="", input_dir="", output_dir="", save_output=False, **reqDict)
  454. return models.ExtrasSingleImageResponse(image=encode_pil_to_base64(result[0][0]), html_info=result[1])
  455. def extras_batch_images_api(self, req: models.ExtrasBatchImagesRequest):
  456. reqDict = setUpscalers(req)
  457. image_list = reqDict.pop('imageList', [])
  458. image_folder = [decode_base64_to_image(x.data) for x in image_list]
  459. with self.queue_lock:
  460. result = postprocessing.run_extras(extras_mode=1, image_folder=image_folder, image="", input_dir="", output_dir="", save_output=False, **reqDict)
  461. return models.ExtrasBatchImagesResponse(images=list(map(encode_pil_to_base64, result[0])), html_info=result[1])
  462. def pnginfoapi(self, req: models.PNGInfoRequest):
  463. image = decode_base64_to_image(req.image.strip())
  464. if image is None:
  465. return models.PNGInfoResponse(info="")
  466. geninfo, items = images.read_info_from_image(image)
  467. if geninfo is None:
  468. geninfo = ""
  469. params = infotext_utils.parse_generation_parameters(geninfo)
  470. script_callbacks.infotext_pasted_callback(geninfo, params)
  471. return models.PNGInfoResponse(info=geninfo, items=items, parameters=params)
  472. def progressapi(self, req: models.ProgressRequest = Depends()):
  473. # copy from check_progress_call of ui.py
  474. if shared.state.job_count == 0:
  475. return models.ProgressResponse(progress=0, eta_relative=0, state=shared.state.dict(), textinfo=shared.state.textinfo)
  476. # avoid dividing zero
  477. progress = 0.01
  478. if shared.state.job_count > 0:
  479. progress += shared.state.job_no / shared.state.job_count
  480. if shared.state.sampling_steps > 0:
  481. progress += 1 / shared.state.job_count * shared.state.sampling_step / shared.state.sampling_steps
  482. time_since_start = time.time() - shared.state.time_start
  483. eta = (time_since_start/progress)
  484. eta_relative = eta-time_since_start
  485. progress = min(progress, 1)
  486. shared.state.set_current_image()
  487. current_image = None
  488. if shared.state.current_image and not req.skip_current_image:
  489. current_image = encode_pil_to_base64(shared.state.current_image)
  490. return models.ProgressResponse(progress=progress, eta_relative=eta_relative, state=shared.state.dict(), current_image=current_image, textinfo=shared.state.textinfo, current_task=current_task)
  491. def interrogateapi(self, interrogatereq: models.InterrogateRequest):
  492. image_b64 = interrogatereq.image
  493. if image_b64 is None:
  494. raise HTTPException(status_code=404, detail="Image not found")
  495. img = decode_base64_to_image(image_b64)
  496. img = img.convert('RGB')
  497. # Override object param
  498. with self.queue_lock:
  499. if interrogatereq.model == "clip":
  500. processed = shared.interrogator.interrogate(img)
  501. elif interrogatereq.model == "deepdanbooru":
  502. processed = deepbooru.model.tag(img)
  503. else:
  504. raise HTTPException(status_code=404, detail="Model not found")
  505. return models.InterrogateResponse(caption=processed)
  506. def interruptapi(self):
  507. shared.state.interrupt()
  508. return {}
  509. def unloadapi(self):
  510. sd_models.unload_model_weights()
  511. return {}
  512. def reloadapi(self):
  513. sd_models.send_model_to_device(shared.sd_model)
  514. return {}
  515. def skip(self):
  516. shared.state.skip()
  517. def get_config(self):
  518. options = {}
  519. for key in shared.opts.data.keys():
  520. metadata = shared.opts.data_labels.get(key)
  521. if(metadata is not None):
  522. options.update({key: shared.opts.data.get(key, shared.opts.data_labels.get(key).default)})
  523. else:
  524. options.update({key: shared.opts.data.get(key, None)})
  525. return options
  526. def set_config(self, req: dict[str, Any]):
  527. checkpoint_name = req.get("sd_model_checkpoint", None)
  528. if checkpoint_name is not None and checkpoint_name not in sd_models.checkpoint_aliases:
  529. raise RuntimeError(f"model {checkpoint_name!r} not found")
  530. for k, v in req.items():
  531. shared.opts.set(k, v, is_api=True)
  532. shared.opts.save(shared.config_filename)
  533. return
  534. def get_cmd_flags(self):
  535. return vars(shared.cmd_opts)
  536. def get_samplers(self):
  537. return [{"name": sampler[0], "aliases":sampler[2], "options":sampler[3]} for sampler in sd_samplers.all_samplers]
  538. def get_upscalers(self):
  539. return [
  540. {
  541. "name": upscaler.name,
  542. "model_name": upscaler.scaler.model_name,
  543. "model_path": upscaler.data_path,
  544. "model_url": None,
  545. "scale": upscaler.scale,
  546. }
  547. for upscaler in shared.sd_upscalers
  548. ]
  549. def get_latent_upscale_modes(self):
  550. return [
  551. {
  552. "name": upscale_mode,
  553. }
  554. for upscale_mode in [*(shared.latent_upscale_modes or {})]
  555. ]
  556. def get_sd_models(self):
  557. import modules.sd_models as sd_models
  558. 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()]
  559. def get_sd_vaes(self):
  560. import modules.sd_vae as sd_vae
  561. return [{"model_name": x, "filename": sd_vae.vae_dict[x]} for x in sd_vae.vae_dict.keys()]
  562. def get_hypernetworks(self):
  563. return [{"name": name, "path": shared.hypernetworks[name]} for name in shared.hypernetworks]
  564. def get_face_restorers(self):
  565. return [{"name":x.name(), "cmd_dir": getattr(x, "cmd_dir", None)} for x in shared.face_restorers]
  566. def get_realesrgan_models(self):
  567. return [{"name":x.name,"path":x.data_path, "scale":x.scale} for x in get_realesrgan_models(None)]
  568. def get_prompt_styles(self):
  569. styleList = []
  570. for k in shared.prompt_styles.styles:
  571. style = shared.prompt_styles.styles[k]
  572. styleList.append({"name":style[0], "prompt": style[1], "negative_prompt": style[2]})
  573. return styleList
  574. def get_embeddings(self):
  575. db = sd_hijack.model_hijack.embedding_db
  576. def convert_embedding(embedding):
  577. return {
  578. "step": embedding.step,
  579. "sd_checkpoint": embedding.sd_checkpoint,
  580. "sd_checkpoint_name": embedding.sd_checkpoint_name,
  581. "shape": embedding.shape,
  582. "vectors": embedding.vectors,
  583. }
  584. def convert_embeddings(embeddings):
  585. return {embedding.name: convert_embedding(embedding) for embedding in embeddings.values()}
  586. return {
  587. "loaded": convert_embeddings(db.word_embeddings),
  588. "skipped": convert_embeddings(db.skipped_embeddings),
  589. }
  590. def refresh_embeddings(self):
  591. with self.queue_lock:
  592. sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings(force_reload=True)
  593. def refresh_checkpoints(self):
  594. with self.queue_lock:
  595. shared.refresh_checkpoints()
  596. def refresh_vae(self):
  597. with self.queue_lock:
  598. shared_items.refresh_vae_list()
  599. def create_embedding(self, args: dict):
  600. try:
  601. shared.state.begin(job="create_embedding")
  602. filename = create_embedding(**args) # create empty embedding
  603. sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings() # reload embeddings so new one can be immediately used
  604. return models.CreateResponse(info=f"create embedding filename: {filename}")
  605. except AssertionError as e:
  606. return models.TrainResponse(info=f"create embedding error: {e}")
  607. finally:
  608. shared.state.end()
  609. def create_hypernetwork(self, args: dict):
  610. try:
  611. shared.state.begin(job="create_hypernetwork")
  612. filename = create_hypernetwork(**args) # create empty embedding
  613. return models.CreateResponse(info=f"create hypernetwork filename: {filename}")
  614. except AssertionError as e:
  615. return models.TrainResponse(info=f"create hypernetwork error: {e}")
  616. finally:
  617. shared.state.end()
  618. def train_embedding(self, args: dict):
  619. try:
  620. shared.state.begin(job="train_embedding")
  621. apply_optimizations = shared.opts.training_xattention_optimizations
  622. error = None
  623. filename = ''
  624. if not apply_optimizations:
  625. sd_hijack.undo_optimizations()
  626. try:
  627. embedding, filename = train_embedding(**args) # can take a long time to complete
  628. except Exception as e:
  629. error = e
  630. finally:
  631. if not apply_optimizations:
  632. sd_hijack.apply_optimizations()
  633. return models.TrainResponse(info=f"train embedding complete: filename: {filename} error: {error}")
  634. except Exception as msg:
  635. return models.TrainResponse(info=f"train embedding error: {msg}")
  636. finally:
  637. shared.state.end()
  638. def train_hypernetwork(self, args: dict):
  639. try:
  640. shared.state.begin(job="train_hypernetwork")
  641. shared.loaded_hypernetworks = []
  642. apply_optimizations = shared.opts.training_xattention_optimizations
  643. error = None
  644. filename = ''
  645. if not apply_optimizations:
  646. sd_hijack.undo_optimizations()
  647. try:
  648. hypernetwork, filename = train_hypernetwork(**args)
  649. except Exception as e:
  650. error = e
  651. finally:
  652. shared.sd_model.cond_stage_model.to(devices.device)
  653. shared.sd_model.first_stage_model.to(devices.device)
  654. if not apply_optimizations:
  655. sd_hijack.apply_optimizations()
  656. shared.state.end()
  657. return models.TrainResponse(info=f"train embedding complete: filename: {filename} error: {error}")
  658. except Exception as exc:
  659. return models.TrainResponse(info=f"train embedding error: {exc}")
  660. finally:
  661. shared.state.end()
  662. def get_memory(self):
  663. try:
  664. import os
  665. import psutil
  666. process = psutil.Process(os.getpid())
  667. res = process.memory_info() # only rss is cross-platform guaranteed so we dont rely on other values
  668. ram_total = 100 * res.rss / process.memory_percent() # and total memory is calculated as actual value is not cross-platform safe
  669. ram = { 'free': ram_total - res.rss, 'used': res.rss, 'total': ram_total }
  670. except Exception as err:
  671. ram = { 'error': f'{err}' }
  672. try:
  673. import torch
  674. if torch.cuda.is_available():
  675. s = torch.cuda.mem_get_info()
  676. system = { 'free': s[0], 'used': s[1] - s[0], 'total': s[1] }
  677. s = dict(torch.cuda.memory_stats(shared.device))
  678. allocated = { 'current': s['allocated_bytes.all.current'], 'peak': s['allocated_bytes.all.peak'] }
  679. reserved = { 'current': s['reserved_bytes.all.current'], 'peak': s['reserved_bytes.all.peak'] }
  680. active = { 'current': s['active_bytes.all.current'], 'peak': s['active_bytes.all.peak'] }
  681. inactive = { 'current': s['inactive_split_bytes.all.current'], 'peak': s['inactive_split_bytes.all.peak'] }
  682. warnings = { 'retries': s['num_alloc_retries'], 'oom': s['num_ooms'] }
  683. cuda = {
  684. 'system': system,
  685. 'active': active,
  686. 'allocated': allocated,
  687. 'reserved': reserved,
  688. 'inactive': inactive,
  689. 'events': warnings,
  690. }
  691. else:
  692. cuda = {'error': 'unavailable'}
  693. except Exception as err:
  694. cuda = {'error': f'{err}'}
  695. return models.MemoryResponse(ram=ram, cuda=cuda)
  696. def get_extensions_list(self):
  697. from modules import extensions
  698. extensions.list_extensions()
  699. ext_list = []
  700. for ext in extensions.extensions:
  701. ext: extensions.Extension
  702. ext.read_info_from_repo()
  703. if ext.remote is not None:
  704. ext_list.append({
  705. "name": ext.name,
  706. "remote": ext.remote,
  707. "branch": ext.branch,
  708. "commit_hash":ext.commit_hash,
  709. "commit_date":ext.commit_date,
  710. "version":ext.version,
  711. "enabled":ext.enabled
  712. })
  713. return ext_list
  714. def launch(self, server_name, port, root_path):
  715. self.app.include_router(self.router)
  716. uvicorn.run(
  717. self.app,
  718. host=server_name,
  719. port=port,
  720. timeout_keep_alive=shared.cmd_opts.timeout_keep_alive,
  721. root_path=root_path,
  722. ssl_keyfile=shared.cmd_opts.tls_keyfile,
  723. ssl_certfile=shared.cmd_opts.tls_certfile
  724. )
  725. def kill_webui(self):
  726. restart.stop_program()
  727. def restart_webui(self):
  728. if restart.is_restartable():
  729. restart.restart_program()
  730. return Response(status_code=501)
  731. def stop_webui(request):
  732. shared.state.server_command = "stop"
  733. return Response("Stopping.")