kaggle_gpu_server/engine/workflow.py

"""
Workflow Engine — Maps user requests to model pipelines.

Analyzes user intent via keyword matching and optional captioning,
then selects and executes the appropriate pipeline.

Usage:
    engine = WorkflowEngine(registry, vram_manager)
    result_image = engine.process_edit(image, "make me look like an anime character")
    result_image = engine.process_generate("a futuristic city at sunset")
"""

import re
from typing import Optional
from PIL import Image

from .plugin_base import PluginCapability
from .plugin_registry import PluginRegistry
from .vram_manager import VRAMManager
from .pipeline import Pipeline, PipelineStep, PipelineContext


# ═══════════════════════════════════════════════════════════════
# Pre-defined Workflows
# ═══════════════════════════════════════════════════════════════

PIPELINES = {
    # ── Style Transfer (identity-preserving) ──
    "style_transfer": Pipeline("style_transfer", [
        PipelineStep("nsfw_detection", optional=True),
        PipelineStep("captioning", optional=True),
        PipelineStep("face_detection", optional=True),
        PipelineStep("face_recognition", optional=True),
        PipelineStep("image_generation", config={"mode": "style_transfer"}),
        PipelineStep("face_restoration", optional=True),
    ], description="Apply artistic styles while preserving identity"),

    # ── Anime / Cartoon Conversion ──
    "anime_conversion": Pipeline("anime_conversion", [
        PipelineStep("nsfw_detection", optional=True),
        PipelineStep("face_detection", optional=True),
        PipelineStep("face_recognition", optional=True),
        PipelineStep("anime_conversion"),
        PipelineStep("face_restoration", optional=True),
        PipelineStep("upscaling", config={"scale_factor": 1}, optional=True),
    ], description="Convert photo to anime/manga style"),

    # ── Background Swap ──
    "background_swap": Pipeline("background_swap", [
        PipelineStep("nsfw_detection", optional=True),
        PipelineStep("segmentation", config={"target": "foreground"}),
        PipelineStep("depth_estimation", optional=True),
        PipelineStep("image_generation", config={"mode": "background"}),
        PipelineStep("compositing", optional=True),
    ], description="Replace image background"),

    # ── Outfit / Clothing Change ──
    "outfit_change": Pipeline("outfit_change", [
        PipelineStep("nsfw_detection", optional=True),
        PipelineStep("captioning", optional=True),
        PipelineStep("face_detection", optional=True),
        PipelineStep("face_recognition", optional=True),
        PipelineStep("segmentation", config={"target": "clothing"}),
        PipelineStep("inpainting"),
        PipelineStep("face_restoration", optional=True),
    ], description="Change outfit or clothing"),

    # ── Age Transformation ──
    "age_change": Pipeline("age_change", [
        PipelineStep("nsfw_detection", optional=True),
        PipelineStep("face_detection"),
        PipelineStep("face_recognition"),
        PipelineStep("identity_preservation", config={"type": "age"}),
        PipelineStep("face_restoration", optional=True),
    ], description="Age or de-age a person"),

    # ── Face Swap ──
    "face_swap": Pipeline("face_swap", [
        PipelineStep("nsfw_detection", optional=True),
        PipelineStep("face_detection"),
        PipelineStep("face_recognition"),
        PipelineStep("identity_preservation", config={"type": "swap"}),
        PipelineStep("face_restoration", optional=True),
    ], description="Swap faces between images"),

    # ── Object Insertion ──
    "object_insertion": Pipeline("object_insertion", [
        PipelineStep("nsfw_detection", optional=True),
        PipelineStep("captioning", optional=True),
        PipelineStep("depth_estimation", optional=True),
        PipelineStep("inpainting", config={"mode": "insert"}),
    ], description="Insert objects into the scene"),

    # ── Relighting ──
    "relighting": Pipeline("relighting", [
        PipelineStep("nsfw_detection", optional=True),
        PipelineStep("segmentation", config={"target": "foreground"}, optional=True),
        PipelineStep("depth_estimation", optional=True),
        PipelineStep("image_generation", config={"mode": "relight"}),
    ], description="Change lighting conditions"),

    # ── Background Removal ──
    "background_removal": Pipeline("background_removal", [
        PipelineStep("segmentation", config={"target": "foreground"}),
    ], description="Remove background (transparent)"),

    # ── Image Enhancement / Upscaling ──
    "enhance": Pipeline("enhance", [
        PipelineStep("face_restoration", optional=True),
        PipelineStep("upscaling"),
    ], description="Enhance and upscale image"),

    # ── General Edit (catch-all) ──
    "general_edit": Pipeline("general_edit", [
        PipelineStep("nsfw_detection", optional=True),
        PipelineStep("captioning", optional=True),
        PipelineStep("image_generation", config={"mode": "edit"}),
    ], description="General-purpose image editing"),

    # ── Text-to-Image Generation ──
    "generate": Pipeline("generate", [
        PipelineStep("nsfw_detection", config={"mode": "prompt"}, optional=True),
        PipelineStep("image_generation", config={"mode": "txt2img"}),
    ], description="Generate image from text prompt"),
}


