Improve performance of adding tags
zanllp
parent
5e98f68657
commit
c80afce4a9
|
|
@ -49,7 +49,7 @@ from scripts.iib.db.datamodel import (
|
|||
FileInfoDict,
|
||||
Cursor
|
||||
)
|
||||
from scripts.iib.db.update_image_data import update_image_data, rebuild_image_index
|
||||
from scripts.iib.db.update_image_data import update_image_data, rebuild_image_index, add_image_data_single
|
||||
from scripts.iib.logger import logger
|
||||
from scripts.iib.seq import seq
|
||||
import urllib.parse
|
||||
|
|
@ -823,6 +823,7 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
|
|||
if not img:
|
||||
if DbImg.count(conn):
|
||||
update_image_data([os.path.dirname(path)])
|
||||
add_image_data_single(path)
|
||||
img = DbImg.get(conn, path)
|
||||
else:
|
||||
raise HTTPException(
|
||||
|
|
@ -866,8 +867,8 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
|
|||
)
|
||||
img = DbImg.get(conn, path)
|
||||
if not img:
|
||||
if DbImg.count(conn):
|
||||
update_image_data([os.path.dirname(path)])
|
||||
if DbImg.count(conn):
|
||||
add_image_data_single(path)
|
||||
img = DbImg.get(conn, path)
|
||||
else:
|
||||
raise HTTPException(
|
||||
|
|
|
|||
|
|
@ -40,78 +40,14 @@ def update_image_data(search_dirs: List[str], is_rebuild = False):
|
|||
if not is_rebuild and not Folder.check_need_update(conn, folder_path):
|
||||
return
|
||||
print(f"Processing folder: {folder_path}")
|
||||
parsed_params = ImageGenerationParams()
|
||||
for filename in os.listdir(folder_path):
|
||||
file_path = os.path.normpath(os.path.join(folder_path, filename))
|
||||
try:
|
||||
|
||||
if os.path.isdir(file_path):
|
||||
process_folder(file_path)
|
||||
|
||||
elif is_valid_media_path(file_path):
|
||||
img = DbImg.get(conn, file_path)
|
||||
if is_rebuild:
|
||||
info = get_exif_data(file_path)
|
||||
parsed_params = info.params
|
||||
if not img:
|
||||
img = DbImg(
|
||||
file_path,
|
||||
info.raw_info,
|
||||
os.path.getsize(file_path),
|
||||
get_modified_date(file_path),
|
||||
)
|
||||
img.save(conn)
|
||||
else:
|
||||
if img: # 已存在的跳过
|
||||
if img.date == get_modified_date(img.path):
|
||||
continue
|
||||
else:
|
||||
DbImg.safe_batch_remove(conn=conn, image_ids=[img.id])
|
||||
info = get_exif_data(file_path)
|
||||
parsed_params = info.params
|
||||
img = DbImg(
|
||||
file_path,
|
||||
info.raw_info,
|
||||
os.path.getsize(file_path),
|
||||
get_modified_date(file_path),
|
||||
)
|
||||
img.save(conn)
|
||||
|
||||
if not parsed_params:
|
||||
continue
|
||||
meta = parsed_params.meta
|
||||
lora = parsed_params.extra.get("lora", [])
|
||||
lyco = parsed_params.extra.get("lyco", [])
|
||||
pos = parsed_params.pos_prompt
|
||||
size_tag = Tag.get_or_create(
|
||||
conn,
|
||||
str(meta.get("Size-1", 0)) + " * " + str(meta.get("Size-2", 0)),
|
||||
type="size",
|
||||
)
|
||||
safe_save_img_tag(ImageTag(img.id, size_tag.id))
|
||||
|
||||
for k in [
|
||||
"Model",
|
||||
"Sampler",
|
||||
"Source Identifier",
|
||||
"Postprocess upscale by",
|
||||
"Postprocess upscaler",
|
||||
"Size"
|
||||
]:
|
||||
v = meta.get(k)
|
||||
if not v:
|
||||
continue
|
||||
tag = Tag.get_or_create(conn, str(v), k)
|
||||
safe_save_img_tag(ImageTag(img.id, tag.id))
|
||||
for i in lora:
|
||||
tag = Tag.get_or_create(conn, i["name"], "lora")
|
||||
safe_save_img_tag(ImageTag(img.id, tag.id))
|
||||
for i in lyco:
|
||||
tag = Tag.get_or_create(conn, i["name"], "lyco")
|
||||
safe_save_img_tag(ImageTag(img.id, tag.id))
|
||||
for k in pos:
|
||||
tag = Tag.get_or_create(conn, k, "pos")
|
||||
safe_save_img_tag(ImageTag(img.id, tag.id))
|
||||
build_single_img_idx(conn, file_path, is_rebuild, safe_save_img_tag)
|
||||
# neg暂时跳过感觉个没人会搜索这个
|
||||
except Exception as e:
|
||||
logger.error("Tag generation failed. Skipping this file. file:%s error: %s", file_path, e)
|
||||
|
|
@ -128,6 +64,32 @@ def update_image_data(search_dirs: List[str], is_rebuild = False):
|
|||
tag.save(conn)
|
||||
conn.commit()
|
||||
|
||||
def add_image_data_single(file_path):
|
||||
conn = DataBase.get_conn()
|
||||
tag_incr_count_rec: Dict[int, int] = {}
|
||||
|
||||
def safe_save_img_tag(img_tag: ImageTag):
|
||||
tag_incr_count_rec[img_tag.tag_id] = (
|
||||
tag_incr_count_rec.get(img_tag.tag_id, 0) + 1
|
||||
)
|
||||
img_tag.save_or_ignore(conn)
|
||||
|
||||
file_path = os.path.normpath(file_path)
|
||||
try:
|
||||
if not is_valid_media_path(file_path):
|
||||
return
|
||||
build_single_img_idx(conn, file_path, False, safe_save_img_tag)
|
||||
# neg暂时跳过感觉个没人会搜索这个
|
||||
except Exception as e:
|
||||
logger.error("Tag generation failed. Skipping this file. file:%s error: %s", file_path, e)
|
||||
conn.commit()
|
||||
|
||||
for tag_id in tag_incr_count_rec:
|
||||
tag = Tag.get(conn, tag_id)
|
||||
tag.count += tag_incr_count_rec[tag_id]
|
||||
tag.save(conn)
|
||||
conn.commit()
|
||||
|
||||
def rebuild_image_index(search_dirs: List[str]):
|
||||
conn = DataBase.get_conn()
|
||||
with closing(conn.cursor()) as cur:
|
||||
|
|
@ -140,4 +102,70 @@ def rebuild_image_index(search_dirs: List[str]):
|
|||
)
|
||||
cur.execute("""DELETE FROM tag WHERE tag.type <> 'custom'""")
|
||||
conn.commit()
|
||||
update_image_data(search_dirs=search_dirs, is_rebuild=True)
|
||||
update_image_data(search_dirs=search_dirs, is_rebuild=True)
|
||||
|
||||
def build_single_img_idx(conn, file_path, is_rebuild, safe_save_img_tag):
|
||||
img = DbImg.get(conn, file_path)
|
||||
parsed_params = None
|
||||
if is_rebuild:
|
||||
info = get_exif_data(file_path)
|
||||
parsed_params = info.params
|
||||
if not img:
|
||||
img = DbImg(
|
||||
file_path,
|
||||
info.raw_info,
|
||||
os.path.getsize(file_path),
|
||||
get_modified_date(file_path),
|
||||
)
|
||||
img.save(conn)
|
||||
else:
|
||||
if img: # 已存在的跳过
|
||||
if img.date == get_modified_date(img.path):
|
||||
return
|
||||
else:
|
||||
DbImg.safe_batch_remove(conn=conn, image_ids=[img.id])
|
||||
info = get_exif_data(file_path)
|
||||
parsed_params = info.params
|
||||
img = DbImg(
|
||||
file_path,
|
||||
info.raw_info,
|
||||
os.path.getsize(file_path),
|
||||
get_modified_date(file_path),
|
||||
)
|
||||
img.save(conn)
|
||||
|
||||
if not parsed_params:
|
||||
return
|
||||
meta = parsed_params.meta
|
||||
lora = parsed_params.extra.get("lora", [])
|
||||
lyco = parsed_params.extra.get("lyco", [])
|
||||
pos = parsed_params.pos_prompt
|
||||
size_tag = Tag.get_or_create(
|
||||
conn,
|
||||
str(meta.get("Size-1", 0)) + " * " + str(meta.get("Size-2", 0)),
|
||||
type="size",
|
||||
)
|
||||
safe_save_img_tag(ImageTag(img.id, size_tag.id))
|
||||
|
||||
for k in [
|
||||
"Model",
|
||||
"Sampler",
|
||||
"Source Identifier",
|
||||
"Postprocess upscale by",
|
||||
"Postprocess upscaler",
|
||||
"Size"
|
||||
]:
|
||||
v = meta.get(k)
|
||||
if not v:
|
||||
continue
|
||||
tag = Tag.get_or_create(conn, str(v), k)
|
||||
safe_save_img_tag(ImageTag(img.id, tag.id))
|
||||
for i in lora:
|
||||
tag = Tag.get_or_create(conn, i["name"], "lora")
|
||||
safe_save_img_tag(ImageTag(img.id, tag.id))
|
||||
for i in lyco:
|
||||
tag = Tag.get_or_create(conn, i["name"], "lyco")
|
||||
safe_save_img_tag(ImageTag(img.id, tag.id))
|
||||
for k in pos:
|
||||
tag = Tag.get_or_create(conn, k, "pos")
|
||||
safe_save_img_tag(ImageTag(img.id, tag.id))
|
||||
Loading…
Reference in New Issue