- header_words.py: modelspec.trigger_phrase, else the tags most training pictures share (ss_tag_frequency, generic ones left out); the scanner keeps them as anomalous_header_words when Civitai gives no trigger words, metadata gives headerWords, model search finds them. - model_kind: Flux Schnell told from Dev (no guidance_in), Z-Image from Lumina 2 (caption embedder width, a GGUF by parameter count); RULES 3 so kept header answers are read again. - api/model_guide.py + model_profiles.json: GET /anomalous/model_guide for the model details (example settings, header words, the text encoder and VAE a diffusion models file needs, found by name), and get_model's header_words and usage for the AI. - model_advice.example_settings gives the author's example settings as values; author_examples says them in words as before. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
499 lines
26 KiB
Python
499 lines
26 KiB
Python
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
|
|
import header_words
|
|
|
|
|
|
# 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 (<base>.<ext>) 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": "<p>Automatically inferred by Anomalous Local Engine.</p>",
|
|
"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)
|
|
# A LoRA Civitai gives no trigger words for: what its own header suggests, kept apart.
|
|
info_data.pop("anomalous_header_words", None)
|
|
if not any(str(word).strip() for word in info_data.get("trainedWords") or []):
|
|
words = header_words.of_file(file_path)
|
|
if words:
|
|
info_data["anomalous_header_words"] = words
|
|
# 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
|
|
# <model>.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()
|