# ═══════════════════════════════════════════════════════════════
# Intent Detection — Keyword-based request classification
# ═══════════════════════════════════════════════════════════════

# Patterns: (regex_pattern, pipeline_name, priority)
# Higher priority = matched first
INTENT_PATTERNS = [
    # Background operations
    (r"\b(remove|delete|erase)\s*(the\s+)?(background|bg)\b", "background_removal", 100),
    (r"\b(transparent|no\s*background|cut\s*out)\b", "background_removal", 95),
    (r"\b(change|replace|swap|new)\s*(the\s+)?(background|bg|scene|place|location)\b", "background_swap", 90),
    (r"\b(background|bg)\s*(to|with|into)\b", "background_swap", 85),
    (r"\bput\s*(me|him|her|them|this)\s*(in|at|on)\b", "background_swap", 80),

    # Anime / style conversion
    (r"\b(anime|manga|japanese\s*anime)\b", "anime_conversion", 90),
    (r"\b(chibi|kawaii)\b", "anime_conversion", 85),
    (r"\b(cartoon|oil\s*paint|watercolor|pencil|sketch|pop\s*art|comic|cyberpunk|fantasy|3d\s*pixar|pixar|disney)\b", "style_transfer", 80),
    (r"\b(style\s*of|in\s*the\s*style|artistic|art\s*style)\b", "style_transfer", 75),

    # Outfit / clothing
    (r"\b(change|replace|swap|new)\s*(the\s+)?(outfit|clothes|clothing|dress|shirt|suit|wear)\b", "outfit_change", 90),
    (r"\b(wear|wearing|dressed\s*in|put\s*on)\b", "outfit_change", 80),

    # Age transformation
    (r"\b(older|younger|age|aged|aging|baby|child|elderly|teenager|teen|old\s*version|young\s*version)\b", "age_change", 85),
    (r"\b(how\s*(would|will)\s*(i|they|he|she)\s*look\s*(when|at|in))\b", "age_change", 80),

    # Face swap
    (r"\b(face\s*swap|swap\s*face|swap\s*my\s*face|put\s*my\s*face)\b", "face_swap", 90),

    # Object insertion
    (r"\b(add|insert|place|put)\s+(a\s+|an\s+|the\s+)?\w+\s+(in|on|to|into|next\s*to|beside|near)\b", "object_insertion", 70),

    # Relighting
    (r"\b(relight|lighting|sunset\s*light|golden\s*hour|dramatic\s*light|studio\s*light|neon\s*light)\b", "relighting", 75),
    (r"\b(change|add|make)\s*(the\s+)?(light|lighting|illumination)\b", "relighting", 70),

    # Enhancement / upscaling
    (r"\b(enhance|upscale|upres|super\s*resolution|sharpen|improve\s*quality|hd|4k|8k)\b", "enhance", 85),
    (r"\b(fix|restore|clean|denoise|deblur)\s*(the\s+)?(face|image|photo|picture)\b", "enhance", 75),

    # General edit (lowest priority catch-all)
    (r"\b(edit|modify|transform|convert|make|turn)\b", "general_edit", 10),
]

# Compile regex patterns
_COMPILED_PATTERNS = [
    (re.compile(pattern, re.IGNORECASE), pipeline_name, priority)
    for pattern, pipeline_name, priority in INTENT_PATTERNS
]


