diff --git a/scripts/iib/api.py b/scripts/iib/api.py index d4e849d..361f847 100644 --- a/scripts/iib/api.py +++ b/scripts/iib/api.py @@ -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( diff --git a/scripts/iib/db/update_image_data.py b/scripts/iib/db/update_image_data.py index 7f5c2a8..3979c93 100644 --- a/scripts/iib/db/update_image_data.py +++ b/scripts/iib/db/update_image_data.py @@ -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) \ No newline at end of file + 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)) \ No newline at end of file