File size: 5,729 Bytes
8e68b99 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 |
from contextlib import closing
from typing import Dict, List
from scripts.iib.db.datamodel import Image as DbImg, Tag, ImageTag, DataBase, Folder
import os
from PIL import Image
from scripts.iib.tool import (
read_sd_webui_gen_info_from_image,
parse_generation_parameters,
is_valid_image_path,
get_modified_date,
is_dev,
get_comfyui_exif_data,
comfyui_exif_data_to_str,
is_img_created_by_comfyui,
is_img_created_by_comfyui_with_webui_gen_info
)
from scripts.iib.logger import logger
# 定义一个函数来获取图片文件的EXIF数据
def get_exif_data(file_path):
info = ''
params = None
try:
with Image.open(file_path) as img:
if is_img_created_by_comfyui(img):
if is_img_created_by_comfyui_with_webui_gen_info(img):
info = read_sd_webui_gen_info_from_image(img, file_path)
params = parse_generation_parameters(info)
else:
params = get_comfyui_exif_data(img)
info = comfyui_exif_data_to_str(params)
else:
info = read_sd_webui_gen_info_from_image(img, file_path)
params = parse_generation_parameters(info)
except Exception as e:
if is_dev:
logger.error("get_exif_data %s", e)
return params, info
def update_image_data(search_dirs: List[str], is_rebuild = False):
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) # 原先用来处理一些意外，但是写的正确完全没问题,去掉了try catch
# 递归处理每个文件夹
def process_folder(folder_path: str):
if not is_rebuild and not Folder.check_need_update(conn, folder_path):
return
print(f"Processing folder: {folder_path}")
for filename in os.listdir(folder_path):
file_path = os.path.normpath(os.path.join(folder_path, filename))
if os.path.isdir(file_path):
process_folder(file_path)
elif is_valid_image_path(file_path):
img = DbImg.get(conn, file_path)
if is_rebuild:
parsed_params, info = get_exif_data(file_path)
if not img:
img = DbImg(
file_path,
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])
parsed_params, info = get_exif_data(file_path)
img = DbImg(
file_path,
info,
os.path.getsize(file_path),
get_modified_date(file_path),
)
img.save(conn)
if not parsed_params:
continue
meta = parsed_params.get("meta", {})
lora = parsed_params.get("lora", [])
lyco = parsed_params.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",
"Postprocess upscale by",
"Postprocess upscaler",
]:
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))
# neg暂时跳过感觉个没人会搜索这个
# 提交对数据库的更改
Folder.update_modified_date_or_create(conn, folder_path)
conn.commit()
for dir in search_dirs:
process_folder(dir)
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:
cur.execute(
"""DELETE FROM image_tag
WHERE image_tag.tag_id IN (
SELECT tag.id FROM tag WHERE tag.type <> 'custom'
)
"""
)
cur.execute("""DELETE FROM tag WHERE tag.type <> 'custom'""")
conn.commit()
update_image_data(search_dirs=search_dirs, is_rebuild=True) |