def analyze_request(prompt: str, has_image: bool = True) -> str:
    """
    Analyze a user's request and determine which pipeline to use.
    
    Args:
        prompt: The user's text instruction
        has_image: Whether an input image was provided
        
    Returns:
        Pipeline name (e.g., "anime_conversion", "background_swap")
    """
    if not prompt.strip():
        return "enhance" if has_image else "generate"

    if not has_image:
        return "generate"

    # Score each pattern
    matches = []
    for regex, pipeline_name, priority in _COMPILED_PATTERNS:
        if regex.search(prompt):
            matches.append((priority, pipeline_name))

    if matches:
        # Return highest priority match
        matches.sort(reverse=True)
        return matches[0][1]

    # Default fallback
    return "general_edit"


# ═══════════════════════════════════════════════════════════════
# Style Detection
# ═══════════════════════════════════════════════════════════════

STYLE_INSTRUCTIONS = {
    "cartoon": "Transform this photo into a high-quality cartoon illustration with vibrant colors, clean outlines, and expressive features while preserving the person's identity and facial features exactly",
    "oil_painting": "Transform this photo into a masterful oil painting with rich brush strokes, detailed textures, warm classical art tones, and canvas texture while preserving the person's identity exactly",
    "pencil": "Transform this photo into a detailed pencil sketch with realistic graphite shading, crosshatching, fine lines, and paper texture while preserving the person's identity exactly",
    "watercolor": "Transform this photo into a beautiful watercolor painting with soft brush strokes, color blending, fluid artistic washes, and delicate transparency while preserving the person's identity exactly",
    "anime": "Transform this photo into a high-quality anime character illustration with vibrant colors, beautiful anime eyes, clean line art, and studio ghibli style while preserving the person's identity and facial structure exactly",
    "pop_art": "Transform this photo into bold pop art style with vivid flat colors, screen print halftone aesthetic, high contrast, and Andy Warhol inspired design while preserving the person's identity exactly",
    "sketch": "Transform this photo into a clean minimalist line sketch with precise thin lines, minimal shading, on a clean white background while preserving the person's identity exactly",
    "3d_pixar": "Transform this photo into a 3D Pixar/Disney animation style character with smooth plastic-like skin, bright vibrant colors, cinematic lighting, and cute character design while preserving the person's identity exactly",
    "comic": "Transform this photo into a comic book illustration with bold ink outlines, halftone dot shading, retro comic coloring, and dynamic superhero aesthetic while preserving the person's identity exactly",
    "cyberpunk": "Transform this photo into cyberpunk style with neon glowing lights, futuristic high-tech aesthetic, dark synthwave atmosphere, and holographic effects while preserving the person's identity exactly",
    "fantasy": "Transform this photo into fantasy art style with magical atmosphere, glowing ethereal particles, enchanted lighting, and mythical aesthetic while preserving the person's identity exactly",
    "chibi": "Transform this photo into a cute chibi/kawaii character with oversized head, tiny body, cute expressions, bright colorful anime style while preserving the person's facial features",
}

STYLE_ALIASES = {
    "cartoon": "cartoon", "oil painting": "oil_painting", "oil_painting": "oil_painting",
    "pencil": "pencil", "watercolor": "watercolor", "anime": "anime",
    "pop art": "pop_art", "pop_art": "pop_art", "sketch": "sketch",
    "3d pixar": "3d_pixar", "3d_pixar": "3d_pixar", "pixar": "3d_pixar", "disney": "3d_pixar",
    "comic": "comic", "cyberpunk": "cyberpunk", "fantasy": "fantasy", "chibi": "chibi",
}


def detect_style(prompt: str) -> Optional[str]:
    """Detect style name from prompt. Returns canonical style name or None."""
    if not prompt:
        return None
    normalized = prompt.strip().lower()
    # Direct match
    if normalized in STYLE_INSTRUCTIONS:
        return normalized
    # Alias match
    for keyword, canonical in STYLE_ALIASES.items():
        if keyword in normalized:
            return canonical
    return None


# ═══════════════════════════════════════════════════════════════
# Workflow Engine — Main orchestrator
# ═══════════════════════════════════════════════════════════════

