"""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")