mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-12-17 01:52:59 +08:00
optimization: initial scan speed(batching tags)
This commit is contained in:
parent
5b6810a2c6
commit
85ef08449d
@ -26,6 +26,7 @@ from .database.helpers import (
|
|||||||
ensure_tags_exist,
|
ensure_tags_exist,
|
||||||
escape_like_prefix,
|
escape_like_prefix,
|
||||||
fast_asset_file_check,
|
fast_asset_file_check,
|
||||||
|
insert_tags_from_batch,
|
||||||
remove_missing_tag_for_asset_id,
|
remove_missing_tag_for_asset_id,
|
||||||
)
|
)
|
||||||
from .database.models import Asset, AssetCacheState, AssetInfo
|
from .database.models import Asset, AssetCacheState, AssetInfo
|
||||||
@ -146,6 +147,8 @@ async def sync_seed_assets(roots: list[schemas_in.RootType]) -> None:
|
|||||||
async with await create_session() as sess:
|
async with await create_session() as sess:
|
||||||
if tag_pool:
|
if tag_pool:
|
||||||
await ensure_tags_exist(sess, tag_pool, tag_type="user")
|
await ensure_tags_exist(sess, tag_pool, tag_type="user")
|
||||||
|
|
||||||
|
pending_tag_links: list[dict] = []
|
||||||
for ap, sz, mt, name, tags in new_specs:
|
for ap, sz, mt, name, tags in new_specs:
|
||||||
await seed_from_path(
|
await seed_from_path(
|
||||||
sess,
|
sess,
|
||||||
@ -155,12 +158,17 @@ async def sync_seed_assets(roots: list[schemas_in.RootType]) -> None:
|
|||||||
info_name=name,
|
info_name=name,
|
||||||
tags=tags,
|
tags=tags,
|
||||||
owner_id="",
|
owner_id="",
|
||||||
skip_tag_ensure=True,
|
collected_tag_rows=pending_tag_links,
|
||||||
)
|
)
|
||||||
|
|
||||||
created += 1
|
created += 1
|
||||||
if created % 500 == 0:
|
if created % 500 == 0:
|
||||||
|
if pending_tag_links:
|
||||||
|
await insert_tags_from_batch(sess, tag_rows=pending_tag_links)
|
||||||
|
pending_tag_links.clear()
|
||||||
await sess.commit()
|
await sess.commit()
|
||||||
|
if pending_tag_links:
|
||||||
|
await insert_tags_from_batch(sess, tag_rows=pending_tag_links)
|
||||||
await sess.commit()
|
await sess.commit()
|
||||||
finally:
|
finally:
|
||||||
LOGGER.info(
|
LOGGER.info(
|
||||||
|
|||||||
@ -6,6 +6,7 @@ from .projection import is_scalar, project_kv
|
|||||||
from .tags import (
|
from .tags import (
|
||||||
add_missing_tag_for_asset_id,
|
add_missing_tag_for_asset_id,
|
||||||
ensure_tags_exist,
|
ensure_tags_exist,
|
||||||
|
insert_tags_from_batch,
|
||||||
remove_missing_tag_for_asset_id,
|
remove_missing_tag_for_asset_id,
|
||||||
)
|
)
|
||||||
|
|
||||||
@ -19,5 +20,6 @@ __all__ = [
|
|||||||
"ensure_tags_exist",
|
"ensure_tags_exist",
|
||||||
"add_missing_tag_for_asset_id",
|
"add_missing_tag_for_asset_id",
|
||||||
"remove_missing_tag_for_asset_id",
|
"remove_missing_tag_for_asset_id",
|
||||||
|
"insert_tags_from_batch",
|
||||||
"visible_owner_clause",
|
"visible_owner_clause",
|
||||||
]
|
]
|
||||||
|
|||||||
@ -88,3 +88,19 @@ async def remove_missing_tag_for_asset_id(
|
|||||||
AssetInfoTag.tag_name == "missing",
|
AssetInfoTag.tag_name == "missing",
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def insert_tags_from_batch(session: AsyncSession, *, tag_rows: list[dict]) -> None:
|
||||||
|
if session.bind.dialect.name == "sqlite":
|
||||||
|
ins_links = (
|
||||||
|
d_sqlite.insert(AssetInfoTag)
|
||||||
|
.values(tag_rows)
|
||||||
|
.on_conflict_do_nothing(index_elements=[AssetInfoTag.asset_info_id, AssetInfoTag.tag_name])
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
ins_links = (
|
||||||
|
d_pg.insert(AssetInfoTag)
|
||||||
|
.values(tag_rows)
|
||||||
|
.on_conflict_do_nothing(index_elements=[AssetInfoTag.asset_info_id, AssetInfoTag.tag_name])
|
||||||
|
)
|
||||||
|
await session.execute(ins_links)
|
||||||
|
|||||||
@ -67,7 +67,7 @@ async def seed_from_path(
|
|||||||
info_name: str,
|
info_name: str,
|
||||||
tags: Sequence[str],
|
tags: Sequence[str],
|
||||||
owner_id: str = "",
|
owner_id: str = "",
|
||||||
skip_tag_ensure: bool = False,
|
collected_tag_rows: list[dict],
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Creates Asset(hash=NULL), AssetCacheState(file_path), and AssetInfo exist for the path."""
|
"""Creates Asset(hash=NULL), AssetCacheState(file_path), and AssetInfo exist for the path."""
|
||||||
locator = os.path.abspath(abs_path)
|
locator = os.path.abspath(abs_path)
|
||||||
@ -154,8 +154,6 @@ async def seed_from_path(
|
|||||||
if info_inserted:
|
if info_inserted:
|
||||||
want = normalize_tags(tags)
|
want = normalize_tags(tags)
|
||||||
if want:
|
if want:
|
||||||
if not skip_tag_ensure:
|
|
||||||
await ensure_tags_exist(session, want, tag_type="user")
|
|
||||||
tag_rows = [
|
tag_rows = [
|
||||||
{
|
{
|
||||||
"asset_info_id": new_info_id,
|
"asset_info_id": new_info_id,
|
||||||
@ -165,19 +163,7 @@ async def seed_from_path(
|
|||||||
}
|
}
|
||||||
for t in want
|
for t in want
|
||||||
]
|
]
|
||||||
if dialect == "sqlite":
|
collected_tag_rows.extend(tag_rows)
|
||||||
ins_links = (
|
|
||||||
d_sqlite.insert(AssetInfoTag)
|
|
||||||
.values(tag_rows)
|
|
||||||
.on_conflict_do_nothing(index_elements=[AssetInfoTag.asset_info_id, AssetInfoTag.tag_name])
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
ins_links = (
|
|
||||||
d_pg.insert(AssetInfoTag)
|
|
||||||
.values(tag_rows)
|
|
||||||
.on_conflict_do_nothing(index_elements=[AssetInfoTag.asset_info_id, AssetInfoTag.tag_name])
|
|
||||||
)
|
|
||||||
await session.execute(ins_links)
|
|
||||||
|
|
||||||
if fname: # simple filename projection with single row
|
if fname: # simple filename projection with single row
|
||||||
meta_row = {
|
meta_row = {
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user