""" Plugin Registry — Auto-discovery and management of model plugins. Scans the `plugins/` directory for Python modules that subclass ModelPlugin. Plugins are registered by their `capability` and can be retrieved by name or capability. Usage: registry = PluginRegistry(vram_manager) registry.discover_plugins() face_det = registry.get("face_detection") face_det.run({"image": pil_img}) all_plugins = registry.list_plugins() """ import importlib import pkgutil import os from typing import Optional from .plugin_base import ModelPlugin, PluginCapability, PluginInfo from .vram_manager import VRAMManager class PluginRegistry: """ Central registry for model plugins. Auto-discovers plugins from the `plugins/` package and manages them through the VRAMManager for VRAM-aware loading/unloading. """ def __init__(self, vram_manager: VRAMManager, config: Optional[dict] = None): self.vram_manager = vram_manager self.config = config or {} # capability -> list of plugin instances (ordered by priority) self._by_capability: dict[PluginCapability, list[ModelPlugin]] = {} # name -> plugin instance self._by_name: dict[str, ModelPlugin] = {} # Disabled plugin names from config self._disabled: set[str] = set() # Parse disabled plugins from config plugins_config = self.config.get("plugins", {}) for cap_name, cap_config in plugins_config.items(): if isinstance(cap_config, dict) and not cap_config.get("enabled", True): self._disabled.add(cap_name) # Parse engine mode ("all", "generation", "editing") engine_config = self.config.get("engine", {}) self.mode = os.environ.get("ENGINE_MODE", engine_config.get("mode", "all")).lower() print(f"⚙️ PluginRegistry initialized in '{self.mode}' mode") def discover_plugins(self, package_name: str = "kaggle_gpu_server.plugins") -> int: """ Auto-discover and register all ModelPlugin subclasses from the plugins package. Returns the number of plugins discovered. """ count = 0 try: package = importlib.import_module(package_name) except ImportError as e: print(f"❌ Cannot import plugins package '{package_name}': {e}") # Try relative import try: plugins_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "plugins") if os.path.isdir(plugins_dir): import sys parent = os.path.dirname(os.path.dirname(os.path.dirname(__file__))) if parent not in sys.path: sys.path.insert(0, parent) package = importlib.import_module(package_name) else: return 0 except ImportError: return 0 package_path = os.path.dirname(package.__file__) for _, module_name, _ in pkgutil.iter_modules([package_path]): if module_name.startswith("_"): continue try: full_name = f"{package_name}.{module_name}" module = importlib.import_module(full_name) # Find all ModelPlugin subclasses in the module for attr_name in dir(module): attr = getattr(module, attr_name) if ( isinstance(attr, type) and issubclass(attr, ModelPlugin) and attr is not ModelPlugin and attr.name # Must have a name set ): plugin = attr() self.register(plugin) count += 1 except Exception as e: print(f"⚠️ Failed to load plugin module '{module_name}': {e}") print(f"📦 Discovered {count} plugins from {package_name}") return count def register(self, plugin: ModelPlugin) -> None: """Register a plugin instance.""" if plugin.name in self._disabled: print(f" Skip disabled plugin: {plugin.name}") return # Filter plugins depending on mode to optimize VRAM if self.mode == "generation": allowed_caps = { PluginCapability.IMAGE_GENERATION, PluginCapability.CAPTIONING, PluginCapability.NSFW_DETECTION, PluginCapability.UPSCALING } if plugin.capability not in allowed_caps: print(f" Skipping non-generation plugin: {plugin.name}") return elif self.mode == "editing": disallowed_caps = { PluginCapability.IMAGE_GENERATION } if plugin.capability in disallowed_caps: print(f" Skipping generation plugin: {plugin.name}") return # Register by name self._by_name[plugin.name] = plugin # Register by capability cap = plugin.capability if cap not in self._by_capability: self._by_capability[cap] = [] self._by_capability[cap].append(plugin) print(f" ✅ Registered: {plugin}") def get(self, capability: str, preferred_name: Optional[str] = None) -> Optional[ModelPlugin]: """ Get a plugin by capability name. Args: capability: The capability name (e.g., "face_detection") preferred_name: Optional specific plugin name to prefer Returns: The plugin instance, or None if no plugin has this capability. """ # Try as PluginCapability enum try: cap = PluginCapability(capability) except ValueError: # Try direct name lookup return self._by_name.get(capability) plugins = self._by_capability.get(cap, []) if not plugins: return None # If a preferred name is given, try that first if preferred_name: for p in plugins: if p.name == preferred_name: return p # Check config for default plugins_config = self.config.get("plugins", {}) cap_config = plugins_config.get(capability, {}) if isinstance(cap_config, dict): default_name = cap_config.get("default") if default_name: for p in plugins: if p.name == default_name: return p # Return first registered plugin for this capability return plugins[0] def get_by_name(self, name: str) -> Optional[ModelPlugin]: """Get a plugin by its exact name.""" return self._by_name.get(name) def ensure_loaded(self, capability: str, preferred_name: Optional[str] = None) -> ModelPlugin: """ Get a plugin and ensure it's loaded (via VRAMManager). Raises RuntimeError if plugin not found or load fails. """ plugin = self.get(capability, preferred_name) if plugin is None: raise RuntimeError(f"No plugin registered for capability '{capability}'") if not plugin.is_loaded: if not self.vram_manager.request_load(plugin): raise RuntimeError(f"Failed to load plugin '{plugin.name}'") return plugin def list_plugins(self) -> list[PluginInfo]: """List all registered plugins with their status.""" return [p.get_info() for p in self._by_name.values()] def list_capabilities(self) -> dict[str, list[str]]: """List all capabilities and which plugins provide them.""" return { cap.value: [p.name for p in plugins] for cap, plugins in self._by_capability.items() } def get_loaded_plugins(self) -> list[str]: """Get names of currently loaded plugins.""" return [name for name, p in self._by_name.items() if p.is_loaded] def __repr__(self): total = len(self._by_name) loaded = len(self.get_loaded_plugins()) caps = len(self._by_capability) return f"PluginRegistry({total} plugins, {loaded} loaded, {caps} capabilities)"