Spaces:
Sleeping
Sleeping
File size: 8,538 Bytes
c3649b4 bd1c096 c3649b4 bd1c096 c3649b4 bd1c096 c3649b4 bd1c096 c3649b4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 | """
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)"
|