Spaces:
Sleeping
Sleeping
File size: 8,737 Bytes
c3649b4 6fe6222 c3649b4 a8cef2f c3649b4 6fe6222 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 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 | """
Plugin Base — Abstract base class for all model plugins in the NanoBanana Engine.
Every model plugin MUST subclass `ModelPlugin` and implement:
- load() → download/load weights into VRAM
- unload() → free VRAM and release resources
- run() → execute inference
Plugins are auto-discovered from the `plugins/` directory by the PluginRegistry.
"""
from abc import ABC, abstractmethod
from enum import Enum
from dataclasses import dataclass, field
from typing import Any, Optional
import time
class PluginCapability(str, Enum):
"""Capabilities that plugins can provide. Used for routing in the workflow engine."""
FACE_DETECTION = "face_detection"
FACE_RECOGNITION = "face_recognition"
SEGMENTATION = "segmentation"
POSE_ESTIMATION = "pose_estimation"
DEPTH_ESTIMATION = "depth_estimation"
FACE_RESTORATION = "face_restoration"
ANIME_CONVERSION = "anime_conversion"
INPAINTING = "inpainting"
IDENTITY_PRESERVATION = "identity_preservation"
IMAGE_GENERATION = "image_generation"
UPSCALING = "upscaling"
CAPTIONING = "captioning"
NSFW_DETECTION = "nsfw_detection"
COMPOSITING = "compositing"
class ModelCategory(str, Enum):
"""VRAM category — determines loading/eviction strategy."""
LIGHTWEIGHT = "lightweight" # ≤1.5GB, can stay resident
MEDIUM = "medium" # 1.5–4GB, loaded on demand, LRU eviction
HEAVY = "heavy" # >4GB, mutually exclusive — only 1 at a time
@dataclass
class PluginInfo:
"""Metadata about a plugin, returned by list_plugins()."""
name: str
model_id: str
capability: PluginCapability
category: ModelCategory
vram_estimate_mb: int
is_loaded: bool
version: str
description: str
last_used: Optional[float] = None
class ModelPlugin(ABC):
"""
Abstract base class for all NanoBanana model plugins.
Subclasses must set class-level attributes and implement load/unload/run.
Example:
class MyPlugin(ModelPlugin):
name = "my_model"
model_id = "org/model-name"
capability = PluginCapability.FACE_DETECTION
category = ModelCategory.LIGHTWEIGHT
vram_estimate_mb = 100
version = "1.0.0"
description = "Detects faces using MyModel"
def load(self) -> bool: ...
def unload(self) -> None: ...
def run(self, inputs: dict) -> dict: ...
"""
# ── Required class attributes (set in subclass) ──
name: str = ""
model_id: str = ""
capability: PluginCapability = PluginCapability.IMAGE_GENERATION
category: ModelCategory = ModelCategory.LIGHTWEIGHT
vram_estimate_mb: int = 0
version: str = "1.0.0"
description: str = ""
def __init__(self):
self._loaded = False
self._last_used: Optional[float] = None
self._load_time: Optional[float] = None
self._run_count: int = 0
self._total_run_time: float = 0.0
self._device: str = "cpu"
@property
def is_loaded(self) -> bool:
return self._loaded
@property
def last_used(self) -> Optional[float]:
return self._last_used
@property
def avg_run_time(self) -> float:
if self._run_count == 0:
return 0.0
return self._total_run_time / self._run_count
def set_device(self, device: str):
"""Set the compute device (cuda/cpu). Called by VRAMManager before load()."""
self._device = device
@abstractmethod
def load(self) -> bool:
"""
Load model weights into memory/VRAM.
Returns True on success, False on failure.
Should be idempotent — calling load() when already loaded is a no-op.
"""
...
@abstractmethod
def unload(self) -> None:
"""
Release all model resources and free VRAM.
Must set self._loaded = False.
Should be idempotent — calling unload() when not loaded is a no-op.
"""
...
@abstractmethod
def _execute(self, inputs: dict) -> dict:
"""
Internal execution method. Subclasses implement this.
`inputs` dict varies by capability but always includes:
- "image": PIL.Image (for image-input plugins)
- "prompt": str (for text-input plugins)
Returns a dict with results (e.g., {"image": PIL.Image, "masks": [...], ...})
"""
...
def offload_to_cpu(self):
"""Move all PyTorch modules and custom tensors/pipelines to CPU to free VRAM."""
if not self._loaded:
return
import torch
import gc
# Heavy models use model/sequential offloading or are unloaded completely, so skip manual to("cpu")
if self.category == ModelCategory.HEAVY:
return
print(f"💤 Offloading {self.name} modules to CPU...")
for attr_name in dir(self):
if attr_name.startswith('_'):
continue
try:
attr = getattr(self, attr_name)
if attr is None:
continue
if isinstance(attr, torch.nn.Module):
attr.to("cpu")
elif hasattr(attr, "to") and not isinstance(attr, torch.Tensor):
attr.to("cpu")
except Exception:
pass
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
def load_to_gpu(self):
"""Move all PyTorch modules back to the configured GPU device."""
if not self._loaded or self._device == "cpu":
return
import torch
# Heavy models use model/sequential offloading or are loaded fresh, so skip manual to("cuda")
if self.category == ModelCategory.HEAVY:
return
print(f"🔥 Restoring {self.name} modules to {self._device}...")
for attr_name in dir(self):
if attr_name.startswith('_'):
continue
try:
attr = getattr(self, attr_name)
if attr is None:
continue
if isinstance(attr, torch.nn.Module):
attr.to(self._device)
elif hasattr(attr, "to") and not isinstance(attr, torch.Tensor):
attr.to(self._device)
except Exception:
pass
def run(self, inputs: dict) -> dict:
"""
Execute the plugin with timing and bookkeeping.
Auto-loads if not already loaded. Raises RuntimeError on load failure.
"""
if not self._loaded:
print(f"⚡ Auto-loading {self.name}...")
if not self.load():
raise RuntimeError(f"Failed to auto-load plugin '{self.name}' ({self.model_id})")
start = time.time()
try:
# Wrap execution to optionally run via Hugging Face ZeroGPU spaces decorator
try:
import spaces
# If spaces is available, dynamically decorate and execute on GPU
@spaces.GPU
def run_on_zerogpu():
self.load_to_gpu()
try:
return self._execute(inputs)
finally:
self.offload_to_cpu()
result = run_on_zerogpu()
except ImportError:
# Standard execution if running locally, on Kaggle, or on standard GPU instances
self.load_to_gpu()
try:
result = self._execute(inputs)
finally:
self.offload_to_cpu()
except Exception as e:
print(f"❌ Plugin '{self.name}' failed: {e}")
raise
elapsed = time.time() - start
self._last_used = time.time()
self._run_count += 1
self._total_run_time += elapsed
print(f"✅ {self.name} completed in {elapsed:.2f}s")
return result
def get_info(self) -> PluginInfo:
"""Return metadata about this plugin."""
return PluginInfo(
name=self.name,
model_id=self.model_id,
capability=self.capability,
category=self.category,
vram_estimate_mb=self.vram_estimate_mb,
is_loaded=self._loaded,
version=self.version,
description=self.description,
last_used=self._last_used,
)
def __repr__(self):
status = "✅" if self._loaded else "⭕"
return f"{status} {self.name} ({self.model_id}) [{self.category.value}, ~{self.vram_estimate_mb}MB]"
|