Ben FinkleaandClaude Opus 4.5
69197802b6
When CivitAI hash lookup fails, now tries: 1. Pattern matching for known model families (Gemma, CLIP, SAM, YOLO, etc.) 2. HuggingFace API search by filename This identifies sources for models from HuggingFace, GitHub, and other repositories that aren't on CivitAI. Also fixed lookup getting stuck on hint results by properly filtering out already-checked models regardless of source_type. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
603 lines
22 KiB
Python
603 lines
22 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 that haven't been looked up yet
|
|
# Exclude: models with source_url, or source_type indicating already checked
|
|
all_models = self.store.get_all_models()
|
|
checked_types = {"not_found", "civitai", "huggingface", "github", "hint"}
|
|
models_to_lookup = [
|
|
m for m in all_models
|
|
if not m.get("source_url") and m.get("source_type") not in checked_types
|
|
]
|
|
|
|
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: Smart lookup - tries CivitAI, pattern matching, HuggingFace
|
|
try:
|
|
print(f"[ModelSync] Looking up {model['filename']}...")
|
|
info = await self.tracker.smart_lookup(model["filename"], model.get("file_hash"))
|
|
|
|
if info and info.get("source_url"):
|
|
# Found with a URL
|
|
self.store.update_model(model["id"], {
|
|
"source_url": info.get("source_url"),
|
|
"source_type": info.get("source_type", "unknown"),
|
|
"remote_hash": info.get("remote_hash"),
|
|
})
|
|
source_name = info.get("model_name") or info.get("source_hint") or "Unknown"
|
|
method = info.get("lookup_method", "unknown")
|
|
return web.json_response({
|
|
"status": "found",
|
|
"message": f"Found {model['filename']}: {source_name} ({method})",
|
|
"model": model["filename"],
|
|
"source_name": source_name,
|
|
"source_url": info.get("source_url"),
|
|
"source_type": info.get("source_type"),
|
|
"lookup_method": method,
|
|
"remaining": len(models_to_lookup) - 1,
|
|
"total": len(all_models),
|
|
})
|
|
elif info and info.get("source_hint"):
|
|
# Found a hint but no URL
|
|
self.store.update_model(model["id"], {
|
|
"source_type": info.get("source_type", "hint"),
|
|
"source_url": info.get("source_hint"), # Store hint as URL for reference
|
|
})
|
|
return web.json_response({
|
|
"status": "hint",
|
|
"message": f"{model['filename']} likely from {info.get('source_hint')}",
|
|
"model": model["filename"],
|
|
"source_hint": info.get("source_hint"),
|
|
"lookup_method": info.get("lookup_method"),
|
|
"remaining": len(models_to_lookup) - 1,
|
|
"total": len(all_models),
|
|
})
|
|
else:
|
|
# Not found anywhere
|
|
self.store.update_model(model["id"], {
|
|
"source_type": "not_found",
|
|
})
|
|
return web.json_response({
|
|
"status": "not_found",
|
|
"message": f"{model['filename']} - source unknown",
|
|
"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")
|