diff options
Diffstat (limited to 'modules/api')
| -rw-r--r-- | modules/api/api.py | 58 | ||||
| -rw-r--r-- | modules/api/models.py | 9 |
2 files changed, 57 insertions, 10 deletions
diff --git a/modules/api/api.py b/modules/api/api.py index 606db179..e6edffe7 100644 --- a/modules/api/api.py +++ b/modules/api/api.py @@ -4,6 +4,8 @@ import os import time import datetime import uvicorn +import ipaddress +import requests import gradio as gr from threading import Lock from io import BytesIO @@ -15,7 +17,7 @@ from fastapi.encoders import jsonable_encoder from secrets import compare_digest import modules.shared as shared -from modules import sd_samplers, deepbooru, sd_hijack, images, scripts, ui, postprocessing, errors, restart +from modules import sd_samplers, deepbooru, sd_hijack, images, scripts, ui, postprocessing, errors, restart, shared_items from modules.api import models from modules.shared import opts from modules.processing import StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img, process_images @@ -23,8 +25,7 @@ from modules.textual_inversion.textual_inversion import create_embedding, train_ from modules.textual_inversion.preprocess import preprocess from modules.hypernetworks.hypernetwork import create_hypernetwork, train_hypernetwork from PIL import PngImagePlugin,Image -from modules.sd_models import checkpoints_list, unload_model_weights, reload_model_weights, checkpoint_aliases -from modules.sd_vae import vae_dict +from modules.sd_models import unload_model_weights, reload_model_weights, checkpoint_aliases from modules.sd_models_config import find_checkpoint_config_near_filename from modules.realesrgan_model import get_realesrgan_models from modules import devices @@ -56,7 +57,41 @@ def setUpscalers(req: dict): return reqDict +def verify_url(url): + """Returns True if the url refers to a global resource.""" + + import socket + from urllib.parse import urlparse + try: + parsed_url = urlparse(url) + domain_name = parsed_url.netloc + host = socket.gethostbyname_ex(domain_name) + for ip in host[2]: + ip_addr = ipaddress.ip_address(ip) + if not ip_addr.is_global: + return False + except Exception: + return False + + return True + + def decode_base64_to_image(encoding): + if encoding.startswith("http://") or encoding.startswith("https://"): + if not opts.api_enable_requests: + raise HTTPException(status_code=500, detail="Requests not allowed") + + if opts.api_forbid_local_requests and not verify_url(encoding): + raise HTTPException(status_code=500, detail="Request to local resource not allowed") + + headers = {'user-agent': opts.api_useragent} if opts.api_useragent else {} + response = requests.get(encoding, timeout=30, headers=headers) + try: + image = Image.open(BytesIO(response.content)) + return image + except Exception as e: + raise HTTPException(status_code=500, detail="Invalid image url") from e + if encoding.startswith("data:image/"): encoding = encoding.split(";")[1].split(",")[1] try: @@ -197,6 +232,7 @@ class Api: self.add_api_route("/sdapi/v1/prompt-styles", self.get_prompt_styles, methods=["GET"], response_model=List[models.PromptStyleItem]) self.add_api_route("/sdapi/v1/embeddings", self.get_embeddings, methods=["GET"], response_model=models.EmbeddingsResponse) self.add_api_route("/sdapi/v1/refresh-checkpoints", self.refresh_checkpoints, methods=["POST"]) + self.add_api_route("/sdapi/v1/refresh-vae", self.refresh_vae, methods=["POST"]) self.add_api_route("/sdapi/v1/create/embedding", self.create_embedding, methods=["POST"], response_model=models.CreateResponse) self.add_api_route("/sdapi/v1/create/hypernetwork", self.create_hypernetwork, methods=["POST"], response_model=models.CreateResponse) self.add_api_route("/sdapi/v1/preprocess", self.preprocess, methods=["POST"], response_model=models.PreprocessResponse) @@ -329,6 +365,7 @@ class Api: with self.queue_lock: with closing(StableDiffusionProcessingTxt2Img(sd_model=shared.sd_model, **args)) as p: + p.is_api = True p.scripts = script_runner p.outpath_grids = opts.outdir_txt2img_grids p.outpath_samples = opts.outdir_txt2img_samples @@ -343,6 +380,7 @@ class Api: processed = process_images(p) finally: shared.state.end() + shared.total_tqdm.clear() b64images = list(map(encode_pil_to_base64, processed.images)) if send_images else [] @@ -388,6 +426,7 @@ class Api: with self.queue_lock: with closing(StableDiffusionProcessingImg2Img(sd_model=shared.sd_model, **args)) as p: p.init_images = [decode_base64_to_image(x) for x in init_images] + p.is_api = True p.scripts = script_runner p.outpath_grids = opts.outdir_img2img_grids p.outpath_samples = opts.outdir_img2img_samples @@ -402,6 +441,7 @@ class Api: processed = process_images(p) finally: shared.state.end() + shared.total_tqdm.clear() b64images = list(map(encode_pil_to_base64, processed.images)) if send_images else [] @@ -530,7 +570,7 @@ class Api: raise RuntimeError(f"model {checkpoint_name!r} not found") for k, v in req.items(): - shared.opts.set(k, v) + shared.opts.set(k, v, is_api=True) shared.opts.save(shared.config_filename) return @@ -562,10 +602,12 @@ class Api: ] def get_sd_models(self): - 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 checkpoints_list.values()] + import modules.sd_models as sd_models + 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()] def get_sd_vaes(self): - return [{"model_name": x, "filename": vae_dict[x]} for x in vae_dict.keys()] + import modules.sd_vae as sd_vae + return [{"model_name": x, "filename": sd_vae.vae_dict[x]} for x in sd_vae.vae_dict.keys()] def get_hypernetworks(self): return [{"name": name, "path": shared.hypernetworks[name]} for name in shared.hypernetworks] @@ -608,6 +650,10 @@ class Api: with self.queue_lock: shared.refresh_checkpoints() + def refresh_vae(self): + with self.queue_lock: + shared_items.refresh_vae_list() + def create_embedding(self, args: dict): try: shared.state.begin(job="create_embedding") diff --git a/modules/api/models.py b/modules/api/models.py index 800c9b93..6a574771 100644 --- a/modules/api/models.py +++ b/modules/api/models.py @@ -50,10 +50,12 @@ class PydanticModelGenerator: additional_fields = None, ): def field_type_generator(k, v): - # field_type = str if not overrides.get(k) else overrides[k]["type"] - # print(k, v.annotation, v.default) field_type = v.annotation + if field_type == 'Image': + # images are sent as base64 strings via API + field_type = 'str' + return Optional[field_type] def merge_class_params(class_): @@ -63,7 +65,6 @@ class PydanticModelGenerator: parameters = {**parameters, **inspect.signature(classes.__init__).parameters} return parameters - self._model_name = model_name self._class_data = merge_class_params(class_instance) @@ -72,7 +73,7 @@ class PydanticModelGenerator: field=underscore(k), field_alias=k, field_type=field_type_generator(k, v), - field_value=v.default + field_value=None if isinstance(v.default, property) else v.default ) for (k,v) in self._class_data.items() if k not in API_NOT_ALLOWED ] |