kaggle_gpu_server/engine/plugin_base.py

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