refactor interrogate/caption

Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
2025-02-01 15:40:19 -05:00
parent 654f44f66f
commit 2963ce127c
19 changed files with 239 additions and 159 deletions
+6 -6
View File
@@ -74,8 +74,8 @@ def get_extra_networks(page: Optional[str] = None, name: Optional[str] = None, f
return res
def get_interrogate():
from modules.interrogate.legacy import get_clip_models
return ['clip', 'deepdanbooru'] + get_clip_models()
from modules.interrogate.openclip import refresh_clip_models
return ['clip', 'deepdanbooru'] + refresh_clip_models()
def post_interrogate(req: models.ReqInterrogate):
if req.image is None or len(req.image) < 64:
@@ -84,8 +84,8 @@ def post_interrogate(req: models.ReqInterrogate):
image = image.convert('RGB')
if req.model == "clip":
try:
from modules.interrogate import legacy
caption = legacy.interrogator.interrogate(image)
from modules.interrogate import openclip
caption = openclip.interrogator.interrogate(image)
except Exception as e:
caption = str(e)
return models.ResInterrogate(caption=caption)
@@ -94,8 +94,8 @@ def post_interrogate(req: models.ReqInterrogate):
caption = deepbooru.model.tag(image)
return models.ResInterrogate(caption=caption)
else:
from modules.interrogate.legacy import interrogate_image, analyze_image, get_clip_models
if req.model not in get_clip_models():
from modules.interrogate.openclip import interrogate_image, analyze_image, refresh_clip_models
if req.model not in refresh_clip_models():
raise HTTPException(status_code=404, detail="Model not found")
try:
caption = interrogate_image(image, clip_model=req.clip_model, blip_model=req.blip_model, mode=req.mode)