Adding DeepDanbooru to the interrogation API (67c8e11b) · Commits · github_fork / Stable Diffusion Webui

modules/api/api.py

+14 −2

Original line number	Original line	Diff line number	Diff line
	@@ -15,6 +15,9 @@ from modules.sd_models import checkpoints_list
	from modules.realesrgan_model import get_realesrgan_models		from modules.realesrgan_model import get_realesrgan_models
	from typing import List		from typing import List

			if shared.cmd_opts.deepdanbooru:
			from modules.deepbooru import get_deepbooru_tags

	def upscaler_to_index(name: str):		def upscaler_to_index(name: str):
	try:		try:
	return [x.name.lower() for x in shared.sd_upscalers].index(name.lower())		return [x.name.lower() for x in shared.sd_upscalers].index(name.lower())
	@@ -220,11 +223,20 @@ class Api:
	if image_b64 is None:		if image_b64 is None:
	raise HTTPException(status_code=404, detail="Image not found")		raise HTTPException(status_code=404, detail="Image not found")

	img = self.__base64_to_image(image_b64)		img = decode_base64_to_image(image_b64)
			img = img.convert('RGB')

	# Override object param		# Override object param
	with self.queue_lock:		with self.queue_lock:
			if interrogatereq.model == "clip":
	processed = shared.interrogator.interrogate(img)		processed = shared.interrogator.interrogate(img)
			elif interrogatereq.model == "deepdanbooru":
			if shared.cmd_opts.deepdanbooru:
			processed = get_deepbooru_tags(img)
			else:
			raise HTTPException(status_code=404, detail="Model not found. Add --deepdanbooru when launching for using the model.")
			else:
			raise HTTPException(status_code=404, detail="Model not found")

	return InterrogateResponse(caption=processed)		return InterrogateResponse(caption=processed)

+1 −0

Original line number	Original line	Diff line number	Diff line
	@@ -170,6 +170,7 @@ class ProgressResponse(BaseModel):

	class InterrogateRequest(BaseModel):		class InterrogateRequest(BaseModel):
	image: str = Field(default="", title="Image", description="Image to work on, must be a Base64 string containing the image's data.")		image: str = Field(default="", title="Image", description="Image to work on, must be a Base64 string containing the image's data.")
			model: str = Field(default="clip", title="Model", description="The interrogate model used.")

	class InterrogateResponse(BaseModel):		class InterrogateResponse(BaseModel):
	caption: str = Field(default=None, title="Caption", description="The generated caption for the image.")		caption: str = Field(default=None, title="Caption", description="The generated caption for the image.")