character_forge_image/plugins/comfyui_plugin.py

"""
ComfyUI Backend Plugin

Plugin adapter for ComfyUI local backend with qwen_image_edit_2509.
"""

import sys
import json
import random
from pathlib import Path
from typing import Any, Dict, Optional, List
from PIL import Image

# Add parent directories to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent))

from core.comfyui_client import ComfyUIClient
from config.settings import Settings

# Import from shared plugin system
sys.path.insert(0, str(Path(__file__).parent.parent.parent / 'shared'))
from plugin_system.base_plugin import BaseBackendPlugin


class ComfyUIPlugin(BaseBackendPlugin):
    """Plugin adapter for ComfyUI backend using qwen_image_edit_2509."""

    def __init__(self, config_path: Path):
        """Initialize ComfyUI plugin."""
        super().__init__(config_path)

        # Get settings
        settings = Settings()
        server_address = settings.COMFYUI_BASE_URL.replace("http://", "")

        try:
            self.client = ComfyUIClient(server_address=server_address)
            # Test connection
            healthy, _ = self.client.health_check()
            self.available = healthy
        except Exception as e:
            print(f"Warning: ComfyUI backend not available: {e}")
            self.client = None
            self.available = False

        # Load qwen workflow template
        self.workflow_template = None
        workflow_path = Path(__file__).parent.parent.parent / 'tools' / 'comfyui' / 'workflows' / 'qwen_image_edit.json'
        if workflow_path.exists():
            with open(workflow_path) as f:
                self.workflow_template = json.load(f)
        else:
            print(f"Warning: Workflow template not found at {workflow_path}")

    def health_check(self) -> bool:
        """Check if ComfyUI backend is available."""
        if not self.available or self.client is None:
            return False

        try:
            healthy, _ = self.client.health_check()
            return healthy
        except:
            return False

    def _update_qwen_workflow(
        self,
        workflow: dict,
        prompt: str = None,
        negative_prompt: str = None,
        input_image_filename: str = None,
        seed: int = None,
        width: int = None,
        height: int = None
    ) -> dict:
        """
        Update workflow parameters for qwen_image_edit workflow.

        Node IDs for qwen_image_edit.json:
        - 111: Positive prompt (TextEncodeQwenImageEditPlus)
        - 110: Negative prompt (TextEncodeQwenImageEditPlus)
        - 78: Load Image
        - 3: KSampler (seed)
        - 112: EmptySD3LatentImage (width, height)
        """
        # Clone workflow to avoid modifying original
        wf = json.loads(json.dumps(workflow))

        # Update prompt
        if prompt is not None:
            wf["111"]["inputs"]["prompt"] = prompt

        # Update negative prompt
        if negative_prompt is not None:
            wf["110"]["inputs"]["prompt"] = negative_prompt

        # Update input image
        if input_image_filename is not None:
            wf["78"]["inputs"]["image"] = input_image_filename

        # Update seed
        if seed is not None:
            wf["3"]["inputs"]["seed"] = seed
        else:
            # Random seed if not specified
            wf["3"]["inputs"]["seed"] = random.randint(1, 2**32 - 1)

        # Update dimensions
        if width is not None:
            wf["112"]["inputs"]["width"] = width
        if height is not None:
            wf["112"]["inputs"]["height"] = height

        return wf

    def generate_image(
        self,
        prompt: str,
        input_images: Optional[List[Image.Image]] = None,
        **kwargs
    ) -> Image.Image:
        """
        Generate image using ComfyUI qwen_image_edit_2509 workflow.

        Args:
            prompt: Text prompt for image editing
            input_images: Optional list of input images (uses first image)
            **kwargs: Additional parameters (negative_prompt, seed, width, height)

        Returns:
            Generated PIL Image
        """
        if not self.health_check():
            raise RuntimeError("ComfyUI backend not available")

        if self.workflow_template is None:
            raise RuntimeError("Workflow template not loaded")

        if not input_images or len(input_images) == 0:
            raise ValueError("qwen_image_edit_2509 requires an input image")

        # Upload input image
        input_image = input_images[0]
        uploaded_filename = self.client.upload_image(input_image)

        # Get parameters from kwargs
        negative_prompt = kwargs.get('negative_prompt', '')
        seed = kwargs.get('seed', None)
        width = kwargs.get('width', 1024)
        height = kwargs.get('height', 1024)

        # Update workflow with parameters
        workflow = self._update_qwen_workflow(
            self.workflow_template,
            prompt=prompt,
            negative_prompt=negative_prompt,
            input_image_filename=uploaded_filename,
            seed=seed,
            width=width,
            height=height
        )

        # Execute workflow
        images = self.client.execute_workflow(workflow)

        if not images:
            raise RuntimeError("No images generated")

        # Return first image
        return images[0]

    def get_capabilities(self) -> Dict[str, Any]:
        """Report ComfyUI backend capabilities."""
        return {
            'name': 'ComfyUI Local',
            'type': 'local',
            'supports_input_images': True,
            'supports_multi_image': True,
            'max_input_images': 16,
            'supports_aspect_ratios': True,
            'available_aspect_ratios': ['1:1', '3:4', '4:3', '9:16', '16:9'],
            'supports_guidance_scale': True,
            'supports_inference_steps': True,
            'supports_seed': True,
            'available_models': [
                'qwen_image_edit_2509',  # To be installed
                'flux.1_kontext_ai'  # To be installed
            ],
            'status': 'partial',  # Needs workflow implementation
            'estimated_time_per_image': 3.0,  # seconds (depends on GPU and model)
            'cost_per_image': 0.0,  # Free, local
        }