File size: 6,275 Bytes
5b6e956
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""
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
        }