Ben FinkleaandClaude Opus 4.5
0b6b9066b3
Bug fixes: - Fix menu button not appearing: Use app.registerExtension() for proper ComfyUI initialization timing. The button now appears in the top bar next to the Manager button. - Fix source_hint incorrectly saved as source_url: Pattern-matched hints like "Lightricks/LTX-Video" are not valid URLs and should not be used for update checks. - Fix deprecated tempfile.mktemp(): Use secure mkstemp() instead to avoid race conditions. - Add os import for file descriptor handling in update_scheduler.py Test improvements: - Update Playwright tests to hover over UI to reveal ComfyUI menu - All 26 tests now passing (17 unit + 9 UI) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
563 lines
21 KiB
Python
563 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_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:
|
|
"""
|
|
Unified check: finds sources and checks for updates concurrently.
|
|
Processes models in batches for efficiency.
|
|
"""
|
|
import asyncio
|
|
|
|
data = await request.json() if request.body_exists else {}
|
|
model_ids = data.get("model_ids")
|
|
concurrent_limit = data.get("concurrent", 5)
|
|
|
|
all_models = self.store.get_all_models()
|
|
if model_ids:
|
|
all_models = [m for m in all_models if m["id"] in model_ids]
|
|
|
|
if not all_models:
|
|
return web.json_response({
|
|
"total": 0, "checked": 0, "sources_found": 0,
|
|
"updates_available": 0, "status": "complete"
|
|
})
|
|
|
|
# Source types that indicate we already tried to find the source
|
|
already_looked_up = {"not_found", "civitai", "huggingface", "github", "hint"}
|
|
|
|
stats = {
|
|
"total": len(all_models),
|
|
"hashed": 0,
|
|
"sources_found": 0,
|
|
"checked": 0,
|
|
"updates_available": 0,
|
|
"skipped": 0,
|
|
"errors": 0,
|
|
}
|
|
|
|
semaphore = asyncio.Semaphore(concurrent_limit)
|
|
|
|
async def process_model(model):
|
|
"""Process a single model: hash → find source → check update."""
|
|
async with semaphore:
|
|
model_id = model["id"]
|
|
filename = model["filename"]
|
|
|
|
try:
|
|
# Step 1: Compute hash if missing
|
|
if not model.get("file_hash"):
|
|
try:
|
|
file_hash = compute_sha256_chunked(model["path"])
|
|
self.store.update_model(model_id, {"file_hash": file_hash})
|
|
model["file_hash"] = file_hash
|
|
stats["hashed"] += 1
|
|
except Exception as e:
|
|
print(f"[ModelSync] Hash error {filename}: {e}")
|
|
stats["errors"] += 1
|
|
return
|
|
|
|
# Step 2: Find source if we haven't looked yet
|
|
source_url = model.get("source_url")
|
|
if not source_url and model.get("source_type") not in already_looked_up:
|
|
info = await self.tracker.smart_lookup(filename, model.get("file_hash"))
|
|
|
|
if info and info.get("source_url"):
|
|
self.store.update_model(model_id, {
|
|
"source_url": info.get("source_url"),
|
|
"source_type": info.get("source_type", "unknown"),
|
|
})
|
|
source_url = info.get("source_url")
|
|
model["source_url"] = source_url
|
|
model["source_type"] = info.get("source_type")
|
|
stats["sources_found"] += 1
|
|
elif info and info.get("source_hint"):
|
|
# source_hint is just a hint (e.g., repo name), not a usable URL
|
|
# Only save source_type to indicate we looked but couldn't find exact URL
|
|
self.store.update_model(model_id, {
|
|
"source_type": info.get("source_type", "hint"),
|
|
})
|
|
# Don't set source_url since it's not a valid URL for update checks
|
|
else:
|
|
self.store.update_model(model_id, {"source_type": "not_found"})
|
|
|
|
# Step 3: Check for updates if we have a valid HTTP source URL
|
|
source_url = model.get("source_url")
|
|
if source_url and source_url.startswith("http"):
|
|
remote_info = await self.tracker.check_version(source_url)
|
|
|
|
if remote_info:
|
|
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,
|
|
})
|
|
|
|
stats["checked"] += 1
|
|
if is_update:
|
|
stats["updates_available"] += 1
|
|
else:
|
|
# Couldn't check but mark as checked anyway
|
|
self.store.update_model(model_id, {
|
|
"last_checked": datetime.now().isoformat(),
|
|
"update_available": 0,
|
|
})
|
|
stats["checked"] += 1
|
|
else:
|
|
# No valid URL - mark as checked with no update
|
|
self.store.update_model(model_id, {
|
|
"last_checked": datetime.now().isoformat(),
|
|
})
|
|
stats["skipped"] += 1
|
|
|
|
except Exception as e:
|
|
print(f"[ModelSync] Error processing {filename}: {e}")
|
|
stats["errors"] += 1
|
|
|
|
# Process all models concurrently
|
|
await asyncio.gather(*[process_model(m) for m in all_models])
|
|
|
|
stats["status"] = "complete"
|
|
return web.json_response(stats)
|
|
|
|
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,
|
|
})
|
|
|
|
# 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")
|