import os
import sys
_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
if _SCRIPT_DIR not in sys.path:
sys.path.insert(0, _SCRIPT_DIR)
try:
sys.stdout.reconfigure(encoding="utf-8")
except:
pass
try:
sys.stderr.reconfigure(encoding="utf-8")
except:
pass
import time
import json
import re
import urllib.request
import argparse
import shutil
from model_policies import is_physical_rename_protected
from recycle_bin import move_to_trash
from model_identity import USER_INFO_SUFFIX, computed_file_identity, file_sha256, infer_base_model_from_header, is_unmatched, safetensors_damage, scan_choice, scan_info_path, sidecar_file_hash, sidecar_info as read_local_info, unmatched_reason
from civitai_client import CIVITAI_API_KEY, CivitaiUnreachable, download_media, fetch_civitai_info, fetch_model_description
import model_sidecars as sidecars
# Fixed tuples avoid rebuilding long extension lists for every scanned model.
# Exact-path checks are intentionally used instead of directory-wide globbing.
MEDIA_EXTENSIONS = (".png", ".jpg", ".jpeg", ".webp", ".gif", ".avif", ".mp4", ".webm", ".mov", ".avi")
PREVIEW_SUFFIXES = tuple(f".preview{ext}" for ext in MEDIA_EXTENSIONS)
CIVITAI_BACKUP_SUFFIXES = tuple(f".civitai_bak{ext}" for ext in MEDIA_EXTENSIONS)
COVER_SUFFIXES = MEDIA_EXTENSIONS + PREVIEW_SUFFIXES
SIDECAR_SUFFIXES = (
".info", ".civitai.info", USER_INFO_SUFFIX, ".json", ".txt", ".yaml",
*MEDIA_EXTENSIONS,
*PREVIEW_SUFFIXES,
*CIVITAI_BACKUP_SUFFIXES,
)
def drop_partial(base):
"""A download that did not finish (.) goes."""
for ext in MEDIA_EXTENSIONS:
try:
os.remove(base + ext)
except OSError:
pass
def write_scan_progress(progress_file, phase, total=0, current=0, filename=""):
if not progress_file:
return
temp_file = f"{progress_file}.{os.getpid()}.tmp"
try:
with open(temp_file, 'w', encoding='utf-8') as f:
json.dump({
"phase": phase,
"total": total,
"current": current,
"filename": filename,
}, f, ensure_ascii=False)
os.replace(temp_file, progress_file)
except OSError:
try:
os.remove(temp_file)
except OSError:
pass
def left_out(root, filename, selected):
"""A model the user set to "don't scan" (model_identity.scan_choice), in a scan of its whole folder."""
return not selected and scan_choice(os.path.join(root, os.path.splitext(filename)[0])) == "skip"
def count_scan_files(target_folder, target_files_basenames, unmatched_only=False):
total = 0
selected = set(target_files_basenames)
for root, _, files in os.walk(target_folder):
total += sum(
1 for filename in files
if filename.endswith(".safetensors") and (not selected or filename in selected)
and not left_out(root, filename, selected)
and (not unmatched_only or is_unmatched(read_local_info(os.path.join(root, os.path.splitext(filename)[0]))))
)
return total
def calculate_sha256(file_path: str) -> str:
"""计算文件的 SHA256 哈希值 (用于 Civitai 匹配)"""
print(f"[*] 正在计算 Hash (大文件可能需要几分钟): {os.path.basename(file_path)}")
return file_sha256(file_path)
def sanitize_filename(name: str) -> str:
"""清理文件名中的非法字符"""
name = re.sub(r'[\r\n\t]+', ' ', name)
name = re.sub(r'[\\/*?:"<>|#]', "", name)
return name.strip(' .')
CIVITAI_DOWN_AFTER = 2 # models in a row Civitai did not answer for before the rest of the run stays offline
# Covers in a row that would not download before the rest of the run fetches none: the API can
# answer on civitai.red where the pictures (image.civitai.com) cannot be reached, and each try waits.
COVERS_DOWN_AFTER = 2
def write_report(args, **entry):
"""One line of the scan report (--report-file): a model's outcome, or an event of the run."""
if not args.report_file:
return
try:
with open(args.report_file, 'a', encoding='utf-8') as f:
f.write(json.dumps(entry, ensure_ascii=False) + "\n")
except OSError as e:
print(f"[-] 扫描记录写入失败: {e}")
def local_record(file_path, root, file_hash, reason):
"""What the file alone tells (id -1), with why Civitai's record is missing."""
inferred_base = infer_base_model_from_header(file_path)
stem = os.path.splitext(os.path.basename(file_path))[0]
print(f"[*] 使用本地哈希重建基础元数据 ({os.path.basename(file_path)})")
return {
"id": -1,
"modelId": -1,
"name": stem,
"baseModel": "" if inferred_base == 'Unknown' else inferred_base,
"description": "Automatically inferred by Anomalous Local Engine.
",
"model": {"name": stem, "type": "LORA" if "lora" in root.lower() else "Checkpoint"},
"files": [{"hashes": {"SHA256": file_hash}}],
"anomalous_unmatched_reason": reason,
}
def ask_civitai(args, state, file_hash):
"""(Civitai's record or None, why it is missing: "not_found" | "network" | "offline" | "")."""
if args.offline_only:
return None, "offline"
if state["civitai_down"]:
return None, "network"
try:
data = fetch_civitai_info(file_hash)
except CivitaiUnreachable as unreachable:
state["network_failures"] += 1
if state["network_failures"] >= CIVITAI_DOWN_AFTER:
state["civitai_down"] = True
print("\033[93m[!] 连续几个模型都连不上 Civitai，本次其余模型改为从文件推断，之后可以“重新联网查找”。\033[0m")
write_report(args, event="civitai_down", reason=str(unreachable))
return None, "network"
state["network_failures"] = 0
return data, ("" if data else "not_found")
def scan_one(args, state, rename_log, root, filename, previous):
"""Scans one model; returns its report entry, or None when nothing about it changed."""
file_path = os.path.join(root, filename)
old_base = os.path.splitext(file_path)[0]
online = not args.offline_only and not state["civitai_down"]
# A .civitai.info holding only edits from older versions is not a scan result.
info_exists = scan_info_path(old_base) is not None
# Models Civitai was never asked about (an offline scan, no connection then) are looked up
# by every online scan; models Civitai did not know wait for "look up again".
unasked = is_unmatched(previous) and unmatched_reason(previous) in ("offline", "network")
if args.force_overwrite or args.retry_unmatched or (online and unasked):
info_exists = False
preview_exists = args.skip_media or any(os.path.exists(old_base + ext) for ext in COVER_SUFFIXES)
needs_rename = bool(args.virtual_rename or (not args.skip_rename and args.physical_rename))
if info_exists and preview_exists and not needs_rename:
print(f"[*] 已跳过 (信息满足要求): {filename}")
return None
print(f"\n---> 处理文件: {filename} (位于 {root})")
civitai_data = None
file_hash = None
reason = ""
if info_exists: # a rename, or a missing cover fetched from the saved record: no new lookup
civitai_data = read_local_info(old_base)
if civitai_data:
print(f"[*] 使用本地已有的信息")
if not civitai_data:
if is_unmatched(previous): # the hash from last time, if the file is unchanged
file_hash = sidecar_file_hash(previous, file_path, (previous.get("files") or [{}])[0])[0]
file_hash = file_hash or calculate_sha256(file_path)
civitai_data, reason = ask_civitai(args, state, file_hash)
# A damaged header changes the file's hash and hides its tensors: say so instead of
# "Civitai does not know it", so the user downloads it again.
damage = "" if civitai_data else safetensors_damage(file_path)
if damage:
print(f"\033[91m[-] 文件头损坏 ({damage})，文件可能不完整或被改坏了，建议重新下载: {filename}\033[0m")
reason = "damaged"
# What was read from disk is not fetched again: only fresh Civitai data needs the model
# page's description and the cover.
fresh = civitai_data is not None and file_hash is not None
if not civitai_data:
if args.skip_local_metadata:
print(f"[*] Civitai 获取失败，且禁用了本地元数据解析。跳过 {filename}")
return {"status": "failed", "error": "no_local_metadata", "reason": reason}
civitai_data = local_record(file_path, root, file_hash, reason)
entry = {"status": "matched" if fresh else "inferred" if reason else "unchanged"}
if reason:
entry["reason"] = reason
model_id = civitai_data.get("modelId")
if fresh and model_id and model_id != -1 and not state["civitai_down"]:
description = fetch_model_description(model_id)
if description:
civitai_data["description"] = description
if not isinstance(civitai_data.get("model"), dict):
civitai_data["model"] = {}
civitai_data["model"]["description"] = description
model_name = sanitize_filename(civitai_data.get("model", {}).get("name", "UnknownModel"))
version_name = sanitize_filename(civitai_data.get("name", "UnknownVersion"))
if fresh:
entry["name"] = f"{civitai_data.get('model', {}).get('name', '')} · {civitai_data.get('name', '')}".strip(" ·")
entry["base"] = civitai_data.get("baseModel") or ""
new_filename = f"{model_name}_{version_name}.safetensors"
new_file_path = os.path.join(root, new_filename)
new_base = os.path.splitext(new_file_path)[0]
info_data = civitai_data
if file_hash:
info_data["anomalous_file_identity"] = computed_file_identity(file_path, file_hash)
# What the file alone tells has no standard name: both names are the file's own, so a
# rename would only repeat it (name_name, then name_name_name_name on the next scan).
inferred = is_unmatched(civitai_data)
# Civitai's name for the card; a name set in the model editor lives in
# .anomalous.json and is shown instead.
if inferred:
info_data.pop("anomalous_custom_name", None)
elif args.virtual_rename:
info_data["anomalous_custom_name"] = f"{model_name}_{version_name}"
noted = sidecars.civitai_cover_hash(old_base)
if noted and sidecars.COVER_HASH_KEY not in info_data:
info_data[sidecars.COVER_HASH_KEY] = noted # fresh data keeps knowing the cover until a new one comes
if not args.dry_run:
with open(sidecars.target(old_base, ".info"), 'w', encoding='utf-8') as f:
json.dump(info_data, f, ensure_ascii=True, indent=4)
print(f"[+] 写入纯净版 Civitai 描述 -> .info")
else:
print(f"[Dry-Run] 拟生成标准描述信息 -> .info")
media_url = None
if not args.skip_media:
media_url = next((img.get("url") for img in civitai_data.get("images") or [] if img.get("url")), None)
civitai_covers, user_covers = sidecars.cover_owners(old_base)
# Download when the data is fresh, or when the model has no cover at all.
if media_url and not state["civitai_down"] and not state["covers_down"] and (fresh or not (civitai_covers or user_covers)):
if not args.dry_run:
print(f"[*] 正在下载预览媒体...")
# Fetched under a working name, so a broken download never replaces a cover.
working = old_base + ".amb_download"
fetched = download_media(media_url, working)
if not fetched:
drop_partial(working)
saved_path = None
if fetched and user_covers:
# Civitai's picture is kept aside (in .amb when the setting says), to bring back later.
kept_before = sidecars.backups(old_base)
aside = sidecars.target(old_base, ".civitai_bak" + os.path.splitext(fetched)[1])
os.replace(fetched, aside)
saved_path = aside
try:
older = [path for path in kept_before if path != aside]
if older:
move_to_trash(*older)
except Exception as e:
print(f"[-] 旧的 Civitai 封面备份没有移走: {e}")
elif fetched:
ext = os.path.splitext(fetched)[1]
saved_path = old_base + (ext if ext.startswith('.preview.') else f".preview{ext}")
state["cover_failures"] = 0 if saved_path else state["cover_failures"] + 1
if state["cover_failures"] >= COVERS_DOWN_AFTER:
state["covers_down"] = True
print("\033[93m[!] 连续几张封面都下载不了，本次其余模型不再下载封面（名字和触发词照常获取）。\033[0m")
if saved_path:
print(f"[+] 媒体下载成功 -> {os.path.basename(saved_path)}")
entry["cover"] = "kept" if user_covers else "new"
if user_covers:
# The user's cover always stays; Civitai's image is kept as the backup.
print(f"[*] 保留你自己的封面，Civitai 封面存为备份")
else:
# Earlier Civitai pictures, as covers or kept aside, give way to the new one:
# the cover is Civitai's, so nothing needs keeping aside.
older = [c for c in civitai_covers if c != saved_path] + sidecars.backups(old_base)
try:
if older:
move_to_trash(*older)
except Exception as e:
print(f"[-] 旧的 Civitai 封面没有移走: {e}")
os.replace(fetched, saved_path)
sidecars.note_cover_hash(old_base, saved_path)
else:
print(f"[Dry-Run] 拟下载预览媒体...")
if args.skip_rename or not args.physical_rename:
print(f"[*] 物理重命名已跳过。仅保存 .info 及其可能包含的虚拟重命名。")
elif inferred:
print(f"[*] 没有 Civitai 的标准名称，保留原文件名。")
elif file_path != new_file_path and new_filename != filename:
if os.path.exists(new_file_path):
print(f"[*] 目标文件名已存在，正在验证内容是否完全相同: {filename}")
try:
files_identical = os.path.getsize(file_path) == os.path.getsize(new_file_path)
if files_identical:
files_identical = calculate_sha256(file_path) == calculate_sha256(new_file_path)
except OSError:
files_identical = False
if not files_identical:
print(f"\033[91m[-] 目标名称冲突但文件内容不同，已保留两个模型并跳过重命名: {filename}\033[0m")
return {**entry, "status": "failed", "error": "name_conflict", "renamed_to": new_file_path}
if not args.dry_run:
try:
print(f"[*] Hash 一致，把已确认的重复副本移到回收站: {filename}")
move_to_trash(file_path, *[path for path, _ in sidecars.existing(old_base, SIDECAR_SUFFIXES)])
entry["duplicate_of"] = new_file_path
except Exception as e:
print(f"[-] 重复副本没有移走，两份都保留 (可能文件被占用或这个盘没有回收站): {e}")
return {**entry, "status": "failed", "error": "duplicate_kept"}
else:
print(f"[Dry-Run] Hash 一致，拟删除重复副本及其附属文件: {filename}")
elif not args.dry_run:
os.rename(file_path, new_file_path)
rename_log[file_path] = new_file_path
for path, ext in sidecars.existing(old_base, SIDECAR_SUFFIXES):
os.replace(path, sidecars.follow(path, new_base, ext))
print(f"[+] 物理重命名完成: {filename} ==> {new_filename}")
entry["renamed_to"] = new_file_path
else:
print(f"[Dry-Run] 拟物理重命名文件: {filename} ==> {new_filename}")
print(f"[Dry-Run] 拟连带重命名附属文件 (.info / .png 等)")
else:
print("[*] 文件名已符合规范，无需重命名。")
changed = entry["status"] != "unchanged" or set(entry) & {"cover", "renamed_to", "duplicate_of"}
return entry if changed else None
def main():
parser = argparse.ArgumentParser(description="ComfyUI 模型 Civitai 嗅探与重命名工具")
parser.add_argument("folder", help="要扫描的文件夹路径 (例如: models/checkpoints)")
parser.add_argument("--dry-run", action="store_true", help="空跑模式，仅打印将要执行的操作，不修改任何文件")
parser.add_argument("--undo", action="store_true", help="根据 backup_rename_log.json 恢复文件名")
parser.add_argument("--skip-rename", action="store_true", help="只下载信息文件，不重命名主文件")
parser.add_argument("--virtual-rename", action="store_true", help="虚拟重命名：修改 JSON 注入标准名称，不修改底层物理文件名")
parser.add_argument("--physical-rename", action="store_true", help="物理重命名：真实修改底层的 safetensors 及其附属文件名")
parser.add_argument("--skip-media", action="store_true", help="不下载预览图或视频")
parser.add_argument("--offline-only", action="store_true", help="跳过 Civitai 联网获取，强制使用本地脱机张量推断提取 Base Model")
parser.add_argument("--force-overwrite", action="store_true", help="强制覆盖已存在的信息文件")
parser.add_argument("--retry-unmatched", action="store_true", help="只处理之前没在 Civitai 匹配到的模型，重新联网查找（文件没变时复用哈希，不替换封面）")
parser.add_argument("--skip-local-metadata", action="store_true", help="忽略本地已有的.info / .json文件")
parser.add_argument("--target-files", type=str, default="", help="仅扫描逗号分隔的具体文件(相对路径)")
parser.add_argument("--folder-type", default="", help="由 ComfyUI 传入的模型目录类型，用于执行安全策略")
parser.add_argument("--progress-file", default="", help=argparse.SUPPRESS)
parser.add_argument("--targets-file", default="", help=argparse.SUPPRESS) # JSON list of file names
parser.add_argument("--report-file", default="", help=argparse.SUPPRESS) # one JSON line per changed model
parser.add_argument("--civitai-down", action="store_true", help=argparse.SUPPRESS) # an earlier folder could not reach it
args = parser.parse_args()
if args.offline_only:
# Offline means no network at all: previews and model pages come from Civitai.
args.skip_media = True
target_folder = args.folder
if not os.path.isdir(target_folder):
print(f"[-] 错误: 文件夹不存在 -> {target_folder}")
sys.exit(1)
if args.physical_rename and is_physical_rename_protected(args.folder_type, target_folder):
print("[*] 基础组件目录受保护：VAE、CLIP/Text Encoder 与 CLIP Vision 不执行物理重命名。")
args.physical_rename = False
backup_log_path = os.path.join(target_folder, "backup_rename_log.json")
# ==========================
# 模式一：Undo 回滚模式
# ==========================
if args.undo:
if not os.path.exists(backup_log_path):
print("[-] 未找到备份日志 backup_rename_log.json，无法撤销。")
sys.exit(1)
with open(backup_log_path, 'r', encoding='utf-8') as f:
log_data = json.load(f)
print("[*] 开始回滚文件名...")
for old_path, new_path in log_data.items():
if os.path.exists(new_path):
print(f"[*] 恢复主文件: {os.path.basename(new_path)} -> {os.path.basename(old_path)}")
if not args.dry_run:
os.replace(new_path, old_path)
else:
print(f"[-] 找不到被重命名的文件: {new_path}")
old_base = os.path.splitext(old_path)[0]
new_base = os.path.splitext(new_path)[0]
for new_ext_path, ext in sidecars.existing(new_base, SIDECAR_SUFFIXES):
print(f"[*] 恢复配套文件: {os.path.basename(new_ext_path)} -> {os.path.basename(old_base + ext)}")
if not args.dry_run:
os.replace(new_ext_path, sidecars.follow(new_ext_path, old_base, ext))
print("[+] 回滚完成！")
sys.exit(0)
# ==========================
# 模式二：正常嗅探与重命名
# ==========================
rename_log = {}
if os.path.exists(backup_log_path):
with open(backup_log_path, 'r', encoding='utf-8') as f:
rename_log = json.load(f)
target_files_basenames = []
for targets_path in (os.path.join(target_folder, '.scan_targets.json'), args.targets_file):
if not targets_path or not os.path.exists(targets_path):
continue
try:
with open(targets_path, 'r', encoding='utf-8') as f:
target_files_basenames += [os.path.basename(t.strip()) for t in json.load(f)]
if targets_path != args.targets_file: # the scanner removes its own file
os.remove(targets_path)
except (OSError, ValueError, AttributeError):
pass
if args.target_files:
target_files_basenames.extend([os.path.basename(t.strip()) for t in args.target_files.split(',')])
print(f"[*] 开始扫描文件夹: {target_folder}")
if args.dry_run:
print("==================================================")
print("[警告]: 当前处于 Dry-Run (空跑) 模式，不会修改系统中的任何文件！")
print("==================================================")
if not CIVITAI_API_KEY and not args.offline_only:
print("[!] 未配置 Civitai API Key。部分限制级模型或将无法获取图片。")
write_scan_progress(args.progress_file, "enumerating")
total_files = count_scan_files(target_folder, target_files_basenames, args.retry_unmatched)
write_scan_progress(args.progress_file, "scanning", total_files)
current_file = 0
last_progress_write = 0.0
state = {"network_failures": 0, "civitai_down": args.civitai_down, "cover_failures": 0, "covers_down": False}
counts = {"success": 0, "fail": 0, "unchanged": 0}
for root, _, files in os.walk(target_folder):
for filename in files:
if not filename.endswith(".safetensors"):
continue
if target_files_basenames and filename not in target_files_basenames:
continue
if left_out(root, filename, target_files_basenames):
print(f"[*] 已设为不扫描，跳过: {filename}")
continue
previous = read_local_info(os.path.join(root, os.path.splitext(filename)[0]))
if args.retry_unmatched and not is_unmatched(previous):
continue
current_file += 1
now = time.monotonic()
if current_file == 1 or current_file == total_files or now - last_progress_write >= 0.2:
write_scan_progress(args.progress_file, "scanning", total_files, current_file, filename)
last_progress_write = now
try:
entry = scan_one(args, state, rename_log, root, filename, previous)
except Exception as e: # one model's trouble does not end the scan of the others
print(f"\033[91m[-] 处理失败，已跳过: {filename}: {e}\033[0m")
entry = {"status": "failed", "error": str(e)[:300]}
if entry is None:
counts["unchanged"] += 1
continue
counts["fail" if entry["status"] == "failed" else "success"] += 1
write_report(args, file=os.path.join(root, filename), **entry)
# Save scan results
write_scan_progress(args.progress_file, "complete", total_files, total_files)
write_report(args, event="done", unchanged=counts["unchanged"], civitai_down=state["civitai_down"])
if not args.dry_run:
result_path = os.path.join(target_folder, ".scan_result.json")
try:
with open(result_path, 'w', encoding='utf-8') as f:
json.dump({"success": counts["success"], "fail": counts["fail"]}, f)
except Exception as e:
print(f"[-] 保存统计结果失败: {e}")
if not args.dry_run and rename_log:
with open(backup_log_path, 'w', encoding='utf-8') as f:
json.dump(rename_log, f, ensure_ascii=True, indent=4)
print(f"\n[+] 重命名映射日志已保存至: {backup_log_path}")
if __name__ == "__main__":
main()