init

reelforge/capabilities/image.py

@@ -0,0 +1,141 @@
+"""
+Image Generation Capabilities using ComfyKit
+
+ComfyKit provides unified access to ComfyUI workflows:
+- Local ComfyUI execution
+- RunningHub cloud execution
+- Flexible workflow-based generation
+- Structured result handling
+
+Convention: Tool names must be image_{id}
+"""
+
+import os
+from typing import Any
+
+from comfykit import ComfyKit
+from loguru import logger
+from pydantic import Field
+
+from reelforge.core.mcp_server import reelforge_mcp
+
+
+@reelforge_mcp.tool(
+    description="Generate images using ComfyKit (local ComfyUI or RunningHub cloud)",
+    meta={
+        "reelforge": {
+            "display_name": "ComfyKit Image Generator",
+            "description": "基于 ComfyKit 的图像生成 - 支持本地和云端",
+            "is_default": True,
+        }
+    },
+)
+async def image_comfykit(
+    workflow: str = Field(description="Workflow path, ID, or URL"),
+    comfyui_url: str | None = Field(default=None, description="ComfyUI server URL (default: http://127.0.0.1:8188)"),
+    runninghub_api_key: str | None = Field(default=None, description="RunningHub API key (for cloud execution)"),
+    # Common workflow parameters
+    prompt: str | None = Field(default=None, description="Image generation prompt"),
+    width: int | None = Field(default=None, description="Image width"),
+    height: int | None = Field(default=None, description="Image height"),
+    negative_prompt: str | None = Field(default=None, description="Negative prompt"),
+    steps: int | None = Field(default=None, description="Sampling steps"),
+    seed: int | None = Field(default=None, description="Random seed"),
+    cfg: float | None = Field(default=None, description="CFG scale"),
+    sampler: str | None = Field(default=None, description="Sampler name"),
+) -> str:
+    """
+    Generate image using ComfyKit
+    
+    Supports both local ComfyUI and RunningHub cloud execution.
+    Returns the first generated image URL.
+    
+    Environment variables (optional):
+    - COMFYUI_BASE_URL: ComfyUI server URL
+    - RUNNINGHUB_API_KEY: RunningHub API key
+    
+    Example:
+        # Local ComfyUI (default)
+        image_url = await image_comfykit(
+            workflow="workflows/book_cover.json",
+            title="Atomic Habits",
+            author="James Clear"
+        )
+        
+        # RunningHub cloud
+        image_url = await image_comfykit(
+            workflow="12345",  # RunningHub workflow ID
+            runninghub_api_key="rh-key-xxx",
+            prompt="a beautiful landscape"
+        )
+        
+        # Custom ComfyUI server
+        image_url = await image_comfykit(
+            workflow="workflows/text2img.json",
+            comfyui_url="http://192.168.1.100:8188",
+            prompt="a cute cat"
+        )
+    """
+    logger.debug(f"Generating image with workflow: {workflow}")
+    
+    try:
+        # Initialize ComfyKit
+        kit_config = {}
+        
+        # Local ComfyUI configuration
+        if comfyui_url:
+            kit_config["comfyui_url"] = comfyui_url
+        elif os.getenv("COMFYUI_BASE_URL"):
+            kit_config["comfyui_url"] = os.getenv("COMFYUI_BASE_URL")
+        
+        # RunningHub cloud configuration
+        if runninghub_api_key:
+            kit_config["runninghub_api_key"] = runninghub_api_key
+        elif os.getenv("RUNNINGHUB_API_KEY"):
+            kit_config["runninghub_api_key"] = os.getenv("RUNNINGHUB_API_KEY")
+        
+        kit = ComfyKit(**kit_config)
+        
+        # Build workflow parameters
+        workflow_params = {}
+        if prompt is not None:
+            workflow_params["prompt"] = prompt
+        if width is not None:
+            workflow_params["width"] = width
+        if height is not None:
+            workflow_params["height"] = height
+        if negative_prompt is not None:
+            workflow_params["negative_prompt"] = negative_prompt
+        if steps is not None:
+            workflow_params["steps"] = steps
+        if seed is not None:
+            workflow_params["seed"] = seed
+        if cfg is not None:
+            workflow_params["cfg"] = cfg
+        if sampler is not None:
+            workflow_params["sampler"] = sampler
+        
+        logger.debug(f"Workflow parameters: {workflow_params}")
+        
+        # Execute workflow
+        result = await kit.execute(workflow, workflow_params)
+        
+        # Check execution status
+        if result.status != "completed":
+            error_msg = result.msg or "Unknown error"
+            logger.error(f"Image generation failed: {error_msg}")
+            raise Exception(f"Image generation failed: {error_msg}")
+        
+        # Return first image URL
+        if result.images:
+            image_url = result.images[0]
+            logger.info(f"Generated image: {image_url}")
+            return image_url
+        else:
+            logger.error("No images generated")
+            raise Exception("No images generated")
+    
+    except Exception as e:
+        logger.error(f"ComfyKit error: {e}")
+        raise
+