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