Ben FinkleaandClaude Opus 4.5
261d3c598c
- Added lookup_civitai_by_hash() to query CivitAI API by file hash - Modified lookup-sources endpoint to process one model at a time - UI now shows real-time progress with found/not-found counts - Added Stop button to cancel ongoing lookup - Computes SHA256 hashes for models that don't have them This allows identifying models that were downloaded outside of tracked sources and finding their CivitAI pages automatically. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
578 lines
21 KiB
Python
578 lines
21 KiB
Python
"""API endpoints for ModelSync."""
|
|
|
|
import os
|
|
import asyncio
|
|
from pathlib import Path
|
|
from datetime import datetime
|
|
from typing import Optional
|
|
from aiohttp import web
|
|
|
|
# Import ComfyUI server components
|
|
try:
|
|
from server import PromptServer
|
|
COMFYUI_AVAILABLE = True
|
|
except ImportError:
|
|
COMFYUI_AVAILABLE = False
|
|
PromptServer = None
|
|
|
|
from .lib.metadata_store import MetadataStore
|
|
from .lib.version_tracker import VersionTracker
|
|
from .lib.update_scheduler import UpdateScheduler
|
|
from .lib.sync_manager import SyncManager
|
|
from .lib.hash_utils import compute_sha256_chunked
|
|
|
|
|
|
class ModelSyncAPI:
|
|
"""API handler for ModelSync endpoints."""
|
|
|
|
# Model file extensions to scan
|
|
MODEL_EXTENSIONS = {".safetensors", ".ckpt", ".pt", ".pth", ".bin", ".onnx"}
|
|
|
|
def __init__(self, comfy_models_path: Optional[str] = None):
|
|
"""
|
|
Initialize the ModelSync API.
|
|
|
|
Args:
|
|
comfy_models_path: Path to ComfyUI models directory
|
|
"""
|
|
self.models_path = Path(comfy_models_path) if comfy_models_path else None
|
|
self.store = MetadataStore()
|
|
self.tracker = VersionTracker()
|
|
self.scheduler = UpdateScheduler(
|
|
store=self.store,
|
|
tracker=self.tracker,
|
|
on_update_found=self._on_update_found,
|
|
on_download_complete=self._on_download_complete,
|
|
)
|
|
self.sync_manager = SyncManager(
|
|
store=self.store,
|
|
base_model_path=comfy_models_path,
|
|
)
|
|
|
|
def _on_update_found(self, model: dict):
|
|
"""Callback when an update is found."""
|
|
# Could send WebSocket notification here
|
|
print(f"[ModelSync] Update available: {model['filename']}")
|
|
|
|
def _on_download_complete(self, model_id: int, success: bool):
|
|
"""Callback when a download completes."""
|
|
model = self.store.get_model(model_id)
|
|
if model:
|
|
status = "completed" if success else "failed"
|
|
print(f"[ModelSync] Download {status}: {model['filename']}")
|
|
|
|
def setup_routes(self, app: web.Application):
|
|
"""Register API routes with the web application."""
|
|
app.router.add_get("/modelsync/models", self.list_models)
|
|
app.router.add_post("/modelsync/models/scan", self.scan_models)
|
|
app.router.add_get("/modelsync/models/{id}", self.get_model)
|
|
app.router.add_patch("/modelsync/models/{id}", self.update_model)
|
|
app.router.add_delete("/modelsync/models/{id}", self.delete_model)
|
|
|
|
app.router.add_post("/modelsync/check", self.check_updates)
|
|
app.router.add_post("/modelsync/download/{id}", self.download_update)
|
|
app.router.add_post("/modelsync/lookup-sources", self.lookup_sources)
|
|
|
|
app.router.add_get("/modelsync/config", self.get_config)
|
|
app.router.add_post("/modelsync/config", self.set_config)
|
|
|
|
app.router.add_get("/modelsync/targets", self.list_targets)
|
|
app.router.add_post("/modelsync/targets", self.add_target)
|
|
app.router.add_get("/modelsync/targets/{id}", self.get_target)
|
|
app.router.add_patch("/modelsync/targets/{id}", self.update_target)
|
|
app.router.add_delete("/modelsync/targets/{id}", self.delete_target)
|
|
|
|
app.router.add_post("/modelsync/sync/{id}", self.sync_to_target)
|
|
app.router.add_post("/modelsync/sync", self.sync_to_all)
|
|
app.router.add_get("/modelsync/sync/script/{id}", self.get_sync_script)
|
|
|
|
app.router.add_get("/modelsync/status", self.get_status)
|
|
|
|
# Model endpoints
|
|
|
|
async def list_models(self, request: web.Request) -> web.Response:
|
|
"""List all tracked models."""
|
|
model_type = request.query.get("type")
|
|
filter_updates = request.query.get("updates") == "true"
|
|
filter_auto = request.query.get("auto_update") == "true"
|
|
|
|
if filter_updates:
|
|
models = self.store.get_models_with_updates()
|
|
elif filter_auto:
|
|
models = self.store.get_auto_update_models()
|
|
else:
|
|
models = self.store.get_all_models(model_type=model_type)
|
|
|
|
return web.json_response({"models": models})
|
|
|
|
async def scan_models(self, request: web.Request) -> web.Response:
|
|
"""Scan filesystem and update database."""
|
|
if not self.models_path:
|
|
return web.json_response(
|
|
{"error": "Models path not configured"},
|
|
status=400
|
|
)
|
|
|
|
stats = {"scanned": 0, "new": 0, "updated": 0, "removed": 0}
|
|
found_paths = set()
|
|
|
|
# Scan all model subdirectories
|
|
for model_type_dir in self.models_path.iterdir():
|
|
if not model_type_dir.is_dir():
|
|
continue
|
|
|
|
model_type = model_type_dir.name
|
|
|
|
for model_file in model_type_dir.rglob("*"):
|
|
if model_file.suffix.lower() not in self.MODEL_EXTENSIONS:
|
|
continue
|
|
|
|
if not model_file.is_file():
|
|
continue
|
|
|
|
stats["scanned"] += 1
|
|
path_str = str(model_file)
|
|
found_paths.add(path_str)
|
|
|
|
existing = self.store.get_model_by_path(path_str)
|
|
file_stat = model_file.stat()
|
|
|
|
if existing:
|
|
# Update if size changed
|
|
if existing["file_size"] != file_stat.st_size:
|
|
self.store.update_model(existing["id"], {
|
|
"file_size": file_stat.st_size,
|
|
})
|
|
stats["updated"] += 1
|
|
else:
|
|
# Add new model
|
|
self.store.add_model({
|
|
"path": path_str,
|
|
"filename": model_file.name,
|
|
"model_type": model_type,
|
|
"file_size": file_stat.st_size,
|
|
"install_date": datetime.now().isoformat(),
|
|
})
|
|
stats["new"] += 1
|
|
|
|
# Remove models that no longer exist
|
|
all_models = self.store.get_all_models()
|
|
for model in all_models:
|
|
if model["path"] not in found_paths:
|
|
self.store.delete_model(model["id"])
|
|
stats["removed"] += 1
|
|
|
|
return web.json_response(stats)
|
|
|
|
async def get_model(self, request: web.Request) -> web.Response:
|
|
"""Get a single model by ID."""
|
|
model_id = int(request.match_info["id"])
|
|
model = self.store.get_model(model_id)
|
|
|
|
if not model:
|
|
return web.json_response({"error": "Model not found"}, status=404)
|
|
|
|
return web.json_response(model)
|
|
|
|
async def update_model(self, request: web.Request) -> web.Response:
|
|
"""Update model settings."""
|
|
model_id = int(request.match_info["id"])
|
|
data = await request.json()
|
|
|
|
# Only allow updating specific fields
|
|
allowed_fields = {
|
|
"source_url", "source_type", "auto_update",
|
|
"vram_mb", "exclude_from_sync"
|
|
}
|
|
updates = {k: v for k, v in data.items() if k in allowed_fields}
|
|
|
|
if not updates:
|
|
return web.json_response({"error": "No valid fields to update"}, status=400)
|
|
|
|
success = self.store.update_model(model_id, updates)
|
|
if not success:
|
|
return web.json_response({"error": "Model not found"}, status=404)
|
|
|
|
return web.json_response({"success": True})
|
|
|
|
async def delete_model(self, request: web.Request) -> web.Response:
|
|
"""Remove a model from tracking (does not delete the file)."""
|
|
model_id = int(request.match_info["id"])
|
|
success = self.store.delete_model(model_id)
|
|
|
|
if not success:
|
|
return web.json_response({"error": "Model not found"}, status=404)
|
|
|
|
return web.json_response({"success": True})
|
|
|
|
# Version check endpoints
|
|
|
|
async def check_updates(self, request: web.Request) -> web.Response:
|
|
"""Check models for updates."""
|
|
data = await request.json() if request.body_exists else {}
|
|
model_ids = data.get("model_ids")
|
|
|
|
if model_ids:
|
|
models = [self.store.get_model(mid) for mid in model_ids]
|
|
models = [m for m in models if m and m.get("source_url")]
|
|
else:
|
|
models = self.store.get_all_models()
|
|
models = [m for m in models if m.get("source_url")]
|
|
|
|
if not models:
|
|
return web.json_response({"checked": 0, "updates_available": 0})
|
|
|
|
results = await self.tracker.batch_check(models)
|
|
updates_found = 0
|
|
|
|
for model in models:
|
|
model_id = model["id"]
|
|
if model_id in results:
|
|
remote_info = results[model_id]
|
|
|
|
is_update = self.tracker.is_update_available(
|
|
local_hash=model.get("file_hash"),
|
|
local_size=model.get("file_size"),
|
|
remote_hash=remote_info.get("hash"),
|
|
remote_size=remote_info.get("size"),
|
|
)
|
|
|
|
self.store.update_model(model_id, {
|
|
"remote_hash": remote_info.get("hash"),
|
|
"remote_size": remote_info.get("size"),
|
|
"remote_modified": remote_info.get("modified"),
|
|
"last_checked": datetime.now().isoformat(),
|
|
"update_available": 1 if is_update else 0,
|
|
})
|
|
|
|
if is_update:
|
|
updates_found += 1
|
|
|
|
return web.json_response({
|
|
"checked": len(models),
|
|
"updates_available": updates_found,
|
|
})
|
|
|
|
async def download_update(self, request: web.Request) -> web.Response:
|
|
"""Download update for a specific model."""
|
|
model_id = int(request.match_info["id"])
|
|
model = self.store.get_model(model_id)
|
|
|
|
if not model:
|
|
return web.json_response({"error": "Model not found"}, status=404)
|
|
|
|
if not model.get("source_url"):
|
|
return web.json_response({"error": "No source URL"}, status=400)
|
|
|
|
success = await self.scheduler.download_update(model_id)
|
|
|
|
return web.json_response({
|
|
"success": success,
|
|
"model_id": model_id,
|
|
})
|
|
|
|
async def lookup_sources(self, request: web.Request) -> web.Response:
|
|
"""
|
|
Look up source URL for ONE model at a time by querying CivitAI with file hash.
|
|
Call repeatedly to process all models with progress feedback.
|
|
"""
|
|
data = await request.json() if request.body_exists else {}
|
|
|
|
# Get all models without source URLs
|
|
all_models = self.store.get_all_models()
|
|
models_to_lookup = [m for m in all_models if not m.get("source_url")]
|
|
|
|
if not models_to_lookup:
|
|
return web.json_response({
|
|
"status": "complete",
|
|
"message": "All models already have sources",
|
|
"remaining": 0,
|
|
"total": len(all_models),
|
|
})
|
|
|
|
# Find next model to process (prioritize ones with hashes)
|
|
model = None
|
|
for m in models_to_lookup:
|
|
if m.get("file_hash"):
|
|
model = m
|
|
break
|
|
|
|
# If no model has hash, pick first one to hash
|
|
if not model:
|
|
model = models_to_lookup[0]
|
|
|
|
# Step 1: Compute hash if needed
|
|
if not model.get("file_hash"):
|
|
try:
|
|
print(f"[ModelSync] Hashing {model['filename']}...")
|
|
file_hash = compute_sha256_chunked(model["path"])
|
|
self.store.update_model(model["id"], {"file_hash": file_hash})
|
|
model["file_hash"] = file_hash
|
|
return web.json_response({
|
|
"status": "hashing",
|
|
"message": f"Computed hash for {model['filename']}",
|
|
"model": model["filename"],
|
|
"remaining": len(models_to_lookup),
|
|
"total": len(all_models),
|
|
})
|
|
except Exception as e:
|
|
print(f"[ModelSync] Error hashing {model['filename']}: {e}")
|
|
return web.json_response({
|
|
"status": "error",
|
|
"message": f"Failed to hash {model['filename']}: {str(e)}",
|
|
"remaining": len(models_to_lookup),
|
|
"total": len(all_models),
|
|
})
|
|
|
|
# Step 2: Look up on CivitAI
|
|
try:
|
|
print(f"[ModelSync] Looking up {model['filename']} on CivitAI...")
|
|
info = await self.tracker.lookup_civitai_by_hash(model["file_hash"])
|
|
|
|
if info:
|
|
self.store.update_model(model["id"], {
|
|
"source_url": info.get("source_url"),
|
|
"source_type": info.get("source_type", "civitai"),
|
|
"remote_hash": info.get("remote_hash"),
|
|
})
|
|
return web.json_response({
|
|
"status": "found",
|
|
"message": f"Found {model['filename']} on CivitAI: {info.get('model_name', 'Unknown')}",
|
|
"model": model["filename"],
|
|
"civitai_name": info.get("model_name"),
|
|
"source_url": info.get("source_url"),
|
|
"remaining": len(models_to_lookup) - 1,
|
|
"total": len(all_models),
|
|
})
|
|
else:
|
|
# Mark as checked but not found (use source_type to track)
|
|
self.store.update_model(model["id"], {
|
|
"source_type": "not_found",
|
|
})
|
|
return web.json_response({
|
|
"status": "not_found",
|
|
"message": f"{model['filename']} not found on CivitAI",
|
|
"model": model["filename"],
|
|
"remaining": len(models_to_lookup) - 1,
|
|
"total": len(all_models),
|
|
})
|
|
|
|
except Exception as e:
|
|
print(f"[ModelSync] Error looking up {model['filename']}: {e}")
|
|
return web.json_response({
|
|
"status": "error",
|
|
"message": f"Lookup failed for {model['filename']}: {str(e)}",
|
|
"remaining": len(models_to_lookup),
|
|
"total": len(all_models),
|
|
})
|
|
|
|
# Configuration endpoints
|
|
|
|
async def get_config(self, request: web.Request) -> web.Response:
|
|
"""Get all configuration values."""
|
|
config = self.store.get_all_config()
|
|
return web.json_response(config)
|
|
|
|
async def set_config(self, request: web.Request) -> web.Response:
|
|
"""Set configuration values."""
|
|
data = await request.json()
|
|
|
|
for key, value in data.items():
|
|
self.store.set_config(key, str(value))
|
|
|
|
return web.json_response({"success": True})
|
|
|
|
# Sync target endpoints
|
|
|
|
async def list_targets(self, request: web.Request) -> web.Response:
|
|
"""List all sync targets."""
|
|
targets = self.store.get_all_sync_targets()
|
|
return web.json_response({"targets": targets})
|
|
|
|
async def add_target(self, request: web.Request) -> web.Response:
|
|
"""Add a new sync target."""
|
|
data = await request.json()
|
|
|
|
required = ["name", "model_path"]
|
|
if not all(k in data for k in required):
|
|
return web.json_response(
|
|
{"error": f"Required fields: {required}"},
|
|
status=400
|
|
)
|
|
|
|
try:
|
|
target_id = self.store.add_sync_target(data)
|
|
return web.json_response({"id": target_id, "success": True})
|
|
except Exception as e:
|
|
return web.json_response({"error": str(e)}, status=400)
|
|
|
|
async def get_target(self, request: web.Request) -> web.Response:
|
|
"""Get a sync target by ID."""
|
|
target_id = int(request.match_info["id"])
|
|
target = self.store.get_sync_target(target_id)
|
|
|
|
if not target:
|
|
return web.json_response({"error": "Target not found"}, status=404)
|
|
|
|
return web.json_response(target)
|
|
|
|
async def update_target(self, request: web.Request) -> web.Response:
|
|
"""Update sync target settings."""
|
|
target_id = int(request.match_info["id"])
|
|
data = await request.json()
|
|
|
|
allowed_fields = {"name", "ssh_host", "model_path", "max_vram_mb", "enabled"}
|
|
updates = {k: v for k, v in data.items() if k in allowed_fields}
|
|
|
|
if not updates:
|
|
return web.json_response({"error": "No valid fields"}, status=400)
|
|
|
|
success = self.store.update_sync_target(target_id, updates)
|
|
if not success:
|
|
return web.json_response({"error": "Target not found"}, status=404)
|
|
|
|
return web.json_response({"success": True})
|
|
|
|
async def delete_target(self, request: web.Request) -> web.Response:
|
|
"""Delete a sync target."""
|
|
target_id = int(request.match_info["id"])
|
|
success = self.store.delete_sync_target(target_id)
|
|
|
|
if not success:
|
|
return web.json_response({"error": "Target not found"}, status=404)
|
|
|
|
return web.json_response({"success": True})
|
|
|
|
# Sync operation endpoints
|
|
|
|
async def sync_to_target(self, request: web.Request) -> web.Response:
|
|
"""Sync models to a specific target."""
|
|
target_id = int(request.match_info["id"])
|
|
target = self.store.get_sync_target(target_id)
|
|
|
|
if not target:
|
|
return web.json_response({"error": "Target not found"}, status=404)
|
|
|
|
data = await request.json() if request.body_exists else {}
|
|
dry_run = data.get("dry_run", False)
|
|
|
|
# Run sync in background to not block
|
|
result = self.sync_manager.sync_to_target(target, dry_run=dry_run)
|
|
|
|
return web.json_response({
|
|
"target": target["name"],
|
|
"success": result.success,
|
|
"models_synced": result.models_synced,
|
|
"models_skipped": result.models_skipped,
|
|
"bytes_transferred": result.bytes_transferred,
|
|
"duration_seconds": result.duration_seconds,
|
|
"details": result.details,
|
|
})
|
|
|
|
async def sync_to_all(self, request: web.Request) -> web.Response:
|
|
"""Sync models to all enabled targets."""
|
|
data = await request.json() if request.body_exists else {}
|
|
dry_run = data.get("dry_run", False)
|
|
|
|
results = self.sync_manager.sync_to_all_targets(dry_run=dry_run)
|
|
|
|
return web.json_response({
|
|
"results": [
|
|
{
|
|
"target": r.target_name,
|
|
"success": r.success,
|
|
"models_synced": r.models_synced,
|
|
"models_skipped": r.models_skipped,
|
|
}
|
|
for r in results
|
|
]
|
|
})
|
|
|
|
async def get_sync_script(self, request: web.Request) -> web.Response:
|
|
"""Generate a sync script for a target."""
|
|
target_id = int(request.match_info["id"])
|
|
target = self.store.get_sync_target(target_id)
|
|
|
|
if not target:
|
|
return web.json_response({"error": "Target not found"}, status=404)
|
|
|
|
script = self.sync_manager.generate_sync_script(target)
|
|
|
|
return web.Response(
|
|
text=script,
|
|
content_type="text/x-shellscript",
|
|
headers={
|
|
"Content-Disposition": f'attachment; filename="sync-{target["name"]}.sh"'
|
|
}
|
|
)
|
|
|
|
# Status endpoint
|
|
|
|
async def get_status(self, request: web.Request) -> web.Response:
|
|
"""Get overall ModelSync status."""
|
|
scheduler_state = self.scheduler.state
|
|
sync_summary = self.sync_manager.get_sync_summary()
|
|
|
|
return web.json_response({
|
|
"scheduler": {
|
|
"running": scheduler_state.running,
|
|
"last_check": scheduler_state.last_check.isoformat() if scheduler_state.last_check else None,
|
|
"next_check": scheduler_state.next_check.isoformat() if scheduler_state.next_check else None,
|
|
"models_checked": scheduler_state.models_checked,
|
|
"updates_found": scheduler_state.updates_found,
|
|
"current_downloads": len(scheduler_state.current_downloads),
|
|
},
|
|
"models": {
|
|
"total": self.store.get_model_count(),
|
|
"updates_available": self.store.get_update_count(),
|
|
},
|
|
"sync": sync_summary,
|
|
})
|
|
|
|
|
|
# Global API instance
|
|
_api_instance: Optional[ModelSyncAPI] = None
|
|
|
|
|
|
def get_api() -> Optional[ModelSyncAPI]:
|
|
"""Get the global API instance."""
|
|
return _api_instance
|
|
|
|
|
|
def init_api(models_path: Optional[str] = None) -> ModelSyncAPI:
|
|
"""Initialize and return the global API instance."""
|
|
global _api_instance
|
|
|
|
if _api_instance is None:
|
|
_api_instance = ModelSyncAPI(comfy_models_path=models_path)
|
|
|
|
return _api_instance
|
|
|
|
|
|
def setup_routes():
|
|
"""Setup routes with ComfyUI's PromptServer."""
|
|
if not COMFYUI_AVAILABLE:
|
|
print("[ModelSync] ComfyUI not available, skipping route setup")
|
|
return
|
|
|
|
# Try to detect ComfyUI models path
|
|
models_path = None
|
|
if hasattr(PromptServer, "instance") and PromptServer.instance:
|
|
# Try to get from ComfyUI folder paths
|
|
try:
|
|
import folder_paths
|
|
models_path = folder_paths.models_dir
|
|
except (ImportError, AttributeError):
|
|
pass
|
|
|
|
api = init_api(models_path)
|
|
|
|
if PromptServer.instance:
|
|
api.setup_routes(PromptServer.instance.app)
|
|
print("[ModelSync] API routes registered")
|
|
|
|
# Start scheduler if auto-update is enabled
|
|
if api.store.get_config("auto_update_enabled", "false").lower() == "true":
|
|
api.scheduler.start()
|
|
print("[ModelSync] Background scheduler started")
|