Add a barebones interrogate API

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

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

View File

@@ -63,7 +63,12 @@ class PydanticModelGenerator:
self._model_name = model_name
self._class_data = merge_class_params(class_instance)
if class_instance is not None:
self._class_data = merge_class_params(class_instance)
else:
self._class_data = {}
self._model_def = [
ModelDef(
field=underscore(k),
@@ -105,4 +110,10 @@ StableDiffusionImg2ImgProcessingAPI = PydanticModelGenerator(
"StableDiffusionProcessingImg2Img",
StableDiffusionProcessingImg2Img,
[{"key": "sampler_index", "type": str, "default": "Euler"}, {"key": "init_images", "type": list, "default": None}, {"key": "denoising_strength", "type": float, "default": 0.75}, {"key": "mask", "type": str, "default": None}, {"key": "include_init_images", "type": bool, "default": False, "exclude" : True}]
).generate_model()
InterrogateAPI = PydanticModelGenerator(
"Interrogate",
None,
[{"key": "image", "type": str, "default": None}]
).generate_model()