138 lines
4.4 KiB
Python
Executable File
138 lines
4.4 KiB
Python
Executable File
"""
|
|
Image Generation Tool
|
|
Generate images using ComfyUI.
|
|
"""
|
|
from typing import Dict, Any, Optional
|
|
from loguru import logger
|
|
|
|
from tools.base import BaseTool, ToolResult
|
|
from tools.comfyui.base import ComfyUIClient
|
|
|
|
|
|
class ImageGenerationTool(BaseTool):
|
|
"""Generate images using ComfyUI."""
|
|
|
|
def __init__(self, config: Optional[Dict] = None):
|
|
self.client = ComfyUIClient()
|
|
super().__init__(config)
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return "generate_image"
|
|
|
|
@property
|
|
def description(self) -> str:
|
|
return "Generate an image from a text description. Creates visual content based on your prompt."
|
|
|
|
@property
|
|
def parameters(self) -> Dict[str, Any]:
|
|
return {
|
|
"type": "object",
|
|
"properties": {
|
|
"prompt": {
|
|
"type": "string",
|
|
"description": "Description of the image to generate"
|
|
},
|
|
"negative_prompt": {
|
|
"type": "string",
|
|
"description": "What to avoid in the image (optional)",
|
|
"default": ""
|
|
},
|
|
"size": {
|
|
"type": "string",
|
|
"description": "Image size (e.g., '512x512', '1024x768')",
|
|
"default": "512x512"
|
|
},
|
|
"steps": {
|
|
"type": "integer",
|
|
"description": "Number of generation steps",
|
|
"default": 20
|
|
},
|
|
"cfg_scale": {
|
|
"type": "number",
|
|
"description": "CFG scale for prompt adherence",
|
|
"default": 7.0
|
|
},
|
|
"seed": {
|
|
"type": "integer",
|
|
"description": "Random seed for reproducibility (optional)"
|
|
}
|
|
},
|
|
"required": ["prompt"]
|
|
}
|
|
|
|
async def execute(
|
|
self,
|
|
prompt: str,
|
|
negative_prompt: str = "",
|
|
size: str = "512x512",
|
|
steps: int = 20,
|
|
cfg_scale: float = 7.0,
|
|
seed: Optional[int] = None,
|
|
**kwargs
|
|
) -> ToolResult:
|
|
"""Generate an image."""
|
|
self._log_execution({"prompt": prompt[:100], "size": size, "steps": steps})
|
|
|
|
# Reload config to get latest settings
|
|
self.client.reload_config()
|
|
|
|
# Load the image workflow
|
|
workflow = self.client.load_workflow("image")
|
|
|
|
if not workflow:
|
|
return ToolResult(
|
|
success=False,
|
|
error="Image generation workflow not configured. Please upload a workflow JSON in the admin panel."
|
|
)
|
|
|
|
try:
|
|
# Modify workflow with parameters
|
|
modified_workflow = self.client.modify_workflow(
|
|
workflow,
|
|
prompt=prompt,
|
|
workflow_type="image",
|
|
negative_prompt=negative_prompt,
|
|
size=size,
|
|
steps=steps,
|
|
cfg_scale=cfg_scale,
|
|
seed=seed
|
|
)
|
|
|
|
# Queue the prompt
|
|
prompt_id = await self.client.queue_prompt(modified_workflow)
|
|
logger.info(f"Queued image generation: {prompt_id}")
|
|
|
|
# Wait for completion
|
|
outputs = await self.client.wait_for_completion(
|
|
prompt_id,
|
|
timeout=300 # 5 minutes for image generation
|
|
)
|
|
|
|
# Get output images
|
|
images = await self.client.get_output_images(outputs)
|
|
|
|
if not images:
|
|
return ToolResult(
|
|
success=False,
|
|
error="No images were generated"
|
|
)
|
|
|
|
# Return info about generated images
|
|
result_parts = [f"Successfully generated {len(images)} image(s):"]
|
|
for img in images:
|
|
result_parts.append(f" - {img['filename']}")
|
|
|
|
result = "\n".join(result_parts)
|
|
|
|
self._log_success(result)
|
|
return ToolResult(success=True, data=result)
|
|
|
|
except TimeoutError as e:
|
|
self._log_error(str(e))
|
|
return ToolResult(success=False, error="Image generation timed out")
|
|
|
|
except Exception as e:
|
|
self._log_error(str(e))
|
|
return ToolResult(success=False, error=str(e))
|