424 lines
16 KiB
Python
424 lines
16 KiB
Python
import os
|
|
import sys
|
|
import json
|
|
import urllib.parse
|
|
import subprocess
|
|
import threading
|
|
from aiohttp import web
|
|
import folder_paths
|
|
|
|
def get_metadata(file_path):
|
|
base_path = os.path.splitext(file_path)[0]
|
|
metadata = {
|
|
"name": os.path.basename(base_path),
|
|
"description": "",
|
|
"notes": "",
|
|
"trainedWords": [],
|
|
"baseModel": "",
|
|
"civitai_url": ""
|
|
}
|
|
|
|
info_files = [f"{base_path}.civitai.info", f"{base_path}.info"]
|
|
for info_file in info_files:
|
|
if os.path.exists(info_file):
|
|
try:
|
|
with open(info_file, 'r', encoding='utf-8') as f:
|
|
data = json.load(f)
|
|
|
|
name = data.get("name", "")
|
|
if "model" in data and isinstance(data["model"], dict):
|
|
model_name = data["model"].get("name", "")
|
|
name = f"{model_name} - {name}".strip(' -')
|
|
|
|
description = data.get("description", "") or ""
|
|
if not description and "model" in data and isinstance(data["model"], dict):
|
|
description = data["model"].get("description", "")
|
|
notes = data.get("notes", "")
|
|
trained_words = data.get("trainedWords", [])
|
|
base_model = data.get("baseModel", "")
|
|
|
|
model_id = data.get("modelId", "")
|
|
version_id = data.get("id", "")
|
|
civitai_url = ""
|
|
if model_id:
|
|
# Handle Civitai's new mature content policy
|
|
nsfw_level = data.get("nsfwLevel", 1)
|
|
is_nsfw = False
|
|
if "model" in data and isinstance(data["model"], dict):
|
|
is_nsfw = data["model"].get("nsfw", False)
|
|
|
|
domain = "civitai.red" if (nsfw_level > 1 or is_nsfw) else "civitai.com"
|
|
civitai_url = f"https://{domain}/models/{model_id}"
|
|
if version_id:
|
|
civitai_url += f"?modelVersionId={version_id}"
|
|
|
|
metadata["name"] = name if name else metadata["name"]
|
|
metadata["description"] = description
|
|
metadata["notes"] = notes
|
|
metadata["trainedWords"] = trained_words
|
|
metadata["baseModel"] = base_model
|
|
metadata["civitai_url"] = civitai_url
|
|
break
|
|
except Exception:
|
|
pass
|
|
|
|
return metadata
|
|
|
|
|
|
async def api_get_folders(request):
|
|
types = ['checkpoints', 'loras', 'unet', 'diffusion_models']
|
|
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', '.png', '.jpg', '.jpeg', '.webp', '.mp4', '.webm']:
|
|
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)
|
|
preview_url = f"/anomalous/image?type={q_type}&path_idx={q_idx}&subfolder={q_sub}&filename={q_file}"
|
|
|
|
try:
|
|
size_mb = round(os.path.getsize(file_path) / (1024 * 1024), 1)
|
|
except Exception:
|
|
size_mb = 0
|
|
|
|
models.append({
|
|
"filename": f,
|
|
"size_mb": size_mb,
|
|
"metadata": meta,
|
|
"preview_url": preview_url
|
|
})
|
|
|
|
return web.json_response({"models": models})
|
|
|
|
|
|
async def api_serve_image(request):
|
|
"""Dedicated image serving endpoint for model preview images."""
|
|
folder_type = request.query.get('type', 'checkpoints')
|
|
try:
|
|
path_idx = int(request.query.get('path_idx', 0))
|
|
except:
|
|
path_idx = 0
|
|
subfolder = request.query.get('subfolder', '')
|
|
filename = request.query.get('filename', '')
|
|
|
|
if not filename or '..' in filename or '..' in subfolder:
|
|
return web.Response(status=400, text='Invalid request')
|
|
|
|
try:
|
|
paths = folder_paths.get_folder_paths(folder_type)
|
|
except Exception:
|
|
return web.Response(status=404, text='Folder type not found')
|
|
if not paths or path_idx >= len(paths):
|
|
return web.Response(status=404, text='Folder type not found')
|
|
|
|
base_dir = paths[path_idx]
|
|
if subfolder:
|
|
file_path = os.path.join(base_dir, subfolder, filename)
|
|
else:
|
|
file_path = os.path.join(base_dir, filename)
|
|
|
|
if not os.path.exists(file_path) or not os.path.isfile(file_path):
|
|
return web.Response(status=404, text='Image not found')
|
|
|
|
ext = os.path.splitext(filename)[1].lower()
|
|
content_types = {
|
|
'.png': 'image/png',
|
|
'.jpg': 'image/jpeg',
|
|
'.jpeg': 'image/jpeg',
|
|
'.webp': 'image/webp',
|
|
'.gif': 'image/gif'
|
|
}
|
|
content_type = content_types.get(ext, 'application/octet-stream')
|
|
|
|
return web.FileResponse(file_path, headers={'Content-Type': content_type})
|
|
|
|
|
|
|
|
async def api_scan_status(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({"scanning": False})
|
|
if not paths or path_idx >= len(paths) or '..' in subfolder:
|
|
return web.json_response({"scanning": False})
|
|
|
|
base_dir = paths[path_idx]
|
|
if subfolder == '/':
|
|
target_dir = base_dir
|
|
else:
|
|
target_dir = os.path.join(base_dir, subfolder.strip('/'))
|
|
|
|
marker_file = os.path.join(target_dir, '.scan_in_progress')
|
|
return web.json_response({"scanning": os.path.exists(marker_file)})
|
|
|
|
async def api_scan_folder(request):
|
|
"""Launches the scraper in the background for a specific directory."""
|
|
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
|
|
|
|
if '..' in subfolder:
|
|
return web.json_response({"status": "error", "message": "Invalid 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('/'))
|
|
|
|
if not os.path.exists(target_dir):
|
|
return web.json_response({"status": "error", "message": "Directory does not exist"})
|
|
|
|
# Get the scraper path relative to this script
|
|
script_dir = os.path.dirname(os.path.abspath(__file__))
|
|
scraper_path = os.path.join(script_dir, "scraper.py")
|
|
|
|
if not os.path.exists(scraper_path):
|
|
return web.json_response({"status": "error", "message": "scraper.py not found in extension directory"})
|
|
|
|
print(f"[Anomalous Browser] Starting background scan for: {target_dir}")
|
|
|
|
try:
|
|
marker_file = os.path.join(target_dir, '.scan_in_progress')
|
|
with open(marker_file, 'w') as f: f.write('1')
|
|
def run_bg():
|
|
try:
|
|
subprocess.run(
|
|
[sys.executable, scraper_path, target_dir],
|
|
cwd=script_dir,
|
|
stdout=subprocess.DEVNULL,
|
|
stderr=subprocess.DEVNULL
|
|
)
|
|
finally:
|
|
if os.path.exists(marker_file):
|
|
try: os.remove(marker_file)
|
|
except: pass
|
|
|
|
threading.Thread(target=run_bg, daemon=True).start()
|
|
|
|
return web.json_response({"status": "ok", "message": "Scan started in background. Check console for details."})
|
|
except Exception as e:
|
|
return web.json_response({"status": "error", "message": str(e)})
|
|
|
|
|
|
async def api_get_config(request):
|
|
script_dir = os.path.dirname(os.path.abspath(__file__))
|
|
config_path = os.path.join(script_dir, "config.json")
|
|
has_key = False
|
|
if os.path.exists(config_path):
|
|
try:
|
|
with open(config_path, 'r', encoding='utf-8') as f:
|
|
cfg = json.load(f)
|
|
has_key = bool(cfg.get("CIVITAI_API_KEY", "").strip())
|
|
except:
|
|
pass
|
|
return web.json_response({"has_api_key": has_key})
|
|
|
|
async def api_save_config(request):
|
|
try:
|
|
data = await request.json()
|
|
api_key = data.get("api_key", "").strip()
|
|
script_dir = os.path.dirname(os.path.abspath(__file__))
|
|
config_path = os.path.join(script_dir, "config.json")
|
|
|
|
cfg = {}
|
|
if os.path.exists(config_path):
|
|
try:
|
|
with open(config_path, 'r', encoding='utf-8') as f:
|
|
cfg = json.load(f)
|
|
except:
|
|
pass
|
|
|
|
cfg["CIVITAI_API_KEY"] = api_key
|
|
with open(config_path, 'w', encoding='utf-8') as f:
|
|
json.dump(cfg, f, indent=4)
|
|
|
|
return web.json_response({"status": "ok"})
|
|
except Exception as e:
|
|
return web.json_response({"status": "error", "message": str(e)})
|
|
|
|
|
|
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": "Invalid request parameters"})
|
|
|
|
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"})
|
|
|
|
base_name = os.path.splitext(filename)[0]
|
|
|
|
# extensions that are considered associated with the model
|
|
associated_exts = [
|
|
'.safetensors', '.ckpt', '.pt', '.bin',
|
|
'.info', '.civitai.info',
|
|
'.png', '.preview.png', '.jpg', '.jpeg', '.webp', '.gif',
|
|
'.json', '.txt', '.yaml'
|
|
]
|
|
|
|
deleted_files = []
|
|
for ext in associated_exts:
|
|
file_to_del = os.path.join(target_dir, base_name + ext)
|
|
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}")
|
|
|
|
return web.json_response({"status": "ok", "deleted": deleted_files})
|
|
|
|
except Exception as e:
|
|
import traceback
|
|
traceback.print_exc()
|
|
return web.json_response({"status": "error", "message": str(e)}, status=500)
|
|
|
|
|
|
def setup_routes(app):
|
|
app.router.add_get('/anomalous/folders', api_get_folders)
|
|
app.router.add_get('/anomalous/models', api_get_models)
|
|
app.router.add_get('/anomalous/image', api_serve_image)
|
|
app.router.add_post('/anomalous/scan', api_scan_folder)
|
|
app.router.add_get('/anomalous/scan_status', api_scan_status)
|
|
app.router.add_get('/anomalous/config', api_get_config)
|
|
app.router.add_post('/anomalous/config', api_save_config)
|
|
app.router.add_post('/anomalous/delete_model', api_delete_model)
|