787 lines
36 KiB
Python
787 lines
36 KiB
Python
from .metadata import get_metadata
|
|
import os
|
|
import sys
|
|
import json
|
|
import urllib.parse
|
|
import subprocess
|
|
import threading
|
|
import asyncio
|
|
from aiohttp import web
|
|
import folder_paths
|
|
import struct
|
|
|
|
async def api_get_folders(request):
|
|
types = ['checkpoints', 'loras', 'unet', 'diffusion_models', 'controlnet', 'vae']
|
|
result = []
|
|
seen_dirs = set()
|
|
|
|
for t in types:
|
|
try:
|
|
paths = folder_paths.get_folder_paths(t)
|
|
except Exception:
|
|
continue
|
|
if not paths:
|
|
continue
|
|
|
|
for path_idx, base_dir in enumerate(paths):
|
|
if not os.path.exists(base_dir):
|
|
continue
|
|
real_path = os.path.realpath(base_dir)
|
|
if real_path in seen_dirs:
|
|
continue
|
|
seen_dirs.add(real_path)
|
|
|
|
tree = {}
|
|
for root, dirs, files in os.walk(base_dir):
|
|
has_models = any(f.endswith(('.safetensors', '.ckpt', '.pt')) for f in files)
|
|
rel = os.path.relpath(root, base_dir)
|
|
if rel == '.':
|
|
rel = '/'
|
|
else:
|
|
rel = '/' + rel.replace('\\', '/')
|
|
tree[rel] = {
|
|
"path": rel,
|
|
"name": os.path.basename(root) if rel != '/' else '[Root]',
|
|
"has_models": has_models,
|
|
"model_count": sum(1 for f in files if f.endswith(('.safetensors', '.ckpt', '.pt')))
|
|
}
|
|
|
|
label = t.capitalize()
|
|
if path_idx > 0:
|
|
label += f" ({os.path.basename(base_dir)})"
|
|
|
|
result.append({
|
|
"type": t,
|
|
"path_idx": path_idx,
|
|
"label": label,
|
|
"folders": tree
|
|
})
|
|
|
|
return web.json_response({"folders": result})
|
|
|
|
|
|
async def api_get_models(request):
|
|
folder_type = request.query.get('type', 'checkpoints')
|
|
subfolder = request.query.get('subfolder', '/')
|
|
try:
|
|
path_idx = int(request.query.get('path_idx', 0))
|
|
except:
|
|
path_idx = 0
|
|
|
|
try:
|
|
paths = folder_paths.get_folder_paths(folder_type)
|
|
except Exception:
|
|
return web.json_response({"models": []})
|
|
if not paths or path_idx >= len(paths):
|
|
return web.json_response({"models": []})
|
|
|
|
if '..' in subfolder:
|
|
return web.Response(status=400, text='Invalid subfolder')
|
|
|
|
base_dir = paths[path_idx]
|
|
if subfolder == '/':
|
|
target_dir = base_dir
|
|
rel_subfolder = ""
|
|
else:
|
|
rel_subfolder = subfolder.strip('/')
|
|
target_dir = os.path.join(base_dir, rel_subfolder)
|
|
|
|
if not os.path.exists(target_dir):
|
|
return web.json_response({"models": []})
|
|
|
|
models = []
|
|
try:
|
|
entries = os.listdir(target_dir)
|
|
except Exception:
|
|
return web.json_response({"models": []})
|
|
|
|
for f in sorted(entries):
|
|
if f.endswith('.safetensors') or f.endswith('.ckpt') or f.endswith('.pt'):
|
|
file_path = os.path.join(target_dir, f)
|
|
if not os.path.isfile(file_path):
|
|
continue
|
|
|
|
meta = get_metadata(file_path)
|
|
|
|
base_name = os.path.splitext(f)[0]
|
|
preview_file = None
|
|
for ext in ['.preview.png', '.preview.jpg', '.preview.jpeg', '.preview.webp', '.preview.gif', '.preview.avif', '.preview.mp4', '.preview.webm', '.preview.mov', '.preview.avi', '.png', '.jpg', '.jpeg', '.webp', '.gif', '.avif', '.mp4', '.webm', '.mov', '.avi']:
|
|
if os.path.exists(os.path.join(target_dir, base_name + ext)):
|
|
preview_file = base_name + ext
|
|
break
|
|
|
|
preview_url = ""
|
|
if preview_file:
|
|
q_type = urllib.parse.quote(folder_type)
|
|
q_idx = str(path_idx)
|
|
q_sub = urllib.parse.quote(rel_subfolder)
|
|
q_file = urllib.parse.quote(preview_file)
|
|
try:
|
|
mtime = int(os.path.getmtime(os.path.join(target_dir, preview_file)))
|
|
except:
|
|
mtime = 0
|
|
preview_url = f"/anomalous/image?type={q_type}&path_idx={q_idx}&subfolder={q_sub}&filename={q_file}&t={mtime}"
|
|
|
|
try:
|
|
size_bytes = os.path.getsize(file_path)
|
|
size_mb = round(size_bytes / (1024 * 1024), 1)
|
|
except Exception:
|
|
size_mb = 0
|
|
size_bytes = 0
|
|
|
|
models.append({
|
|
"filename": f,
|
|
"size_mb": size_mb,
|
|
"size_bytes": size_bytes,
|
|
"metadata": meta,
|
|
"preview_url": preview_url
|
|
})
|
|
|
|
return web.json_response({"models": models})
|
|
|
|
|
|
async def api_find_model(request):
|
|
search = request.query.get('search', '').lower()
|
|
if not search:
|
|
return web.json_response({"status": "error", "message": "No search query provided"})
|
|
|
|
for folder_type in folder_paths.folder_names_and_paths.keys():
|
|
paths = folder_paths.get_folder_paths(folder_type)
|
|
if not paths: continue
|
|
for path_idx, base_dir in enumerate(paths):
|
|
if not os.path.exists(base_dir): continue
|
|
for root, dirs, files in os.walk(base_dir):
|
|
for f in files:
|
|
if f.endswith('.safetensors') or f.endswith('.ckpt') or f.endswith('.pt') or f.endswith('.bin') or f.endswith('.sft'):
|
|
search_norm = search.replace(os.sep, '/')
|
|
abs_path_norm = os.path.join(root, f).replace(os.sep, '/').lower()
|
|
if search_norm in f.lower() or search_norm in abs_path_norm:
|
|
rel_subfolder = os.path.relpath(root, base_dir)
|
|
if rel_subfolder == '.': rel_subfolder = '/'
|
|
|
|
file_path = os.path.join(root, f)
|
|
meta = get_metadata(file_path)
|
|
|
|
base_name = os.path.splitext(f)[0]
|
|
preview_file = None
|
|
for ext in ['.preview.png', '.preview.jpg', '.preview.jpeg', '.preview.webp', '.preview.gif', '.preview.avif', '.preview.mp4', '.preview.webm', '.preview.mov', '.preview.avi', '.png', '.jpg', '.jpeg', '.webp', '.gif', '.avif', '.mp4', '.webm', '.mov', '.avi']:
|
|
if os.path.exists(os.path.join(root, base_name + ext)):
|
|
preview_file = base_name + ext
|
|
break
|
|
|
|
preview_url = ""
|
|
if preview_file:
|
|
q_type = urllib.parse.quote(folder_type)
|
|
q_idx = str(path_idx)
|
|
q_sub = urllib.parse.quote(rel_subfolder)
|
|
q_file = urllib.parse.quote(preview_file)
|
|
try: mtime = int(os.path.getmtime(os.path.join(root, preview_file)))
|
|
except: mtime = 0
|
|
preview_url = f"/anomalous/image?type={q_type}&path_idx={q_idx}&subfolder={q_sub}&filename={q_file}&t={mtime}"
|
|
|
|
try: size_mb = round(os.path.getsize(file_path) / (1024 * 1024), 1)
|
|
except: size_mb = 0
|
|
|
|
modelData = {
|
|
"filename": f,
|
|
"size_mb": size_mb,
|
|
"metadata": meta,
|
|
"preview_url": preview_url
|
|
}
|
|
return web.json_response({
|
|
"status": "success",
|
|
"model": modelData,
|
|
"type": folder_type,
|
|
"path_idx": path_idx,
|
|
"subfolder": rel_subfolder
|
|
})
|
|
|
|
return web.json_response({"status": "error", "message": "Model not found"})
|
|
|
|
async def api_delete_model(request):
|
|
try:
|
|
data = await request.json()
|
|
folder_type = data.get('type', 'checkpoints')
|
|
subfolder = data.get('subfolder', '/')
|
|
filename = data.get('filename', '')
|
|
try:
|
|
path_idx = int(data.get('path_idx', 0))
|
|
except:
|
|
path_idx = 0
|
|
|
|
if not filename or '..' in filename or '..' in subfolder:
|
|
return web.json_response({"status": "error", "message": f"Invalid request parameters. filename={repr(filename)}, subfolder={repr(subfolder)}"})
|
|
|
|
try:
|
|
paths = folder_paths.get_folder_paths(folder_type)
|
|
except Exception:
|
|
return web.json_response({"status": "error", "message": "Invalid folder type"})
|
|
|
|
if not paths or path_idx >= len(paths):
|
|
return web.json_response({"status": "error", "message": "Invalid path index"})
|
|
|
|
base_dir = paths[path_idx]
|
|
if subfolder == '/':
|
|
target_dir = base_dir
|
|
else:
|
|
target_dir = os.path.join(base_dir, subfolder.strip('/'))
|
|
|
|
model_path = os.path.join(target_dir, filename)
|
|
if not os.path.exists(model_path):
|
|
return web.json_response({"status": "error", "message": "Model file not found"})
|
|
|
|
# 1. 优先尝试删除你点击的主模型文件
|
|
try:
|
|
os.remove(model_path)
|
|
except Exception as e:
|
|
error_msg = str(e)
|
|
if "being used" in error_msg or "WinError 32" in error_msg or "Permission" in error_msg:
|
|
error_msg = "文件被占用 (正在被 ComfyUI 使用)。请先重启 ComfyUI 或在工作流中卸载该模型后再删除！"
|
|
return web.json_response({"status": "error", "message": f"主模型删除失败: {error_msg}"})
|
|
|
|
base_name = os.path.splitext(filename)[0]
|
|
|
|
# 2. 主模型成功删除后，再清理配套的垃圾文件
|
|
associated_exts = [
|
|
'.safetensors', '.ckpt', '.pt', '.bin',
|
|
'.info', '.civitai.info',
|
|
'.png', '.preview.png', '.jpg', '.jpeg', '.webp', '.gif',
|
|
'.json', '.txt', '.yaml'
|
|
]
|
|
|
|
deleted_files = [filename]
|
|
for ext in associated_exts:
|
|
file_to_del = os.path.join(target_dir, base_name + ext)
|
|
if file_to_del == model_path:
|
|
continue
|
|
if os.path.exists(file_to_del) and os.path.isfile(file_to_del):
|
|
try:
|
|
os.remove(file_to_del)
|
|
deleted_files.append(base_name + ext)
|
|
except Exception as e:
|
|
print(f"[Anomalous Browser] Warning: Failed to delete {file_to_del}: {e}")
|
|
|
|
# 3. 修正：前端期待的成功状态是 "success" 而不是 "ok"
|
|
return web.json_response({"status": "success", "deleted": deleted_files})
|
|
|
|
except Exception as e:
|
|
import traceback
|
|
traceback.print_exc()
|
|
return web.json_response({"status": "error", "message": str(e)}, status=500)
|
|
|
|
async def api_compatible_models(request):
|
|
base_model = request.query.get('base_model', '')
|
|
target_type = request.query.get('target_type', 'loras')
|
|
|
|
if not base_model:
|
|
return web.json_response({"models": []})
|
|
|
|
target_types = [t.strip() for t in target_type.split(',')]
|
|
compatible_models = []
|
|
seen_files = set()
|
|
|
|
for t in target_types:
|
|
try:
|
|
paths = folder_paths.get_folder_paths(t)
|
|
except Exception:
|
|
continue
|
|
|
|
if not paths:
|
|
continue
|
|
|
|
for path_idx, base_dir in enumerate(paths):
|
|
if not os.path.exists(base_dir):
|
|
continue
|
|
|
|
for root, _, files in os.walk(base_dir):
|
|
for f in files:
|
|
if f.endswith('.safetensors') or f.endswith('.ckpt') or f.endswith('.pt'):
|
|
file_path = os.path.join(root, f)
|
|
real_path = os.path.realpath(file_path)
|
|
if real_path in seen_files:
|
|
continue
|
|
seen_files.add(real_path)
|
|
|
|
meta = get_metadata(file_path)
|
|
m_bm = str(meta.get("baseModel", "")).strip().lower().replace(" ", "")
|
|
req_bm = str(base_model).strip().lower().replace(" ", "")
|
|
|
|
if req_bm and m_bm and (req_bm in m_bm or m_bm in req_bm):
|
|
rel_subfolder = os.path.relpath(root, base_dir)
|
|
if rel_subfolder == '.':
|
|
rel_subfolder = '/'
|
|
else:
|
|
rel_subfolder = '/' + rel_subfolder.replace('\\', '/')
|
|
|
|
base_name = os.path.splitext(f)[0]
|
|
preview_file = None
|
|
for ext in ['.preview.png', '.preview.jpg', '.preview.jpeg', '.preview.webp', '.preview.gif', '.preview.avif', '.preview.mp4', '.preview.webm', '.preview.mov', '.preview.avi', '.png', '.jpg', '.jpeg', '.webp', '.gif', '.avif', '.mp4', '.webm', '.mov', '.avi']:
|
|
if os.path.exists(os.path.join(root, base_name + ext)):
|
|
preview_file = base_name + ext
|
|
break
|
|
|
|
preview_url = ""
|
|
if preview_file:
|
|
q_type = urllib.parse.quote(t)
|
|
q_idx = str(path_idx)
|
|
q_sub = urllib.parse.quote(rel_subfolder.strip('/')) if rel_subfolder != '/' else ""
|
|
q_file = urllib.parse.quote(preview_file)
|
|
preview_url = f"/anomalous/image?type={q_type}&path_idx={q_idx}&subfolder={q_sub}&filename={q_file}"
|
|
|
|
try:
|
|
size_bytes = os.path.getsize(file_path)
|
|
size_mb = round(size_bytes / (1024 * 1024), 1)
|
|
except Exception:
|
|
size_mb = 0
|
|
size_bytes = 0
|
|
|
|
compatible_models.append({
|
|
"type": t,
|
|
"path_idx": path_idx,
|
|
"subfolder": rel_subfolder,
|
|
"filename": f,
|
|
"size_mb": size_mb,
|
|
"size_bytes": size_bytes,
|
|
"preview_url": preview_url,
|
|
"metadata": meta
|
|
})
|
|
|
|
return web.json_response({"models": compatible_models})
|
|
|
|
async def api_base_models(request):
|
|
target_types = ['checkpoints', 'loras', 'unet', 'diffusion_models', 'controlnet', 'vae']
|
|
base_models = set()
|
|
seen_files = set()
|
|
|
|
for t in target_types:
|
|
try:
|
|
paths = folder_paths.get_folder_paths(t)
|
|
except Exception:
|
|
continue
|
|
if not paths: continue
|
|
|
|
for base_dir in paths:
|
|
if not os.path.exists(base_dir): continue
|
|
for root, _, files in os.walk(base_dir):
|
|
for f in files:
|
|
if f.endswith('.safetensors') or f.endswith('.ckpt') or f.endswith('.pt'):
|
|
file_path = os.path.join(root, f)
|
|
real_path = os.path.realpath(file_path)
|
|
if real_path in seen_files: continue
|
|
seen_files.add(real_path)
|
|
|
|
meta = get_metadata(file_path)
|
|
m_bm = meta.get("baseModel", "")
|
|
if m_bm and str(m_bm).strip():
|
|
# Remove typical generic strings that might pollute
|
|
clean_bm = str(m_bm).strip()
|
|
base_models.add(clean_bm)
|
|
|
|
return web.json_response({"base_models": sorted(list(base_models))})
|
|
|
|
async def api_resolve_hash(request):
|
|
from .metadata import _extract_safetensors_hash
|
|
target_hash = request.query.get("hash", "").strip().upper()
|
|
size_str = request.query.get("size", "").strip()
|
|
filename_query = request.query.get("filename", "").strip()
|
|
target_size = int(size_str) if size_str.isdigit() else None
|
|
|
|
if not target_hash and not target_size and not filename_query:
|
|
return web.json_response({"found": False})
|
|
|
|
types = ['checkpoints', 'loras', 'unet', 'diffusion_models', 'controlnet', 'vae']
|
|
|
|
# Layer 1: Fast Size Match (O(1) filesystem check)
|
|
if target_size is not None:
|
|
size_matches = []
|
|
seen_realpaths = set()
|
|
for t in types:
|
|
try:
|
|
paths = folder_paths.get_folder_paths(t)
|
|
except Exception:
|
|
continue
|
|
if not paths:
|
|
continue
|
|
|
|
for base_dir in paths:
|
|
if not os.path.exists(base_dir):
|
|
continue
|
|
|
|
for root, dirs, files in os.walk(base_dir):
|
|
for file in files:
|
|
if file.endswith('.safetensors') or file.endswith('.ckpt') or file.endswith('.pt'):
|
|
file_path = os.path.join(root, file)
|
|
try:
|
|
if os.path.getsize(file_path) == target_size:
|
|
real_p = os.path.realpath(file_path)
|
|
if real_p in seen_realpaths:
|
|
continue
|
|
seen_realpaths.add(real_p)
|
|
|
|
rel_path = os.path.relpath(file_path, base_dir)
|
|
if rel_path.startswith('.\\') or rel_path.startswith('./'):
|
|
rel_path = rel_path[2:]
|
|
rel_path = rel_path.replace('\\', '/')
|
|
size_matches.append({
|
|
"type": t,
|
|
"filename": rel_path
|
|
})
|
|
except Exception:
|
|
pass
|
|
|
|
# If exactly ONE file matches the size, it's a definitive match
|
|
if len(size_matches) == 1:
|
|
return web.json_response({
|
|
"found": True,
|
|
"type": size_matches[0]["type"],
|
|
"filename": size_matches[0]["filename"],
|
|
"matched_by_size": True
|
|
})
|
|
|
|
# Layer 2: Metadata Hash Match (Reads .info, falls back to safetensors header)
|
|
if target_hash and target_hash != "UNKNOWN":
|
|
for t in types:
|
|
try:
|
|
paths = folder_paths.get_folder_paths(t)
|
|
except Exception:
|
|
continue
|
|
if not paths:
|
|
continue
|
|
|
|
for base_dir in paths:
|
|
if not os.path.exists(base_dir):
|
|
continue
|
|
|
|
for root, dirs, files in os.walk(base_dir):
|
|
for file in files:
|
|
if file.endswith('.safetensors') or file.endswith('.ckpt') or file.endswith('.pt'):
|
|
file_path = os.path.join(root, file)
|
|
meta = get_metadata(file_path)
|
|
|
|
# Check .info or header hash
|
|
if meta.get("hash", "").upper() == target_hash:
|
|
rel_path = os.path.relpath(file_path, base_dir)
|
|
if rel_path.startswith('.\\') or rel_path.startswith('./'):
|
|
rel_path = rel_path[2:]
|
|
rel_path = rel_path.replace('\\', '/')
|
|
return web.json_response({
|
|
"found": True,
|
|
"type": t,
|
|
"filename": rel_path
|
|
})
|
|
|
|
# Additional direct check just in case get_metadata missed it
|
|
if file.endswith('.safetensors'):
|
|
header_hash = _extract_safetensors_hash(file_path)
|
|
if header_hash and header_hash.upper() == target_hash:
|
|
rel_path = os.path.relpath(file_path, base_dir)
|
|
if rel_path.startswith('.\\') or rel_path.startswith('./'):
|
|
rel_path = rel_path[2:]
|
|
rel_path = rel_path.replace('\\', '/')
|
|
return web.json_response({
|
|
"found": True,
|
|
"type": t,
|
|
"filename": rel_path,
|
|
"matched_by_header": True
|
|
})
|
|
|
|
return web.json_response({"found": False})
|
|
|
|
async def api_get_all_hashes(request):
|
|
"""
|
|
Returns a dictionary of all scanned models with their hash and size.
|
|
Keyed by both relative path and basename for maximum frontend resilience.
|
|
"""
|
|
import asyncio
|
|
|
|
def fetch_all():
|
|
hashes = {}
|
|
types = ['checkpoints', 'loras', 'unet', 'diffusion_models', 'controlnet', 'vae']
|
|
for t in types:
|
|
try:
|
|
paths = folder_paths.get_folder_paths(t)
|
|
if not paths: continue
|
|
for base_dir in paths:
|
|
if not os.path.exists(base_dir): continue
|
|
for root, dirs, files in os.walk(base_dir):
|
|
for file in files:
|
|
if file.endswith('.safetensors') or file.endswith('.ckpt') or file.endswith('.pt'):
|
|
file_path = os.path.join(root, file)
|
|
try:
|
|
size_bytes = os.path.getsize(file_path)
|
|
except Exception:
|
|
size_bytes = 0
|
|
|
|
meta = get_metadata(file_path)
|
|
hash_val = ""
|
|
if meta and meta.get("hash"):
|
|
hash_val = meta["hash"]
|
|
|
|
rel_path = os.path.relpath(file_path, base_dir)
|
|
if rel_path.startswith('.\\') or rel_path.startswith('./'):
|
|
rel_path = rel_path[2:]
|
|
rel_path = rel_path.replace('\\', '/')
|
|
basename = os.path.basename(file_path)
|
|
|
|
val = {"hash": hash_val, "size": size_bytes}
|
|
hashes[rel_path] = val
|
|
hashes[basename] = val
|
|
except Exception:
|
|
pass
|
|
return hashes
|
|
|
|
hashes = await asyncio.to_thread(fetch_all)
|
|
return web.json_response(hashes)
|
|
|
|
async def api_update_metadata(request):
|
|
try:
|
|
data = await request.json()
|
|
folder_type = data.get('type', 'checkpoints')
|
|
subfolder = data.get('subfolder', '/')
|
|
filename = data.get('filename', '')
|
|
custom_name = data.get('custom_name', '')
|
|
custom_notes = data.get('custom_notes', '')
|
|
physical_rename = data.get('physical_rename', False)
|
|
try: path_idx = int(data.get('path_idx', 0))
|
|
except: path_idx = 0
|
|
|
|
if not filename or '..' in filename or '..' in subfolder:
|
|
return web.json_response({"status": "error", "message": f"Invalid request parameters. filename={repr(filename)}, subfolder={repr(subfolder)}"})
|
|
|
|
paths = folder_paths.get_folder_paths(folder_type)
|
|
if not paths or path_idx >= len(paths):
|
|
return web.json_response({"status": "error", "message": "Invalid path index"})
|
|
|
|
base_dir = paths[path_idx]
|
|
target_dir = base_dir if subfolder == '/' else os.path.join(base_dir, subfolder.strip('/'))
|
|
file_path = os.path.join(target_dir, filename)
|
|
|
|
if not os.path.exists(file_path):
|
|
return web.json_response({"status": "error", "message": "Model not found"})
|
|
|
|
base_name = os.path.splitext(file_path)[0]
|
|
ext = os.path.splitext(file_path)[1]
|
|
info_file = f"{base_name}.civitai.info"
|
|
|
|
info_data = {}
|
|
if os.path.exists(info_file):
|
|
parsed = False
|
|
for enc in ['utf-8', 'utf-8-sig', 'mbcs', 'latin-1']:
|
|
try:
|
|
with open(info_file, 'r', encoding=enc) as f:
|
|
info_data = json.load(f)
|
|
parsed = True
|
|
break
|
|
except Exception:
|
|
pass
|
|
if not parsed:
|
|
return web.json_response({"status": "error", "message": "Failed to parse existing .civitai.info file due to encoding or corruption. Rename aborted to prevent data loss."})
|
|
|
|
info_data["anomalous_custom_name"] = custom_name
|
|
info_data["anomalous_custom_notes"] = custom_notes
|
|
|
|
reset_cover = data.get('reset_cover', False)
|
|
if reset_cover:
|
|
# 1. Delete any custom active cover (.preview.*)
|
|
for ext in ['.preview.png', '.preview.jpg', '.preview.jpeg', '.preview.webp', '.preview.gif', '.preview.avif', '.preview.mp4', '.preview.webm', '.preview.mov', '.preview.avi']:
|
|
p = os.path.join(target_dir, f"{base_name}{ext}")
|
|
if os.path.exists(p):
|
|
try: os.remove(p)
|
|
except: pass
|
|
|
|
# 2. If a Civitai backup exists, copy it back to .preview.* so standard ComfyUI can see it
|
|
for ext in ['.png', '.jpg', '.jpeg', '.webp', '.gif', '.avif', '.mp4', '.webm', '.mov', '.avi']:
|
|
civitai_bak = os.path.join(target_dir, f"{base_name}.civitai_bak{ext}")
|
|
if os.path.exists(civitai_bak):
|
|
restored_p = os.path.join(target_dir, f"{base_name}.preview{ext}")
|
|
import shutil
|
|
try: shutil.copy2(civitai_bak, restored_p)
|
|
except: pass
|
|
|
|
with open(info_file, 'w', encoding='utf-8') as f:
|
|
json.dump(info_data, f, indent=4, ensure_ascii=False)
|
|
|
|
new_filename = filename
|
|
|
|
if physical_rename and custom_name:
|
|
import re
|
|
safe_name = re.sub(r'[<>:"/\\|?*]', '_', custom_name)
|
|
new_file_path = os.path.join(target_dir, f"{safe_name}{ext}")
|
|
|
|
if new_file_path != file_path and not os.path.exists(new_file_path):
|
|
os.rename(file_path, new_file_path)
|
|
|
|
new_info = os.path.join(target_dir, f"{safe_name}.civitai.info")
|
|
os.rename(info_file, new_info)
|
|
|
|
preview_png = f"{base_name}.preview.png"
|
|
if os.path.exists(preview_png):
|
|
os.rename(preview_png, os.path.join(target_dir, f"{safe_name}.preview.png"))
|
|
|
|
civitai_bak = f"{base_name}.civitai_bak.png"
|
|
if os.path.exists(civitai_bak):
|
|
os.rename(civitai_bak, os.path.join(target_dir, f"{safe_name}.civitai_bak.png"))
|
|
|
|
new_filename = f"{safe_name}{ext}"
|
|
elif os.path.exists(new_file_path) and new_file_path != file_path:
|
|
return web.json_response({"status": "error", "message": "A file with the target physical name already exists."})
|
|
|
|
return web.json_response({"status": "success", "new_filename": new_filename})
|
|
except Exception as e:
|
|
return web.json_response({"status": "error", "message": str(e)})
|
|
|
|
async def _handle_custom_cover(target_dir, filename, save_func, source_ext='.png'):
|
|
base_name = os.path.splitext(filename)[0]
|
|
|
|
# Always save custom covers with a .preview.[ext] suffix so standard nodes recognize them as covers.
|
|
if source_ext.startswith('.preview.'):
|
|
preview_ext = source_ext
|
|
else:
|
|
preview_ext = f".preview{source_ext}"
|
|
|
|
dest_path = os.path.join(target_dir, f"{base_name}{preview_ext}")
|
|
|
|
# Delete any existing .preview.* files to ensure only one custom cover is active
|
|
for ext in ['.preview.png', '.preview.jpg', '.preview.jpeg', '.preview.webp', '.preview.gif', '.preview.avif', '.preview.mp4', '.preview.webm', '.preview.mov', '.preview.avi']:
|
|
p = os.path.join(target_dir, f"{base_name}{ext}")
|
|
if os.path.exists(p) and p != dest_path:
|
|
try: os.remove(p)
|
|
except: pass
|
|
|
|
await save_func(dest_path)
|
|
|
|
async def api_set_custom_cover(request):
|
|
try:
|
|
data = await request.json()
|
|
folder_type = data.get('type', 'checkpoints')
|
|
subfolder = data.get('subfolder', '/')
|
|
filename = data.get('filename', '')
|
|
source_image = data.get('source_image', '')
|
|
try: path_idx = int(data.get('path_idx', 0))
|
|
except: path_idx = 0
|
|
|
|
if not filename or not source_image or '..' in filename or '..' in subfolder:
|
|
return web.json_response({"status": "error", "message": f"Invalid request parameters. filename={repr(filename)}, subfolder={repr(subfolder)}"})
|
|
|
|
paths = folder_paths.get_folder_paths(folder_type)
|
|
if not paths or path_idx >= len(paths):
|
|
return web.json_response({"status": "error", "message": "Invalid path index"})
|
|
|
|
target_dir = paths[path_idx] if subfolder == '/' else os.path.join(paths[path_idx], subfolder.strip('/'))
|
|
|
|
# Resolve source_image
|
|
output_dir = folder_paths.get_output_directory()
|
|
src_path = os.path.join(output_dir, source_image)
|
|
if not os.path.exists(src_path):
|
|
return web.json_response({"status": "error", "message": "Source image not found in output directory"})
|
|
|
|
source_ext = os.path.splitext(src_path)[1].lower()
|
|
|
|
async def save_copy(dest_path):
|
|
import shutil
|
|
import asyncio
|
|
await asyncio.to_thread(shutil.copy2, src_path, dest_path)
|
|
|
|
await _handle_custom_cover(target_dir, filename, save_copy, source_ext)
|
|
|
|
return web.json_response({"status": "success"})
|
|
except Exception as e:
|
|
return web.json_response({"status": "error", "message": str(e)})
|
|
|
|
async def api_upload_custom_cover(request):
|
|
try:
|
|
data = await request.post()
|
|
folder_type = data.get('type', 'checkpoints')
|
|
subfolder = data.get('subfolder', '/')
|
|
filename = data.get('filename', '')
|
|
try: path_idx = int(data.get('path_idx', 0))
|
|
except: path_idx = 0
|
|
|
|
image_field = data.get('image')
|
|
|
|
if not filename or '..' in filename or '..' in subfolder or image_field is None:
|
|
return web.json_response({"status": "error", "message": f"Invalid request parameters"})
|
|
|
|
paths = folder_paths.get_folder_paths(folder_type)
|
|
if not paths or path_idx >= len(paths):
|
|
return web.json_response({"status": "error", "message": "Invalid path index"})
|
|
|
|
target_dir = paths[path_idx] if subfolder == '/' else os.path.join(paths[path_idx], subfolder.strip('/'))
|
|
|
|
image_data = image_field.file.read()
|
|
|
|
upload_filename = image_field.filename
|
|
source_ext = os.path.splitext(upload_filename)[1].lower()
|
|
if not source_ext:
|
|
source_ext = '.png'
|
|
|
|
async def save_upload(dest_path):
|
|
def write_file():
|
|
with open(dest_path, 'wb') as f:
|
|
f.write(image_data)
|
|
import asyncio
|
|
await asyncio.to_thread(write_file)
|
|
|
|
await _handle_custom_cover(target_dir, filename, save_upload, source_ext)
|
|
|
|
return web.json_response({"status": "success"})
|
|
except Exception as e:
|
|
return web.json_response({"status": "error", "message": str(e)})
|
|
|
|
|
|
import struct
|
|
|
|
async def api_resolve_paths_to_previews(request):
|
|
try:
|
|
data = await request.json()
|
|
paths = data.get('paths', [])
|
|
except:
|
|
return web.json_response({"previews": {}})
|
|
|
|
previews = {}
|
|
all_models = {}
|
|
|
|
for folder_type in folder_paths.folder_names_and_paths.keys():
|
|
folder_dirs = folder_paths.get_folder_paths(folder_type)
|
|
if not folder_dirs: continue
|
|
for path_idx, base_dir in enumerate(folder_dirs):
|
|
if not os.path.exists(base_dir): continue
|
|
for root, dirs, files in os.walk(base_dir):
|
|
for f in files:
|
|
if f.endswith('.safetensors') or f.endswith('.ckpt') or f.endswith('.pt') or f.endswith('.bin') or f.endswith('.sft'):
|
|
rel_path = os.path.relpath(os.path.join(root, f), base_dir).replace(os.sep, '/')
|
|
base_name = os.path.splitext(f)[0]
|
|
preview_file = None
|
|
for ext in ['.preview.png', '.preview.jpg', '.preview.jpeg', '.preview.webp', '.preview.gif', '.preview.avif', '.preview.mp4', '.preview.webm', '.preview.mov', '.preview.avi', '.png', '.jpg', '.jpeg', '.webp', '.gif', '.avif', '.mp4', '.webm', '.mov', '.avi']:
|
|
if os.path.exists(os.path.join(root, base_name + ext)):
|
|
preview_file = base_name + ext
|
|
break
|
|
|
|
preview_url = ""
|
|
if preview_file:
|
|
rel_subfolder = os.path.relpath(root, base_dir)
|
|
if rel_subfolder == '.': rel_subfolder = '/'
|
|
q_type = urllib.parse.quote(folder_type)
|
|
q_idx = str(path_idx)
|
|
q_sub = urllib.parse.quote(rel_subfolder)
|
|
q_file = urllib.parse.quote(preview_file)
|
|
try: mtime = int(os.path.getmtime(os.path.join(root, preview_file)))
|
|
except: mtime = 0
|
|
preview_url = f"/anomalous/image?type={q_type}&path_idx={q_idx}&subfolder={q_sub}&filename={q_file}&t={mtime}"
|
|
|
|
all_models[rel_path.lower()] = preview_url
|
|
all_models[f.lower()] = preview_url
|
|
|
|
for p in paths:
|
|
p_norm = p.replace(os.sep, '/').lower()
|
|
if p_norm in all_models:
|
|
previews[p] = all_models[p_norm]
|
|
else:
|
|
fname = p.split('/')[-1].split('\\')[-1].lower()
|
|
if fname in all_models:
|
|
previews[p] = all_models[fname]
|
|
|
|
return web.json_response({"previews": previews})
|
|
|
|
|