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