mirrors/stable-diffusion-webui
1
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2025-12-06 02:30:30 -08:00
Code Issues Projects Releases Packages Wiki Activity

Add a barebones interrogate API

Browse source
This commit is contained in:
Roy Shilkrot 2022-10-27 15:20:15 -04:00
parent 737eb28fac
commit bdc9083798
4 changed files with 45 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

25
modules/api/api.py
View file

@ -1,4 +1,4 @@
from modules.api.models import StableDiffusionTxt2ImgProcessingAPI, StableDiffusionImg2ImgProcessingAPI
from modules.api.models import StableDiffusionTxt2ImgProcessingAPI, StableDiffusionImg2ImgProcessingAPI, InterrogateAPI
from modules.processing import StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img, process_images
from modules.sd_samplers import all_samplers
from modules.extras import run_pnginfo
@ -25,6 +25,11 @@ class ImageToImageResponse(BaseModel):
parameters: Json
info: Json
class InterrogateResponse(BaseModel):
caption: str = Field(default=None, title="Caption", description="The generated caption for the image.")
parameters: Json
info: Json
class Api:
def __init__(self, app, queue_lock):
@ -33,6 +38,7 @@ class Api:
self.queue_lock = queue_lock
self.app.add_api_route("/sdapi/v1/txt2img", self.text2imgapi, methods=["POST"])
self.app.add_api_route("/sdapi/v1/img2img", self.img2imgapi, methods=["POST"])
self.app.add_api_route("/sdapi/v1/interrogate", self.interrogateapi, methods=["POST"])
def __base64_to_image(self, base64_string):
# if has a comma, deal with prefix
@ -118,6 +124,23 @@ class Api:
return ImageToImageResponse(images=b64images, parameters=json.dumps(vars(img2imgreq)), info=processed.js())
def interrogateapi(self, interrogatereq: InterrogateAPI):
image_b64 = interrogatereq.image
if image_b64 is None:
raise HTTPException(status_code=404, detail="Image not found")
populate = interrogatereq.copy(update={ # Override __init__ params
}
)
img = self.__base64_to_image(image_b64)
# Override object param
with self.queue_lock:
processed = shared.interrogator.interrogate(img)
return InterrogateResponse(caption=processed, parameters=json.dumps(vars(interrogatereq)), info=None)
def extrasapi(self):
raise NotImplementedError