class WorkflowEngine:
    """
    High-level orchestrator that ties together request analysis,
    pipeline selection, and execution.
    """

    def __init__(self, registry: PluginRegistry, vram_manager: VRAMManager):
        self.registry = registry
        self.vram_manager = vram_manager

    def process_edit(
        self,
        image: Image.Image,
        instruction: str,
        style: str = "",
        pipeline_override: str = "",
        **kwargs,
    ) -> Image.Image:
        """
        Process an image edit request.
        
        Args:
            image: Input PIL Image
            instruction: User's edit instruction
            style: Optional style name override
            pipeline_override: Force a specific pipeline (bypass auto-detection)
            **kwargs: Extra context (seed, guidance_scale, etc.)
            
        Returns:
            Output PIL Image
        """
        # Determine style
        detected_style = detect_style(style or instruction)
        
        # Determine pipeline
        if pipeline_override and pipeline_override in PIPELINES:
            pipeline_name = pipeline_override
        elif detected_style and detected_style == "anime":
            pipeline_name = "anime_conversion"
        elif detected_style:
            pipeline_name = "style_transfer"
        else:
            pipeline_name = analyze_request(instruction, has_image=True)

        pipeline = PIPELINES.get(pipeline_name)
        if pipeline is None:
            pipeline = PIPELINES["general_edit"]
            pipeline_name = "general_edit"

        print(f"\n🧠 Request: \"{instruction}\"")
        print(f"   Detected style: {detected_style or 'none'}")
        print(f"   Selected pipeline: {pipeline_name}")

        # Build context
        # Construct the prompt for generation steps
        if detected_style and detected_style in STYLE_INSTRUCTIONS:
            generation_prompt = STYLE_INSTRUCTIONS[detected_style]
        else:
            generation_prompt = instruction

        context = PipelineContext(
            image=image,
            prompt=generation_prompt,
            style=detected_style or "",
            **{k: v for k, v in kwargs.items() if hasattr(PipelineContext, k)},
        )

        # Execute pipeline
        context = pipeline.execute(context, self.registry, self.vram_manager)

        if context.nsfw_detected:
            raise ValueError("Content blocked: NSFW content detected in the image or request.")

        if context.output_image is not None:
            return context.output_image
        elif context.image is not None:
            return context.image
        else:
            raise RuntimeError(f"Pipeline '{pipeline_name}' produced no output image. Errors: {context.errors}")

    def process_generate(
        self,
        prompt: str,
        **kwargs,
    ) -> Image.Image:
        """
        Process a text-to-image generation request.
        
        Args:
            prompt: Text description of the image to generate
            **kwargs: Extra context (seed, guidance_scale, num_steps, etc.)
            
        Returns:
            Generated PIL Image
        """
        pipeline = PIPELINES["generate"]

        print(f"\n🧠 Generate: \"{prompt}\"")

        context = PipelineContext(
            prompt=prompt,
            **{k: v for k, v in kwargs.items() if hasattr(PipelineContext, k)},
        )

        context = pipeline.execute(context, self.registry, self.vram_manager)

        if context.nsfw_detected:
            raise ValueError("Content blocked: NSFW content detected in the prompt or generated image.")

        if context.output_image is not None:
            return context.output_image
        else:
            raise RuntimeError(f"Generation pipeline produced no output. Errors: {context.errors}")

    def process_remove_background(self, image: Image.Image) -> Image.Image:
        """Remove background from image."""
        pipeline = PIPELINES["background_removal"]
        context = PipelineContext(image=image)
        context = pipeline.execute(context, self.registry, self.vram_manager)
        return context.output_image or context.image

    def process_enhance(self, image: Image.Image, scale_factor: int = 2) -> Image.Image:
        """Enhance and upscale image."""
        pipeline = PIPELINES["enhance"]
        context = PipelineContext(image=image, scale_factor=scale_factor)
        context = pipeline.execute(context, self.registry, self.vram_manager)
        return context.output_image or context.image

    def get_available_pipelines(self) -> dict:
        """Return all available pipelines and their descriptions."""
        return {
            name: {
                "description": p.description,
                "steps": [s.capability for s in p.steps],
            }
            for name, p in PIPELINES.items()
        }

    def get_status(self) -> dict:
        """Get engine status including VRAM and plugin info."""
        return {
            "vram": self.vram_manager.get_status(),
            "plugins": [p.__dict__ for p in self.registry.list_plugins()],
            "pipelines": list(PIPELINES.keys()),
            "capabilities": self.registry.list_capabilities(),
        }