120 lines
3.8 KiB
Python
Executable File
120 lines
3.8 KiB
Python
Executable File
"""
|
|
Video Generation Tool
|
|
Generate videos 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 VideoGenerationTool(BaseTool):
|
|
"""Generate videos using ComfyUI."""
|
|
|
|
def __init__(self, config: Optional[Dict] = None):
|
|
self.client = ComfyUIClient()
|
|
super().__init__(config)
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return "generate_video"
|
|
|
|
@property
|
|
def description(self) -> str:
|
|
return "Generate a video from a text description. Creates animated visual content."
|
|
|
|
@property
|
|
def parameters(self) -> Dict[str, Any]:
|
|
return {
|
|
"type": "object",
|
|
"properties": {
|
|
"prompt": {
|
|
"type": "string",
|
|
"description": "Description of the video to generate"
|
|
},
|
|
"negative_prompt": {
|
|
"type": "string",
|
|
"description": "What to avoid in the video (optional)",
|
|
"default": ""
|
|
},
|
|
"frames": {
|
|
"type": "integer",
|
|
"description": "Number of frames to generate",
|
|
"default": 24
|
|
},
|
|
"seed": {
|
|
"type": "integer",
|
|
"description": "Random seed for reproducibility (optional)"
|
|
}
|
|
},
|
|
"required": ["prompt"]
|
|
}
|
|
|
|
async def execute(
|
|
self,
|
|
prompt: str,
|
|
negative_prompt: str = "",
|
|
frames: int = 24,
|
|
seed: Optional[int] = None,
|
|
**kwargs
|
|
) -> ToolResult:
|
|
"""Generate a video."""
|
|
self._log_execution({"prompt": prompt[:100], "frames": frames})
|
|
|
|
# Reload config to get latest settings
|
|
self.client.reload_config()
|
|
|
|
# Load the video workflow
|
|
workflow = self.client.load_workflow("video")
|
|
|
|
if not workflow:
|
|
return ToolResult(
|
|
success=False,
|
|
error="Video 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="video",
|
|
negative_prompt=negative_prompt,
|
|
frames=frames,
|
|
seed=seed
|
|
)
|
|
|
|
# Queue the prompt
|
|
prompt_id = await self.client.queue_prompt(modified_workflow)
|
|
logger.info(f"Queued video generation: {prompt_id}")
|
|
|
|
# Wait for completion (longer timeout for videos)
|
|
outputs = await self.client.wait_for_completion(
|
|
prompt_id,
|
|
timeout=600 # 10 minutes for video generation
|
|
)
|
|
|
|
# Get output files
|
|
videos = await self.client.get_output_files(outputs, "videos")
|
|
|
|
if not videos:
|
|
return ToolResult(
|
|
success=False,
|
|
error="No video was generated"
|
|
)
|
|
|
|
result = f"Successfully generated video with {len(videos)} output(s):\n"
|
|
result += "\n".join(f" - {v.get('filename', 'video')}" for v in videos)
|
|
|
|
self._log_success(result)
|
|
return ToolResult(success=True, data=result)
|
|
|
|
except TimeoutError as e:
|
|
self._log_error(str(e))
|
|
return ToolResult(success=False, error="Video generation timed out")
|
|
|
|
except Exception as e:
|
|
self._log_error(str(e))
|
|
return ToolResult(success=False, error=str(e))
|