Initial release — OpenMontage: the first open-source agentic video production system
11 production pipelines, 47 tools, 124 agent skills. Supports cloud APIs (fal.ai, OpenAI, ElevenLabs, Suno, HeyGen, Runway) and free local providers (diffusers, Piper TTS, WAN 2.1, Hunyuan, CogVideo). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1 @@
|
||||
"""Analysis tools for content inspection, transcription, and scene detection."""
|
||||
@@ -0,0 +1,236 @@
|
||||
"""Frame sampler tool wrapping FFmpeg.
|
||||
|
||||
Extracts representative frames from video for AI analysis, thumbnails,
|
||||
or quality inspection. Supports interval-based, count-based, and
|
||||
timestamp-based extraction strategies.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
ToolResult,
|
||||
ToolStability,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class FrameSampler(BaseTool):
|
||||
name = "frame_sampler"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.CORE
|
||||
capability = "analysis"
|
||||
provider = "ffmpeg"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
|
||||
dependencies = ["cmd:ffmpeg"]
|
||||
install_instructions = (
|
||||
"Install FFmpeg: https://ffmpeg.org/download.html"
|
||||
)
|
||||
agent_skills = ["ffmpeg"]
|
||||
|
||||
capabilities = [
|
||||
"extract_frames_interval",
|
||||
"extract_frames_count",
|
||||
"extract_frames_timestamps",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["input_path", "strategy"],
|
||||
"properties": {
|
||||
"input_path": {"type": "string"},
|
||||
"strategy": {
|
||||
"type": "string",
|
||||
"enum": ["interval", "count", "timestamps"],
|
||||
},
|
||||
"interval_seconds": {
|
||||
"type": "number",
|
||||
"minimum": 0.1,
|
||||
"description": "Seconds between frames (for interval strategy)",
|
||||
},
|
||||
"count": {
|
||||
"type": "integer",
|
||||
"minimum": 1,
|
||||
"description": "Total frames to extract (for count strategy)",
|
||||
},
|
||||
"timestamps": {
|
||||
"type": "array",
|
||||
"items": {"type": "number"},
|
||||
"description": "Specific timestamps in seconds (for timestamps strategy)",
|
||||
},
|
||||
"output_dir": {"type": "string"},
|
||||
"format": {"type": "string", "enum": ["png", "jpg"], "default": "jpg"},
|
||||
"quality": {"type": "integer", "minimum": 1, "maximum": 31, "default": 2},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(cpu_cores=1, ram_mb=512, vram_mb=0, disk_mb=500)
|
||||
idempotency_key_fields = ["input_path", "strategy", "interval_seconds", "count"]
|
||||
side_effects = ["writes frame images to output_dir"]
|
||||
user_visible_verification = ["Inspect extracted frames for representative coverage"]
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
input_path = Path(inputs["input_path"])
|
||||
if not input_path.exists():
|
||||
return ToolResult(success=False, error=f"Input not found: {input_path}")
|
||||
|
||||
strategy = inputs["strategy"]
|
||||
fmt = inputs.get("format", "jpg")
|
||||
quality = inputs.get("quality", 2)
|
||||
output_dir = Path(inputs.get("output_dir", input_path.parent / "frames"))
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
start = time.time()
|
||||
|
||||
try:
|
||||
if strategy == "interval":
|
||||
frames = self._extract_interval(input_path, output_dir, fmt, quality, inputs)
|
||||
elif strategy == "count":
|
||||
frames = self._extract_count(input_path, output_dir, fmt, quality, inputs)
|
||||
elif strategy == "timestamps":
|
||||
frames = self._extract_timestamps(input_path, output_dir, fmt, quality, inputs)
|
||||
else:
|
||||
return ToolResult(success=False, error=f"Unknown strategy: {strategy}")
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=str(e))
|
||||
|
||||
elapsed = time.time() - start
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"strategy": strategy,
|
||||
"frame_count": len(frames),
|
||||
"frames": frames,
|
||||
"output_dir": str(output_dir),
|
||||
},
|
||||
artifacts=[str(output_dir)],
|
||||
duration_seconds=round(elapsed, 2),
|
||||
)
|
||||
|
||||
def _extract_interval(
|
||||
self,
|
||||
input_path: Path,
|
||||
output_dir: Path,
|
||||
fmt: str,
|
||||
quality: int,
|
||||
inputs: dict,
|
||||
) -> list[dict]:
|
||||
interval = inputs.get("interval_seconds", 5.0)
|
||||
output_pattern = str(output_dir / f"frame_%04d.{fmt}")
|
||||
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", str(input_path),
|
||||
"-vf", f"fps=1/{interval}",
|
||||
]
|
||||
if fmt == "jpg":
|
||||
cmd.extend(["-qscale:v", str(quality)])
|
||||
cmd.append(output_pattern)
|
||||
|
||||
self.run_command(cmd)
|
||||
|
||||
return self._collect_frames(output_dir, fmt, interval)
|
||||
|
||||
def _extract_count(
|
||||
self,
|
||||
input_path: Path,
|
||||
output_dir: Path,
|
||||
fmt: str,
|
||||
quality: int,
|
||||
inputs: dict,
|
||||
) -> list[dict]:
|
||||
count = inputs.get("count", 10)
|
||||
duration = self._get_duration(input_path)
|
||||
if duration <= 0:
|
||||
return []
|
||||
|
||||
interval = duration / count
|
||||
output_pattern = str(output_dir / f"frame_%04d.{fmt}")
|
||||
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", str(input_path),
|
||||
"-vf", f"fps=1/{interval}",
|
||||
"-frames:v", str(count),
|
||||
]
|
||||
if fmt == "jpg":
|
||||
cmd.extend(["-qscale:v", str(quality)])
|
||||
cmd.append(output_pattern)
|
||||
|
||||
self.run_command(cmd)
|
||||
|
||||
return self._collect_frames(output_dir, fmt, interval)
|
||||
|
||||
def _extract_timestamps(
|
||||
self,
|
||||
input_path: Path,
|
||||
output_dir: Path,
|
||||
fmt: str,
|
||||
quality: int,
|
||||
inputs: dict,
|
||||
) -> list[dict]:
|
||||
timestamps = inputs.get("timestamps", [])
|
||||
frames = []
|
||||
|
||||
for i, ts in enumerate(timestamps):
|
||||
output_file = output_dir / f"frame_{i:04d}.{fmt}"
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-ss", str(ts),
|
||||
"-i", str(input_path),
|
||||
"-frames:v", "1",
|
||||
]
|
||||
if fmt == "jpg":
|
||||
cmd.extend(["-qscale:v", str(quality)])
|
||||
cmd.append(str(output_file))
|
||||
|
||||
self.run_command(cmd)
|
||||
|
||||
if output_file.exists():
|
||||
frames.append({
|
||||
"path": str(output_file),
|
||||
"timestamp_seconds": ts,
|
||||
"index": i,
|
||||
})
|
||||
|
||||
return frames
|
||||
|
||||
def _get_duration(self, input_path: Path) -> float:
|
||||
"""Get video duration in seconds via ffprobe."""
|
||||
cmd = [
|
||||
"ffprobe",
|
||||
"-v", "quiet",
|
||||
"-show_entries", "format=duration",
|
||||
"-of", "json",
|
||||
str(input_path),
|
||||
]
|
||||
result = self.run_command(cmd)
|
||||
data = json.loads(result.stdout)
|
||||
return float(data.get("format", {}).get("duration", 0))
|
||||
|
||||
def _collect_frames(
|
||||
self, output_dir: Path, fmt: str, interval: float
|
||||
) -> list[dict]:
|
||||
"""Collect extracted frame files and build metadata."""
|
||||
frames = []
|
||||
pattern = f"frame_*.{fmt}"
|
||||
for i, path in enumerate(sorted(output_dir.glob(pattern))):
|
||||
frames.append({
|
||||
"path": str(path),
|
||||
"timestamp_seconds": round(i * interval, 3),
|
||||
"index": i,
|
||||
})
|
||||
return frames
|
||||
@@ -0,0 +1,262 @@
|
||||
"""Scene detection tool wrapping PySceneDetect.
|
||||
|
||||
Detects scene boundaries and shot changes in video. Falls back to
|
||||
FFmpeg-based detection if PySceneDetect is not installed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
ToolResult,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class SceneDetect(BaseTool):
|
||||
name = "scene_detect"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.CORE
|
||||
capability = "analysis"
|
||||
provider = "ffmpeg"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
|
||||
dependencies = ["cmd:ffmpeg"]
|
||||
install_instructions = (
|
||||
"FFmpeg is required. For better detection install PySceneDetect:\n"
|
||||
"pip install scenedetect[opencv]"
|
||||
)
|
||||
agent_skills = ["ffmpeg"]
|
||||
|
||||
capabilities = [
|
||||
"detect_scenes",
|
||||
"detect_content_changes",
|
||||
"detect_threshold",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["input_path"],
|
||||
"properties": {
|
||||
"input_path": {"type": "string"},
|
||||
"method": {
|
||||
"type": "string",
|
||||
"enum": ["content", "threshold", "adaptive"],
|
||||
"default": "content",
|
||||
},
|
||||
"threshold": {
|
||||
"type": "number",
|
||||
"description": "Detection threshold (method-dependent)",
|
||||
},
|
||||
"min_scene_length_seconds": {
|
||||
"type": "number",
|
||||
"minimum": 0.1,
|
||||
"default": 1.0,
|
||||
},
|
||||
"output_path": {"type": "string", "description": "Path for scene list JSON"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(cpu_cores=2, ram_mb=1024, vram_mb=0, disk_mb=100)
|
||||
idempotency_key_fields = ["input_path", "method", "threshold"]
|
||||
side_effects = ["writes scene list JSON to output_path"]
|
||||
user_visible_verification = [
|
||||
"Spot-check detected scene boundaries against the video",
|
||||
]
|
||||
|
||||
def _has_pyscenedetect(self) -> bool:
|
||||
try:
|
||||
import scenedetect # noqa: F401
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
input_path = Path(inputs["input_path"])
|
||||
if not input_path.exists():
|
||||
return ToolResult(success=False, error=f"Input not found: {input_path}")
|
||||
|
||||
start = time.time()
|
||||
|
||||
if self._has_pyscenedetect():
|
||||
scenes = self._detect_pyscenedetect(inputs)
|
||||
else:
|
||||
scenes = self._detect_ffmpeg(inputs)
|
||||
|
||||
elapsed = time.time() - start
|
||||
|
||||
# Write scene list
|
||||
output_path = Path(
|
||||
inputs.get("output_path", str(input_path.with_suffix(".scenes.json")))
|
||||
)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_text(json.dumps({"scenes": scenes}, indent=2), encoding="utf-8")
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"scene_count": len(scenes),
|
||||
"scenes": scenes,
|
||||
"method": "pyscenedetect" if self._has_pyscenedetect() else "ffmpeg",
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
duration_seconds=round(elapsed, 2),
|
||||
)
|
||||
|
||||
def _detect_pyscenedetect(self, inputs: dict[str, Any]) -> list[dict]:
|
||||
"""Use PySceneDetect for scene detection."""
|
||||
from scenedetect import open_video, SceneManager
|
||||
from scenedetect.detectors import ContentDetector, ThresholdDetector, AdaptiveDetector
|
||||
|
||||
input_path = str(inputs["input_path"])
|
||||
method = inputs.get("method", "content")
|
||||
threshold = inputs.get("threshold")
|
||||
min_scene_len = inputs.get("min_scene_length_seconds", 1.0)
|
||||
|
||||
video = open_video(input_path)
|
||||
scene_manager = SceneManager()
|
||||
|
||||
if method == "content":
|
||||
detector = ContentDetector(
|
||||
threshold=threshold or 27.0,
|
||||
min_scene_len=int(min_scene_len * video.frame_rate),
|
||||
)
|
||||
elif method == "threshold":
|
||||
detector = ThresholdDetector(
|
||||
threshold=threshold or 12.0,
|
||||
min_scene_len=int(min_scene_len * video.frame_rate),
|
||||
)
|
||||
elif method == "adaptive":
|
||||
detector = AdaptiveDetector(
|
||||
adaptive_threshold=threshold or 3.0,
|
||||
min_scene_len=int(min_scene_len * video.frame_rate),
|
||||
)
|
||||
else:
|
||||
detector = ContentDetector()
|
||||
|
||||
scene_manager.add_detector(detector)
|
||||
scene_manager.detect_scenes(video)
|
||||
scene_list = scene_manager.get_scene_list()
|
||||
|
||||
scenes = []
|
||||
for i, (scene_start, scene_end) in enumerate(scene_list):
|
||||
scenes.append({
|
||||
"index": i,
|
||||
"start_seconds": round(scene_start.get_seconds(), 3),
|
||||
"end_seconds": round(scene_end.get_seconds(), 3),
|
||||
"duration_seconds": round(
|
||||
scene_end.get_seconds() - scene_start.get_seconds(), 3
|
||||
),
|
||||
})
|
||||
|
||||
return scenes
|
||||
|
||||
def _detect_ffmpeg(self, inputs: dict[str, Any]) -> list[dict]:
|
||||
"""Fallback: use FFmpeg scene change filter."""
|
||||
input_path = str(inputs["input_path"])
|
||||
threshold = inputs.get("threshold", 0.3)
|
||||
min_scene_len = inputs.get("min_scene_length_seconds", 1.0)
|
||||
|
||||
cmd = [
|
||||
"ffprobe",
|
||||
"-v", "quiet",
|
||||
"-show_entries", "frame=pts_time",
|
||||
"-of", "json",
|
||||
"-f", "lavfi",
|
||||
f"movie='{input_path.replace(chr(92), '/').replace(':', chr(92)+':')}',select='gt(scene,{threshold})'",
|
||||
]
|
||||
|
||||
try:
|
||||
result = self.run_command(cmd, timeout=120)
|
||||
data = json.loads(result.stdout)
|
||||
except Exception:
|
||||
# If ffprobe lavfi approach fails, try a simpler method
|
||||
return self._detect_ffmpeg_simple(input_path, threshold, min_scene_len)
|
||||
|
||||
change_points = [0.0]
|
||||
for frame in data.get("frames", []):
|
||||
ts = float(frame.get("pts_time", 0))
|
||||
if ts - change_points[-1] >= min_scene_len:
|
||||
change_points.append(ts)
|
||||
|
||||
# Get total duration
|
||||
dur_cmd = [
|
||||
"ffprobe", "-v", "quiet",
|
||||
"-show_entries", "format=duration",
|
||||
"-of", "json", input_path,
|
||||
]
|
||||
dur_result = self.run_command(dur_cmd)
|
||||
total_dur = float(json.loads(dur_result.stdout)["format"]["duration"])
|
||||
change_points.append(total_dur)
|
||||
|
||||
scenes = []
|
||||
for i in range(len(change_points) - 1):
|
||||
start = change_points[i]
|
||||
end = change_points[i + 1]
|
||||
scenes.append({
|
||||
"index": i,
|
||||
"start_seconds": round(start, 3),
|
||||
"end_seconds": round(end, 3),
|
||||
"duration_seconds": round(end - start, 3),
|
||||
})
|
||||
|
||||
return scenes
|
||||
|
||||
def _detect_ffmpeg_simple(
|
||||
self, input_path: str, threshold: float, min_scene_len: float
|
||||
) -> list[dict]:
|
||||
"""Simplest fallback: split into uniform segments."""
|
||||
dur_cmd = [
|
||||
"ffprobe", "-v", "quiet",
|
||||
"-show_entries", "format=duration",
|
||||
"-of", "json", input_path,
|
||||
]
|
||||
dur_result = self.run_command(dur_cmd)
|
||||
total_dur = float(json.loads(dur_result.stdout)["format"]["duration"])
|
||||
|
||||
# Use select filter to find scene changes via stdout
|
||||
cmd = [
|
||||
"ffmpeg",
|
||||
"-i", input_path,
|
||||
"-vf", f"select='gt(scene,{threshold})',showinfo",
|
||||
"-f", "null", "-",
|
||||
]
|
||||
try:
|
||||
result = self.run_command(cmd, timeout=120)
|
||||
output = result.stderr
|
||||
except Exception:
|
||||
output = ""
|
||||
|
||||
import re
|
||||
change_points = [0.0]
|
||||
for match in re.finditer(r"pts_time:(\d+\.?\d*)", output):
|
||||
ts = float(match.group(1))
|
||||
if ts - change_points[-1] >= min_scene_len:
|
||||
change_points.append(ts)
|
||||
change_points.append(total_dur)
|
||||
|
||||
scenes = []
|
||||
for i in range(len(change_points) - 1):
|
||||
start = change_points[i]
|
||||
end = change_points[i + 1]
|
||||
scenes.append({
|
||||
"index": i,
|
||||
"start_seconds": round(start, 3),
|
||||
"end_seconds": round(end, 3),
|
||||
"duration_seconds": round(end - start, 3),
|
||||
})
|
||||
|
||||
return scenes
|
||||
@@ -0,0 +1,251 @@
|
||||
"""Transcription tool wrapping faster-whisper / WhisperX.
|
||||
|
||||
Provides speech-to-text with word-level timestamps and optional speaker
|
||||
diarization. Falls back gracefully when GPU or diarization dependencies
|
||||
are not available.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ResumeSupport,
|
||||
ToolResult,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class Transcriber(BaseTool):
|
||||
name = "transcriber"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.CORE
|
||||
capability = "analysis"
|
||||
provider = "whisperx"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
|
||||
dependencies = ["python:faster_whisper"]
|
||||
install_instructions = (
|
||||
"pip install faster-whisper # CPU mode\n"
|
||||
"pip install faster-whisper[gpu] # GPU mode (requires CUDA)\n"
|
||||
"pip install whisperx # For diarization support"
|
||||
)
|
||||
agent_skills = ["speech-to-text"]
|
||||
|
||||
capabilities = [
|
||||
"transcribe",
|
||||
"word_timestamps",
|
||||
"diarization",
|
||||
"language_detection",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["input_path"],
|
||||
"properties": {
|
||||
"input_path": {"type": "string", "description": "Path to audio or video file"},
|
||||
"model_size": {
|
||||
"type": "string",
|
||||
"enum": ["tiny", "base", "small", "medium", "large-v2", "large-v3"],
|
||||
"default": "base",
|
||||
},
|
||||
"language": {"type": "string", "description": "ISO 639-1 language code, or null for auto-detect"},
|
||||
"diarize": {"type": "boolean", "default": False},
|
||||
"output_dir": {"type": "string", "description": "Directory for output files"},
|
||||
},
|
||||
}
|
||||
|
||||
output_schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"segments": {"type": "array"},
|
||||
"word_timestamps": {"type": "array"},
|
||||
"language": {"type": "string"},
|
||||
"duration_seconds": {"type": "number"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=2,
|
||||
ram_mb=2048,
|
||||
vram_mb=0, # CPU by default; GPU optional
|
||||
disk_mb=500,
|
||||
network_required=False,
|
||||
)
|
||||
|
||||
retry_policy = RetryPolicy(max_retries=1, retryable_errors=["MemoryError"])
|
||||
resume_support = ResumeSupport.FROM_START
|
||||
idempotency_key_fields = ["input_path", "model_size", "language"]
|
||||
side_effects = ["writes transcript JSON to output_dir"]
|
||||
fallback = None
|
||||
user_visible_verification = [
|
||||
"Check transcript text against source audio",
|
||||
"Verify word timestamps align with speech",
|
||||
]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
try:
|
||||
import faster_whisper # noqa: F401
|
||||
return ToolStatus.AVAILABLE
|
||||
except ImportError:
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def _has_diarization(self) -> bool:
|
||||
try:
|
||||
import whisperx # noqa: F401
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
def estimate_runtime(self, inputs: dict[str, Any]) -> float:
|
||||
"""Rough estimate: ~0.5x real-time on CPU for 'base' model."""
|
||||
return 60.0 # conservative default
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
input_path = Path(inputs["input_path"])
|
||||
model_size = inputs.get("model_size", "base")
|
||||
language = inputs.get("language")
|
||||
diarize = inputs.get("diarize", False)
|
||||
output_dir = Path(inputs.get("output_dir", input_path.parent))
|
||||
|
||||
if not input_path.exists():
|
||||
return ToolResult(success=False, error=f"Input file not found: {input_path}")
|
||||
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
try:
|
||||
from faster_whisper import WhisperModel
|
||||
except ImportError:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="faster-whisper is not installed. Run: pip install faster-whisper",
|
||||
)
|
||||
|
||||
start = time.time()
|
||||
|
||||
# Load model (CPU by default, CUDA if available)
|
||||
try:
|
||||
import torch
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
compute_type = "float16" if device == "cuda" else "int8"
|
||||
except ImportError:
|
||||
device = "cpu"
|
||||
compute_type = "int8"
|
||||
|
||||
model = WhisperModel(model_size, device=device, compute_type=compute_type)
|
||||
|
||||
# Transcribe
|
||||
segments_iter, info = model.transcribe(
|
||||
str(input_path),
|
||||
language=language,
|
||||
word_timestamps=True,
|
||||
vad_filter=True,
|
||||
)
|
||||
|
||||
segments = []
|
||||
word_timestamps = []
|
||||
|
||||
for seg in segments_iter:
|
||||
seg_data = {
|
||||
"id": seg.id,
|
||||
"start": round(seg.start, 3),
|
||||
"end": round(seg.end, 3),
|
||||
"text": seg.text.strip(),
|
||||
}
|
||||
|
||||
if seg.words:
|
||||
words = []
|
||||
for w in seg.words:
|
||||
word_entry = {
|
||||
"word": w.word,
|
||||
"start": round(w.start, 3),
|
||||
"end": round(w.end, 3),
|
||||
"probability": round(w.probability, 3),
|
||||
}
|
||||
words.append(word_entry)
|
||||
word_timestamps.append(word_entry)
|
||||
seg_data["words"] = words
|
||||
|
||||
segments.append(seg_data)
|
||||
|
||||
detected_language = language or info.language
|
||||
duration = info.duration
|
||||
|
||||
# Optional diarization pass
|
||||
if diarize and self._has_diarization():
|
||||
segments = self._apply_diarization(
|
||||
str(input_path), segments, detected_language
|
||||
)
|
||||
|
||||
elapsed = time.time() - start
|
||||
|
||||
result_data = {
|
||||
"segments": segments,
|
||||
"word_timestamps": word_timestamps,
|
||||
"language": detected_language,
|
||||
"duration_seconds": round(duration, 3),
|
||||
"model_size": model_size,
|
||||
"device": device,
|
||||
}
|
||||
|
||||
# Write transcript JSON
|
||||
output_path = output_dir / f"{input_path.stem}_transcript.json"
|
||||
output_path.write_text(json.dumps(result_data, indent=2), encoding="utf-8")
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data=result_data,
|
||||
artifacts=[str(output_path)],
|
||||
duration_seconds=round(elapsed, 2),
|
||||
)
|
||||
|
||||
def _apply_diarization(
|
||||
self,
|
||||
audio_path: str,
|
||||
segments: list[dict],
|
||||
language: str,
|
||||
) -> list[dict]:
|
||||
"""Apply WhisperX diarization to assign speaker labels."""
|
||||
try:
|
||||
import whisperx
|
||||
|
||||
# Load audio for alignment
|
||||
audio = whisperx.load_audio(audio_path)
|
||||
|
||||
# Align segments with word timestamps
|
||||
align_model, align_metadata = whisperx.load_align_model(
|
||||
language_code=language, device="cpu"
|
||||
)
|
||||
aligned = whisperx.align(
|
||||
segments, align_model, align_metadata, audio, device="cpu"
|
||||
)
|
||||
|
||||
# Diarize
|
||||
import os
|
||||
hf_token = os.environ.get("HF_TOKEN")
|
||||
if not hf_token:
|
||||
# Can't diarize without HuggingFace token for pyannote
|
||||
return segments
|
||||
|
||||
diarize_model = whisperx.DiarizationPipeline(
|
||||
use_auth_token=hf_token, device="cpu"
|
||||
)
|
||||
diarize_segments = diarize_model(audio)
|
||||
result = whisperx.assign_word_speakers(diarize_segments, aligned)
|
||||
|
||||
return result.get("segments", segments)
|
||||
except Exception:
|
||||
# Diarization is best-effort; return original segments on failure
|
||||
return segments
|
||||
@@ -0,0 +1,597 @@
|
||||
"""Video/image understanding tool using vision-language models.
|
||||
|
||||
Analyzes images or video frames using CLIP, BLIP-2, or LLaVA. Supports
|
||||
frame description, visual question answering, quality assessment, and
|
||||
scene classification. Primary use case: visual QA for automated video review.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
VIDEO_EXTENSIONS = {".mp4", ".mov", ".avi", ".webm"}
|
||||
IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".webp"}
|
||||
|
||||
SCENE_CATEGORIES = [
|
||||
"indoor", "outdoor", "landscape", "cityscape", "portrait",
|
||||
"action", "close-up", "aerial", "underwater", "night",
|
||||
"studio", "nature", "urban", "abstract", "text-overlay",
|
||||
]
|
||||
|
||||
|
||||
class VideoUnderstand(BaseTool):
|
||||
name = "video_understand"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.ANALYZE
|
||||
capability = "analysis"
|
||||
provider = "transformers"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
runtime = ToolRuntime.LOCAL_GPU
|
||||
|
||||
dependencies = ["python:transformers", "python:torch"]
|
||||
install_instructions = (
|
||||
"pip install transformers torch # For CLIP/BLIP-2 visual understanding"
|
||||
)
|
||||
agent_skills = ["video-understand"]
|
||||
|
||||
capabilities = [
|
||||
"image_description",
|
||||
"visual_qa",
|
||||
"quality_assessment",
|
||||
"scene_classification",
|
||||
"object_detection",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["input_path"],
|
||||
"properties": {
|
||||
"input_path": {
|
||||
"type": "string",
|
||||
"description": "Path to image or video file",
|
||||
},
|
||||
"query": {
|
||||
"type": "string",
|
||||
"description": "Question to answer about the visual content (for VQA mode)",
|
||||
},
|
||||
"mode": {
|
||||
"type": "string",
|
||||
"enum": ["describe", "qa", "quality", "classify"],
|
||||
"default": "describe",
|
||||
"description": (
|
||||
"describe: generate caption; qa: answer query about content; "
|
||||
"quality: assess technical quality; classify: classify scene type"
|
||||
),
|
||||
},
|
||||
"model": {
|
||||
"type": "string",
|
||||
"enum": ["clip", "blip2", "llava"],
|
||||
"default": "clip",
|
||||
"description": "Which vision-language model to use",
|
||||
},
|
||||
"frame_indices": {
|
||||
"type": "array",
|
||||
"items": {"type": "integer"},
|
||||
"description": (
|
||||
"For video input, which frames to analyze. "
|
||||
"If not provided, samples key frames at even intervals."
|
||||
),
|
||||
},
|
||||
"max_frames": {
|
||||
"type": "integer",
|
||||
"default": 5,
|
||||
"description": "Maximum number of frames to analyze from video",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
output_schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"frames": {
|
||||
"type": "array",
|
||||
"description": "Per-frame analysis results",
|
||||
},
|
||||
"summary": {"type": "string"},
|
||||
"mode": {"type": "string"},
|
||||
"model": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=2,
|
||||
ram_mb=4096,
|
||||
vram_mb=2048,
|
||||
disk_mb=1000,
|
||||
network_required=False,
|
||||
)
|
||||
|
||||
idempotency_key_fields = ["input_path", "mode", "model", "query"]
|
||||
side_effects = []
|
||||
fallback = None
|
||||
user_visible_verification = [
|
||||
"Compare generated descriptions against actual visual content",
|
||||
"Verify quality scores match perceived image quality",
|
||||
]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
"""Check if transformers and torch are both importable."""
|
||||
try:
|
||||
import transformers # noqa: F401
|
||||
import torch # noqa: F401
|
||||
return ToolStatus.AVAILABLE
|
||||
except ImportError:
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_runtime(self, inputs: dict[str, Any]) -> float:
|
||||
"""Estimate runtime in seconds based on mode and frame count."""
|
||||
max_frames = inputs.get("max_frames", 5)
|
||||
mode = inputs.get("mode", "describe")
|
||||
if mode == "quality":
|
||||
return max_frames * 0.5 # quality metrics are fast
|
||||
return max_frames * 5.0 # VLM inference per frame
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
input_path = Path(inputs["input_path"])
|
||||
mode = inputs.get("mode", "describe")
|
||||
model_name = inputs.get("model", "clip")
|
||||
query = inputs.get("query")
|
||||
frame_indices = inputs.get("frame_indices")
|
||||
max_frames = inputs.get("max_frames", 5)
|
||||
|
||||
if not input_path.exists():
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"Input file not found: {input_path}",
|
||||
)
|
||||
|
||||
suffix = input_path.suffix.lower()
|
||||
is_video = suffix in VIDEO_EXTENSIONS
|
||||
is_image = suffix in IMAGE_EXTENSIONS
|
||||
|
||||
if not is_video and not is_image:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=(
|
||||
f"Unsupported file type: {suffix}. "
|
||||
f"Supported: {sorted(VIDEO_EXTENSIONS | IMAGE_EXTENSIONS)}"
|
||||
),
|
||||
)
|
||||
|
||||
if mode == "qa" and not query:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="Query is required for 'qa' mode.",
|
||||
)
|
||||
|
||||
start = time.time()
|
||||
|
||||
# --- Load frames ---
|
||||
try:
|
||||
frames = self._load_frames(
|
||||
input_path, is_video, frame_indices, max_frames
|
||||
)
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Failed to load frames: {e}")
|
||||
|
||||
if not frames:
|
||||
return ToolResult(success=False, error="No frames could be extracted.")
|
||||
|
||||
# --- Analyze each frame ---
|
||||
try:
|
||||
if mode == "quality":
|
||||
frame_results = self._analyze_quality(frames)
|
||||
elif mode == "describe":
|
||||
frame_results = self._analyze_describe(frames, model_name)
|
||||
elif mode == "qa":
|
||||
frame_results = self._analyze_qa(frames, model_name, query)
|
||||
elif mode == "classify":
|
||||
frame_results = self._analyze_classify(frames, model_name)
|
||||
else:
|
||||
return ToolResult(success=False, error=f"Unknown mode: {mode}")
|
||||
except ImportError as e:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"Missing dependency for {model_name}: {e}",
|
||||
)
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Analysis failed: {e}")
|
||||
|
||||
elapsed = time.time() - start
|
||||
|
||||
# --- Build summary ---
|
||||
summary = self._build_summary(frame_results, mode)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"frames": frame_results,
|
||||
"summary": summary,
|
||||
"mode": mode,
|
||||
"model": model_name if mode != "quality" else "metrics",
|
||||
"frame_count": len(frame_results),
|
||||
},
|
||||
duration_seconds=round(elapsed, 2),
|
||||
model=model_name if mode != "quality" else None,
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Frame extraction
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _load_frames(
|
||||
self,
|
||||
input_path: Path,
|
||||
is_video: bool,
|
||||
frame_indices: list[int] | None,
|
||||
max_frames: int,
|
||||
) -> list:
|
||||
"""Load PIL Image objects from an image or video file."""
|
||||
from PIL import Image
|
||||
|
||||
if not is_video:
|
||||
return [Image.open(input_path).convert("RGB")]
|
||||
|
||||
return self._extract_video_frames(input_path, frame_indices, max_frames)
|
||||
|
||||
def _extract_video_frames(
|
||||
self,
|
||||
video_path: Path,
|
||||
frame_indices: list[int] | None,
|
||||
max_frames: int,
|
||||
) -> list:
|
||||
"""Extract frames from a video file using ffmpeg."""
|
||||
from PIL import Image
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp = Path(tmp_dir)
|
||||
|
||||
if frame_indices:
|
||||
# Extract specific frames using select filter
|
||||
frames_to_extract = frame_indices[:max_frames]
|
||||
select_expr = "+".join(
|
||||
f"eq(n\\,{idx})" for idx in frames_to_extract
|
||||
)
|
||||
cmd = [
|
||||
"ffmpeg", "-i", str(video_path),
|
||||
"-vf", f"select='{select_expr}'",
|
||||
"-vsync", "vfr",
|
||||
str(tmp / "frame_%04d.png"),
|
||||
"-y", "-loglevel", "error",
|
||||
]
|
||||
else:
|
||||
# Get total frame count first
|
||||
probe_cmd = [
|
||||
"ffmpeg", "-i", str(video_path),
|
||||
"-map", "0:v:0", "-c", "copy", "-f", "null", "-",
|
||||
]
|
||||
# Sample at even intervals using fps filter
|
||||
# Use a select filter that picks frames at even intervals
|
||||
cmd = [
|
||||
"ffmpeg", "-i", str(video_path),
|
||||
"-frames:v", str(max_frames),
|
||||
"-vf", f"thumbnail={max_frames}",
|
||||
str(tmp / "frame_%04d.png"),
|
||||
"-y", "-loglevel", "error",
|
||||
]
|
||||
|
||||
subprocess.run(cmd, capture_output=True, text=True, timeout=60)
|
||||
|
||||
# Load extracted frames
|
||||
frame_files = sorted(tmp.glob("frame_*.png"))
|
||||
images = []
|
||||
for f in frame_files[:max_frames]:
|
||||
images.append(Image.open(f).convert("RGB"))
|
||||
|
||||
return images
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Quality assessment (uses PIL/numpy, no VLM needed)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _analyze_quality(self, frames: list) -> list[dict[str, Any]]:
|
||||
"""Assess technical quality using simple image metrics."""
|
||||
import numpy as np
|
||||
|
||||
results = []
|
||||
for i, img in enumerate(frames):
|
||||
arr = np.array(img, dtype=np.float64)
|
||||
gray = np.mean(arr, axis=2)
|
||||
|
||||
# Blur detection: Laplacian variance (low = blurry)
|
||||
# Manual Laplacian approximation using numpy
|
||||
laplacian_kernel = np.array([[0, 1, 0], [1, -4, 1], [0, 1, 0]])
|
||||
from scipy.signal import convolve2d
|
||||
laplacian = convolve2d(gray, laplacian_kernel, mode="valid")
|
||||
blur_score = float(np.var(laplacian))
|
||||
|
||||
# Brightness: mean pixel value (0-255 scale)
|
||||
brightness = float(np.mean(arr))
|
||||
|
||||
# Contrast: standard deviation of pixel values
|
||||
contrast = float(np.std(arr))
|
||||
|
||||
# Classify quality
|
||||
quality_issues = []
|
||||
if blur_score < 100:
|
||||
quality_issues.append("blurry")
|
||||
if brightness < 40:
|
||||
quality_issues.append("underexposed")
|
||||
elif brightness > 220:
|
||||
quality_issues.append("overexposed")
|
||||
if contrast < 30:
|
||||
quality_issues.append("low_contrast")
|
||||
|
||||
quality_label = "good" if not quality_issues else "issues_detected"
|
||||
|
||||
results.append({
|
||||
"frame_index": i,
|
||||
"blur_score": round(blur_score, 2),
|
||||
"brightness": round(brightness, 2),
|
||||
"contrast": round(contrast, 2),
|
||||
"quality": quality_label,
|
||||
"issues": quality_issues,
|
||||
"resolution": f"{img.width}x{img.height}",
|
||||
})
|
||||
|
||||
return results
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# VLM-based analysis modes
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _load_model(self, model_name: str):
|
||||
"""Load the requested vision-language model and processor."""
|
||||
import torch
|
||||
from transformers import (
|
||||
CLIPProcessor,
|
||||
CLIPModel,
|
||||
BlipProcessor,
|
||||
BlipForConditionalGeneration,
|
||||
Blip2Processor,
|
||||
Blip2ForConditionalGeneration,
|
||||
AutoProcessor,
|
||||
AutoModelForCausalLM,
|
||||
)
|
||||
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
if model_name == "clip":
|
||||
model_id = "openai/clip-vit-base-patch32"
|
||||
processor = CLIPProcessor.from_pretrained(model_id)
|
||||
model = CLIPModel.from_pretrained(model_id).to(device)
|
||||
return model, processor, device
|
||||
|
||||
if model_name == "blip2":
|
||||
model_id = "Salesforce/blip2-opt-2.7b"
|
||||
processor = Blip2Processor.from_pretrained(model_id)
|
||||
model = Blip2ForConditionalGeneration.from_pretrained(
|
||||
model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32
|
||||
).to(device)
|
||||
return model, processor, device
|
||||
|
||||
if model_name == "llava":
|
||||
model_id = "llava-hf/llava-1.5-7b-hf"
|
||||
processor = AutoProcessor.from_pretrained(model_id)
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32
|
||||
).to(device)
|
||||
return model, processor, device
|
||||
|
||||
raise ValueError(f"Unknown model: {model_name}")
|
||||
|
||||
def _analyze_describe(
|
||||
self, frames: list, model_name: str
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Generate captions for each frame."""
|
||||
import torch
|
||||
|
||||
model, processor, device = self._load_model(model_name)
|
||||
results = []
|
||||
|
||||
for i, img in enumerate(frames):
|
||||
if model_name == "clip":
|
||||
# CLIP is not a captioning model; use zero-shot classification
|
||||
# with generic scene descriptions as a caption proxy
|
||||
candidate_texts = [
|
||||
"a photo of a person", "a photo of a landscape",
|
||||
"a photo of an object", "a photo of text",
|
||||
"a photo of an animal", "a photo of a building",
|
||||
"a photo of food", "a photo of a vehicle",
|
||||
"an abstract image", "a dark scene", "a bright scene",
|
||||
]
|
||||
clip_inputs = processor(
|
||||
text=candidate_texts, images=img, return_tensors="pt", padding=True
|
||||
).to(device)
|
||||
with torch.no_grad():
|
||||
outputs = model(**clip_inputs)
|
||||
probs = outputs.logits_per_image.softmax(dim=1)[0]
|
||||
top_idx = probs.argmax().item()
|
||||
caption = candidate_texts[top_idx]
|
||||
confidence = round(probs[top_idx].item(), 3)
|
||||
results.append({
|
||||
"frame_index": i,
|
||||
"description": caption,
|
||||
"confidence": confidence,
|
||||
})
|
||||
|
||||
elif model_name in ("blip2", "llava"):
|
||||
inputs = processor(images=img, return_tensors="pt").to(device)
|
||||
with torch.no_grad():
|
||||
generated_ids = model.generate(**inputs, max_new_tokens=50)
|
||||
caption = processor.batch_decode(
|
||||
generated_ids, skip_special_tokens=True
|
||||
)[0].strip()
|
||||
results.append({
|
||||
"frame_index": i,
|
||||
"description": caption,
|
||||
})
|
||||
|
||||
return results
|
||||
|
||||
def _analyze_qa(
|
||||
self, frames: list, model_name: str, query: str
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Answer a question about each frame."""
|
||||
import torch
|
||||
|
||||
model, processor, device = self._load_model(model_name)
|
||||
results = []
|
||||
|
||||
for i, img in enumerate(frames):
|
||||
if model_name == "clip":
|
||||
# Use query as one candidate and its negation as another
|
||||
candidates = [query, f"not {query}"]
|
||||
clip_inputs = processor(
|
||||
text=candidates, images=img, return_tensors="pt", padding=True
|
||||
).to(device)
|
||||
with torch.no_grad():
|
||||
outputs = model(**clip_inputs)
|
||||
probs = outputs.logits_per_image.softmax(dim=1)[0]
|
||||
yes_prob = probs[0].item()
|
||||
results.append({
|
||||
"frame_index": i,
|
||||
"query": query,
|
||||
"answer": "yes" if yes_prob > 0.5 else "no",
|
||||
"confidence": round(max(yes_prob, 1 - yes_prob), 3),
|
||||
})
|
||||
|
||||
elif model_name in ("blip2", "llava"):
|
||||
prompt = f"Question: {query} Answer:"
|
||||
inputs = processor(
|
||||
images=img, text=prompt, return_tensors="pt"
|
||||
).to(device)
|
||||
with torch.no_grad():
|
||||
generated_ids = model.generate(**inputs, max_new_tokens=50)
|
||||
answer = processor.batch_decode(
|
||||
generated_ids, skip_special_tokens=True
|
||||
)[0].strip()
|
||||
# Remove the prompt echo if present
|
||||
if answer.startswith(prompt):
|
||||
answer = answer[len(prompt):].strip()
|
||||
results.append({
|
||||
"frame_index": i,
|
||||
"query": query,
|
||||
"answer": answer,
|
||||
})
|
||||
|
||||
return results
|
||||
|
||||
def _analyze_classify(
|
||||
self, frames: list, model_name: str
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Classify each frame into scene categories."""
|
||||
import torch
|
||||
|
||||
model, processor, device = self._load_model(model_name)
|
||||
results = []
|
||||
|
||||
for i, img in enumerate(frames):
|
||||
if model_name == "clip":
|
||||
candidate_texts = [f"a {cat} scene" for cat in SCENE_CATEGORIES]
|
||||
clip_inputs = processor(
|
||||
text=candidate_texts, images=img, return_tensors="pt", padding=True
|
||||
).to(device)
|
||||
with torch.no_grad():
|
||||
outputs = model(**clip_inputs)
|
||||
probs = outputs.logits_per_image.softmax(dim=1)[0]
|
||||
|
||||
scored = sorted(
|
||||
zip(SCENE_CATEGORIES, probs.tolist()),
|
||||
key=lambda x: x[1],
|
||||
reverse=True,
|
||||
)
|
||||
results.append({
|
||||
"frame_index": i,
|
||||
"top_category": scored[0][0],
|
||||
"confidence": round(scored[0][1], 3),
|
||||
"categories": [
|
||||
{"label": label, "score": round(score, 3)}
|
||||
for label, score in scored[:5]
|
||||
],
|
||||
})
|
||||
|
||||
elif model_name in ("blip2", "llava"):
|
||||
prompt = (
|
||||
"Classify this image into one of these categories: "
|
||||
+ ", ".join(SCENE_CATEGORIES)
|
||||
+ ". Category:"
|
||||
)
|
||||
inputs = processor(
|
||||
images=img, text=prompt, return_tensors="pt"
|
||||
).to(device)
|
||||
with torch.no_grad():
|
||||
generated_ids = model.generate(**inputs, max_new_tokens=20)
|
||||
category = processor.batch_decode(
|
||||
generated_ids, skip_special_tokens=True
|
||||
)[0].strip()
|
||||
if category.startswith(prompt):
|
||||
category = category[len(prompt):].strip()
|
||||
results.append({
|
||||
"frame_index": i,
|
||||
"top_category": category,
|
||||
})
|
||||
|
||||
return results
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Summary
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _build_summary(
|
||||
self, frame_results: list[dict[str, Any]], mode: str
|
||||
) -> str:
|
||||
"""Build a human-readable summary from per-frame results."""
|
||||
n = len(frame_results)
|
||||
|
||||
if mode == "describe":
|
||||
descriptions = [r.get("description", "") for r in frame_results]
|
||||
if n == 1:
|
||||
return descriptions[0]
|
||||
return f"Analyzed {n} frames. Descriptions: " + "; ".join(descriptions)
|
||||
|
||||
if mode == "qa":
|
||||
answers = [r.get("answer", "") for r in frame_results]
|
||||
if n == 1:
|
||||
return answers[0]
|
||||
return f"Analyzed {n} frames. Answers: " + "; ".join(answers)
|
||||
|
||||
if mode == "quality":
|
||||
issues_all = []
|
||||
for r in frame_results:
|
||||
issues_all.extend(r.get("issues", []))
|
||||
if not issues_all:
|
||||
return f"All {n} frame(s) passed quality checks."
|
||||
unique_issues = sorted(set(issues_all))
|
||||
return (
|
||||
f"Analyzed {n} frame(s). Issues found: {', '.join(unique_issues)}."
|
||||
)
|
||||
|
||||
if mode == "classify":
|
||||
categories = [r.get("top_category", "unknown") for r in frame_results]
|
||||
if n == 1:
|
||||
return f"Scene classified as: {categories[0]}"
|
||||
return (
|
||||
f"Analyzed {n} frames. Scene categories: "
|
||||
+ ", ".join(categories)
|
||||
)
|
||||
|
||||
return f"Analyzed {n} frame(s)."
|
||||
@@ -0,0 +1 @@
|
||||
"""Audio tools — TTS providers, audio processing, and music generation."""
|
||||
@@ -0,0 +1,186 @@
|
||||
"""Audio enhancement tool for noise reduction and cleanup.
|
||||
|
||||
Provides noise reduction, normalization, and EQ via FFmpeg audio
|
||||
filters. Optional pedalboard integration for higher-quality
|
||||
processing when available.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
ToolResult,
|
||||
ToolStability,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
PRESETS = {
|
||||
"clean_speech": {
|
||||
"description": "Noise gate + highpass + compressor + limiter for clean dialogue",
|
||||
"af": (
|
||||
"highpass=f=80,"
|
||||
"lowpass=f=13000,"
|
||||
"agate=threshold=0.01:ratio=2:attack=5:release=50,"
|
||||
"acompressor=threshold=-20dB:ratio=3:attack=5:release=100,"
|
||||
"loudnorm=I=-16:LRA=11:TP=-1.5"
|
||||
),
|
||||
},
|
||||
"noise_reduce": {
|
||||
"description": "Aggressive noise reduction for noisy environments",
|
||||
"af": (
|
||||
"afftdn=nf=-25:nt=w,"
|
||||
"highpass=f=100,"
|
||||
"loudnorm=I=-16:LRA=11:TP=-1.5"
|
||||
),
|
||||
},
|
||||
"normalize_only": {
|
||||
"description": "Loudness normalization without other processing",
|
||||
"af": "loudnorm=I=-16:LRA=11:TP=-1.5",
|
||||
},
|
||||
"podcast": {
|
||||
"description": "Podcast-style processing: de-ess, compress, normalize",
|
||||
"af": (
|
||||
"highpass=f=80,"
|
||||
"acompressor=threshold=-18dB:ratio=4:attack=5:release=100:makeup=2,"
|
||||
"loudnorm=I=-16:LRA=7:TP=-1.5"
|
||||
),
|
||||
},
|
||||
"broadcast": {
|
||||
"description": "Broadcast-standard processing with tight dynamics",
|
||||
"af": (
|
||||
"highpass=f=80,"
|
||||
"lowpass=f=15000,"
|
||||
"acompressor=threshold=-24dB:ratio=4:attack=5:release=80:makeup=3,"
|
||||
"alimiter=limit=0.95:attack=1:release=10,"
|
||||
"loudnorm=I=-24:LRA=7:TP=-2"
|
||||
),
|
||||
},
|
||||
"voice_clarity": {
|
||||
"description": "Boost vocal presence with EQ and light compression",
|
||||
"af": (
|
||||
"highpass=f=80,"
|
||||
"equalizer=f=200:t=q:w=1.5:g=-3,"
|
||||
"equalizer=f=3000:t=q:w=1.0:g=3,"
|
||||
"equalizer=f=5000:t=q:w=1.5:g=2,"
|
||||
"acompressor=threshold=-20dB:ratio=2.5:attack=10:release=100,"
|
||||
"loudnorm=I=-16:LRA=11:TP=-1.5"
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class AudioEnhance(BaseTool):
|
||||
name = "audio_enhance"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.CORE
|
||||
capability = "audio_processing"
|
||||
provider = "ffmpeg"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
|
||||
dependencies = ["cmd:ffmpeg"]
|
||||
install_instructions = "Install FFmpeg: https://ffmpeg.org/download.html"
|
||||
agent_skills = ["ffmpeg", "elevenlabs"]
|
||||
|
||||
capabilities = [
|
||||
"noise_reduction",
|
||||
"normalization",
|
||||
"compression",
|
||||
"eq",
|
||||
"speech_cleanup",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["input_path"],
|
||||
"properties": {
|
||||
"input_path": {"type": "string"},
|
||||
"output_path": {"type": "string"},
|
||||
"preset": {
|
||||
"type": "string",
|
||||
"enum": list(PRESETS.keys()),
|
||||
"default": "clean_speech",
|
||||
},
|
||||
"custom_af": {
|
||||
"type": "string",
|
||||
"description": "Custom FFmpeg audio filter string",
|
||||
},
|
||||
"audio_codec": {"type": "string", "default": "aac"},
|
||||
"audio_bitrate": {"type": "string", "default": "192k"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(cpu_cores=1, ram_mb=512, vram_mb=0, disk_mb=500)
|
||||
idempotency_key_fields = ["input_path", "preset", "custom_af"]
|
||||
side_effects = ["writes enhanced audio/video to output_path"]
|
||||
user_visible_verification = [
|
||||
"Listen to enhanced audio and compare with original",
|
||||
"Verify speech is clear without artifacts or pumping",
|
||||
]
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
input_path = Path(inputs["input_path"])
|
||||
if not input_path.exists():
|
||||
return ToolResult(success=False, error=f"Input not found: {input_path}")
|
||||
|
||||
output_path = Path(
|
||||
inputs.get("output_path", str(input_path.with_stem(f"{input_path.stem}_enhanced")))
|
||||
)
|
||||
audio_codec = inputs.get("audio_codec", "aac")
|
||||
audio_bitrate = inputs.get("audio_bitrate", "192k")
|
||||
|
||||
af = inputs.get("custom_af")
|
||||
if not af:
|
||||
preset_name = inputs.get("preset", "clean_speech")
|
||||
preset = PRESETS.get(preset_name)
|
||||
if not preset:
|
||||
return ToolResult(success=False, error=f"Unknown preset: {preset_name}")
|
||||
af = preset["af"]
|
||||
|
||||
start = time.time()
|
||||
|
||||
# Determine if input is video or audio-only
|
||||
is_video = input_path.suffix.lower() in {".mp4", ".mkv", ".avi", ".mov", ".webm"}
|
||||
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", str(input_path),
|
||||
"-af", af,
|
||||
]
|
||||
if is_video:
|
||||
cmd.extend(["-c:v", "copy"])
|
||||
cmd.extend(["-c:a", audio_codec, "-b:a", audio_bitrate])
|
||||
cmd.append(str(output_path))
|
||||
|
||||
try:
|
||||
self.run_command(cmd)
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"FFmpeg failed: {e}")
|
||||
|
||||
elapsed = time.time() - start
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"input": str(input_path),
|
||||
"output": str(output_path),
|
||||
"preset": inputs.get("preset"),
|
||||
"filter": af,
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
duration_seconds=round(elapsed, 2),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def list_presets() -> dict[str, str]:
|
||||
"""Return available presets and their descriptions."""
|
||||
return {name: p["description"] for name, p in PRESETS.items()}
|
||||
@@ -0,0 +1,560 @@
|
||||
"""Audio mixer tool wrapping FFmpeg and pydub.
|
||||
|
||||
Mixes speech, music, and SFX tracks with support for ducking, fades,
|
||||
and volume normalization. Falls back to FFmpeg-only mode if pydub is
|
||||
not installed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
ToolResult,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class AudioMixer(BaseTool):
|
||||
name = "audio_mixer"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.CORE
|
||||
capability = "audio_processing"
|
||||
provider = "ffmpeg"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
|
||||
dependencies = ["cmd:ffmpeg"]
|
||||
install_instructions = (
|
||||
"FFmpeg is required. pydub is optional for advanced mixing:\n"
|
||||
"pip install pydub"
|
||||
)
|
||||
agent_skills = ["ffmpeg", "video_toolkit"]
|
||||
|
||||
capabilities = ["mix", "duck", "fade", "normalize", "extract_audio"]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["operation"],
|
||||
"properties": {
|
||||
"operation": {
|
||||
"type": "string",
|
||||
"enum": ["mix", "duck", "extract", "full_mix"],
|
||||
"description": (
|
||||
"mix: layer multiple tracks with volume/delay/fades. "
|
||||
"duck: lower music volume when speech is present. "
|
||||
"extract: extract audio from video file. "
|
||||
"full_mix: combine narration tracks + music with ducking + normalize "
|
||||
"in a single call (preferred for compose-director)."
|
||||
),
|
||||
},
|
||||
"tracks": {
|
||||
"type": "array",
|
||||
"description": (
|
||||
"Audio tracks for mix/duck operations (advanced format). "
|
||||
"For duck, each track needs a 'role' of 'speech' or 'music'. "
|
||||
"For the simple duck API, use primary_audio/secondary_audio instead."
|
||||
),
|
||||
"items": {
|
||||
"type": "object",
|
||||
"required": ["path", "role"],
|
||||
"properties": {
|
||||
"path": {"type": "string"},
|
||||
"role": {
|
||||
"type": "string",
|
||||
"enum": ["speech", "music", "sfx", "primary", "secondary"],
|
||||
},
|
||||
"volume": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"maximum": 1.0,
|
||||
"default": 1.0,
|
||||
},
|
||||
"start_seconds": {"type": "number", "minimum": 0},
|
||||
"fade_in_seconds": {"type": "number", "minimum": 0},
|
||||
"fade_out_seconds": {"type": "number", "minimum": 0},
|
||||
},
|
||||
},
|
||||
},
|
||||
"primary_audio": {
|
||||
"type": "string",
|
||||
"description": (
|
||||
"Path to primary/speech audio track (duck operation, simple format). "
|
||||
"This is the track that stays at full volume (e.g. narration/dialogue). "
|
||||
"Use with secondary_audio as an alternative to the tracks array."
|
||||
),
|
||||
},
|
||||
"secondary_audio": {
|
||||
"type": "string",
|
||||
"description": (
|
||||
"Path to secondary/music audio track (duck operation, simple format). "
|
||||
"This track gets ducked (volume lowered) when primary audio is present. "
|
||||
"Use with primary_audio as an alternative to the tracks array."
|
||||
),
|
||||
},
|
||||
"duck_level": {
|
||||
"type": "number",
|
||||
"description": (
|
||||
"Ducking attenuation in dB for the secondary track (duck operation, "
|
||||
"simple format). Negative values reduce volume, e.g. -12 means duck "
|
||||
"by 12dB. Converted to a linear ratio internally. Default: -12."
|
||||
),
|
||||
"default": -12,
|
||||
},
|
||||
"input_path": {"type": "string", "description": "Input for extract operation"},
|
||||
"output_path": {"type": "string"},
|
||||
"ducking": {
|
||||
"type": "object",
|
||||
"description": (
|
||||
"Advanced ducking parameters. Works with both the simple "
|
||||
"(primary_audio/secondary_audio) and advanced (tracks) formats."
|
||||
),
|
||||
"properties": {
|
||||
"enabled": {"type": "boolean", "default": True},
|
||||
"music_volume_during_speech": {
|
||||
"type": "number", "minimum": 0, "maximum": 1.0, "default": 0.15,
|
||||
},
|
||||
"attack_ms": {"type": "number", "default": 200},
|
||||
"release_ms": {"type": "number", "default": 500},
|
||||
},
|
||||
},
|
||||
"normalize": {"type": "boolean", "default": True},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(cpu_cores=2, ram_mb=1024, vram_mb=0, disk_mb=500)
|
||||
idempotency_key_fields = ["operation", "tracks", "ducking"]
|
||||
side_effects = ["writes mixed audio file to output_path"]
|
||||
user_visible_verification = [
|
||||
"Listen to mixed output and verify speech clarity and music ducking",
|
||||
]
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
operation = inputs["operation"]
|
||||
start = time.time()
|
||||
|
||||
try:
|
||||
if operation == "mix":
|
||||
result = self._mix(inputs)
|
||||
elif operation == "duck":
|
||||
result = self._duck(inputs)
|
||||
elif operation == "extract":
|
||||
result = self._extract(inputs)
|
||||
elif operation == "full_mix":
|
||||
result = self._full_mix(inputs)
|
||||
else:
|
||||
return ToolResult(success=False, error=f"Unknown operation: {operation}")
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=str(e))
|
||||
|
||||
result.duration_seconds = round(time.time() - start, 2)
|
||||
return result
|
||||
|
||||
def _mix(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
"""Mix multiple audio tracks into one output."""
|
||||
tracks = inputs.get("tracks", [])
|
||||
if not tracks:
|
||||
return ToolResult(success=False, error="No tracks provided")
|
||||
|
||||
output_path = Path(inputs.get("output_path", "mixed_audio.wav"))
|
||||
normalize = inputs.get("normalize", True)
|
||||
|
||||
# Validate all inputs exist
|
||||
for t in tracks:
|
||||
if not Path(t["path"]).exists():
|
||||
return ToolResult(success=False, error=f"Track not found: {t['path']}")
|
||||
|
||||
# Build FFmpeg complex filter for mixing
|
||||
filter_parts = []
|
||||
input_args = []
|
||||
|
||||
for i, track in enumerate(tracks):
|
||||
input_args.extend(["-i", track["path"]])
|
||||
volume = track.get("volume", 1.0)
|
||||
delay_ms = int(track.get("start_seconds", 0) * 1000)
|
||||
fade_in = track.get("fade_in_seconds", 0)
|
||||
fade_out = track.get("fade_out_seconds", 0)
|
||||
|
||||
filters = []
|
||||
if volume != 1.0:
|
||||
filters.append(f"volume={volume}")
|
||||
if delay_ms > 0:
|
||||
filters.append(f"adelay={delay_ms}|{delay_ms}")
|
||||
if fade_in > 0:
|
||||
filters.append(f"afade=t=in:d={fade_in}")
|
||||
if fade_out > 0:
|
||||
filters.append(f"afade=t=out:d={fade_out}")
|
||||
|
||||
if filters:
|
||||
filter_chain = ",".join(filters)
|
||||
filter_parts.append(f"[{i}:a]{filter_chain}[a{i}]")
|
||||
else:
|
||||
filter_parts.append(f"[{i}:a]acopy[a{i}]")
|
||||
|
||||
# Amix all processed streams
|
||||
mix_inputs = "".join(f"[a{i}]" for i in range(len(tracks)))
|
||||
filter_parts.append(
|
||||
f"{mix_inputs}amix=inputs={len(tracks)}:duration=longest:dropout_transition=2[mixed]"
|
||||
)
|
||||
|
||||
if normalize:
|
||||
filter_parts.append("[mixed]loudnorm=I=-16:LRA=11:TP=-1.5[out]")
|
||||
out_label = "[out]"
|
||||
else:
|
||||
out_label = "[mixed]"
|
||||
|
||||
filter_complex = ";".join(filter_parts)
|
||||
|
||||
cmd = ["ffmpeg", "-y"]
|
||||
cmd.extend(input_args)
|
||||
cmd.extend(["-filter_complex", filter_complex])
|
||||
cmd.extend(["-map", out_label, str(output_path)])
|
||||
|
||||
self.run_command(cmd)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"operation": "mix",
|
||||
"track_count": len(tracks),
|
||||
"output": str(output_path),
|
||||
"normalized": normalize,
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
)
|
||||
|
||||
def _duck(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
"""Apply ducking: lower music volume when speech is present.
|
||||
|
||||
Accepts two input formats:
|
||||
|
||||
Simple format (preferred for agents):
|
||||
{
|
||||
"operation": "duck",
|
||||
"primary_audio": "speech.mp3",
|
||||
"secondary_audio": "music.mp3",
|
||||
"duck_level": -12,
|
||||
"output_path": "out.wav"
|
||||
}
|
||||
|
||||
Advanced format (tracks array):
|
||||
{
|
||||
"operation": "duck",
|
||||
"tracks": [
|
||||
{"path": "speech.mp3", "role": "primary"}, # or "speech"
|
||||
{"path": "music.mp3", "role": "secondary"} # or "music"
|
||||
],
|
||||
"output_path": "out.wav"
|
||||
}
|
||||
"""
|
||||
ducking = inputs.get("ducking", {})
|
||||
output_path = Path(inputs.get("output_path", "ducked_audio.wav"))
|
||||
|
||||
# --- Resolve speech/music paths from either input format ---
|
||||
speech_path = None
|
||||
music_path = None
|
||||
|
||||
# Simple format: primary_audio / secondary_audio
|
||||
if "primary_audio" in inputs or "secondary_audio" in inputs:
|
||||
speech_path = inputs.get("primary_audio")
|
||||
music_path = inputs.get("secondary_audio")
|
||||
# If duck_level (dB) is provided, convert to linear ratio for
|
||||
# music_volume_during_speech. e.g. -12 dB -> 10^(-12/20) ~ 0.25
|
||||
if "duck_level" in inputs and "ducking" not in inputs:
|
||||
import math
|
||||
db = inputs["duck_level"]
|
||||
ducking = dict(ducking) # copy so we don't mutate caller
|
||||
ducking.setdefault(
|
||||
"music_volume_during_speech",
|
||||
round(math.pow(10, db / 20), 4),
|
||||
)
|
||||
|
||||
# Advanced format: tracks array with role field
|
||||
tracks = inputs.get("tracks", [])
|
||||
if tracks and speech_path is None and music_path is None:
|
||||
# Support both naming conventions: speech/music and primary/secondary
|
||||
speech_tracks = [
|
||||
t for t in tracks if t.get("role") in ("speech", "primary")
|
||||
]
|
||||
music_tracks = [
|
||||
t for t in tracks if t.get("role") in ("music", "secondary")
|
||||
]
|
||||
if speech_tracks:
|
||||
speech_path = speech_tracks[0]["path"]
|
||||
if music_tracks:
|
||||
music_path = music_tracks[0]["path"]
|
||||
|
||||
if not speech_path or not music_path:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=(
|
||||
"Ducking requires a primary (speech) and secondary (music) track. "
|
||||
"Provide either primary_audio/secondary_audio params, or a tracks "
|
||||
"array with role='speech'/'primary' and role='music'/'secondary'."
|
||||
),
|
||||
)
|
||||
|
||||
# Use FFmpeg sidechaincompress for ducking
|
||||
music_vol = ducking.get("music_volume_during_speech", 0.15)
|
||||
attack = ducking.get("attack_ms", 200) / 1000
|
||||
release = ducking.get("release_ms", 500) / 1000
|
||||
|
||||
# Sidechain compress: use speech as the key signal to duck music
|
||||
filter_complex = (
|
||||
f"[1:a]sidechaincompress="
|
||||
f"threshold=0.02:ratio=9:attack={attack}:release={release}:"
|
||||
f"level_sc=1:mix=0.9[ducked];"
|
||||
f"[ducked]volume={music_vol * 3}[music_out];" # compensate sidechain level
|
||||
f"[0:a][music_out]amix=inputs=2:duration=longest[out]"
|
||||
)
|
||||
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", speech_path,
|
||||
"-i", music_path,
|
||||
"-filter_complex", filter_complex,
|
||||
"-map", "[out]",
|
||||
str(output_path),
|
||||
]
|
||||
|
||||
self.run_command(cmd)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"operation": "duck",
|
||||
"speech_track": speech_path,
|
||||
"music_track": music_path,
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
)
|
||||
|
||||
def _extract(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
"""Extract audio from a video file."""
|
||||
input_path = Path(inputs["input_path"])
|
||||
if not input_path.exists():
|
||||
return ToolResult(success=False, error=f"Input not found: {input_path}")
|
||||
|
||||
output_path = Path(
|
||||
inputs.get("output_path", str(input_path.with_suffix(".wav")))
|
||||
)
|
||||
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", str(input_path),
|
||||
"-vn",
|
||||
"-acodec", "pcm_s16le",
|
||||
"-ar", "16000",
|
||||
"-ac", "1",
|
||||
str(output_path),
|
||||
]
|
||||
|
||||
self.run_command(cmd)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"operation": "extract",
|
||||
"input": str(input_path),
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
)
|
||||
|
||||
def _full_mix(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
"""One-call mix: layer narration tracks, add music with ducking, normalize.
|
||||
|
||||
This is the preferred operation for the compose-director skill.
|
||||
It combines mix + duck + normalize in a single FFmpeg filter graph.
|
||||
|
||||
Input format:
|
||||
{
|
||||
"operation": "full_mix",
|
||||
"tracks": [
|
||||
{"path": "narration_s1.mp3", "role": "speech", "start_seconds": 0},
|
||||
{"path": "narration_s2.mp3", "role": "speech", "start_seconds": 10.5},
|
||||
{"path": "music.mp3", "role": "music", "volume": 0.3}
|
||||
],
|
||||
"ducking": {
|
||||
"enabled": true,
|
||||
"music_volume_during_speech": 0.15,
|
||||
"attack_ms": 200,
|
||||
"release_ms": 500
|
||||
},
|
||||
"normalize": true,
|
||||
"output_path": "mixed_audio.wav"
|
||||
}
|
||||
"""
|
||||
tracks = inputs.get("tracks", [])
|
||||
if not tracks:
|
||||
return ToolResult(success=False, error="No tracks provided for full_mix")
|
||||
|
||||
output_path = Path(inputs.get("output_path", "full_mix_output.wav"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
normalize = inputs.get("normalize", True)
|
||||
ducking = inputs.get("ducking", {"enabled": True})
|
||||
|
||||
speech_tracks = [t for t in tracks if t.get("role") in ("speech", "primary")]
|
||||
music_tracks = [t for t in tracks if t.get("role") in ("music", "secondary")]
|
||||
sfx_tracks = [t for t in tracks if t.get("role") == "sfx"]
|
||||
all_tracks = speech_tracks + music_tracks + sfx_tracks
|
||||
|
||||
if not all_tracks:
|
||||
return ToolResult(success=False, error="No valid tracks (need speech/music/sfx roles)")
|
||||
|
||||
# Validate all files exist
|
||||
for t in all_tracks:
|
||||
if not Path(t["path"]).exists():
|
||||
return ToolResult(success=False, error=f"Track not found: {t['path']}")
|
||||
|
||||
# Build FFmpeg inputs and filter graph
|
||||
input_args = []
|
||||
filter_parts = []
|
||||
|
||||
for i, track in enumerate(all_tracks):
|
||||
input_args.extend(["-i", track["path"]])
|
||||
volume = track.get("volume", 1.0)
|
||||
delay_ms = int(track.get("start_seconds", 0) * 1000)
|
||||
fade_in = track.get("fade_in_seconds", 0)
|
||||
fade_out = track.get("fade_out_seconds", 0)
|
||||
|
||||
filters = []
|
||||
if volume != 1.0:
|
||||
filters.append(f"volume={volume}")
|
||||
if delay_ms > 0:
|
||||
filters.append(f"adelay={delay_ms}|{delay_ms}")
|
||||
if fade_in > 0:
|
||||
filters.append(f"afade=t=in:d={fade_in}")
|
||||
if fade_out > 0:
|
||||
filters.append(f"afade=t=out:d={fade_out}")
|
||||
|
||||
if filters:
|
||||
filter_chain = ",".join(filters)
|
||||
filter_parts.append(f"[{i}:a]{filter_chain}[a{i}]")
|
||||
else:
|
||||
filter_parts.append(f"[{i}:a]acopy[a{i}]")
|
||||
|
||||
# If ducking is enabled and we have both speech and music, apply sidechain
|
||||
duck_enabled = ducking.get("enabled", True) if isinstance(ducking, dict) else bool(ducking)
|
||||
|
||||
if duck_enabled and speech_tracks and music_tracks:
|
||||
# Mix speech tracks together first
|
||||
speech_indices = list(range(len(speech_tracks)))
|
||||
speech_labels = "".join(f"[a{i}]" for i in speech_indices)
|
||||
|
||||
if len(speech_tracks) > 1:
|
||||
filter_parts.append(
|
||||
f"{speech_labels}amix=inputs={len(speech_tracks)}:duration=longest[speech_mix]"
|
||||
)
|
||||
speech_out = "[speech_mix]"
|
||||
else:
|
||||
speech_out = f"[a{speech_indices[0]}]"
|
||||
|
||||
# Mix music tracks together
|
||||
music_start = len(speech_tracks)
|
||||
music_indices = list(range(music_start, music_start + len(music_tracks)))
|
||||
music_labels = "".join(f"[a{i}]" for i in music_indices)
|
||||
|
||||
if len(music_tracks) > 1:
|
||||
filter_parts.append(
|
||||
f"{music_labels}amix=inputs={len(music_tracks)}:duration=longest[music_mix]"
|
||||
)
|
||||
music_in = "[music_mix]"
|
||||
else:
|
||||
music_in = f"[a{music_indices[0]}]"
|
||||
|
||||
# Apply sidechain ducking
|
||||
duck_params = ducking if isinstance(ducking, dict) else {}
|
||||
attack = duck_params.get("attack_ms", 200) / 1000
|
||||
release = duck_params.get("release_ms", 500) / 1000
|
||||
music_vol = duck_params.get("music_volume_during_speech", 0.15)
|
||||
|
||||
filter_parts.append(
|
||||
f"{music_in}{speech_out}sidechaincompress="
|
||||
f"threshold=0.02:ratio=9:attack={attack}:release={release}:"
|
||||
f"level_sc=1:mix=0.9[ducked_music];"
|
||||
f"[ducked_music]volume={music_vol * 3}[music_out]"
|
||||
)
|
||||
|
||||
# Duplicate speech for final mix (sidechain consumes it as key)
|
||||
filter_parts.append(
|
||||
f"{speech_out}acopy[speech_dup]" if speech_out.startswith("[a") else ""
|
||||
)
|
||||
# Re-mix speech path: we need speech audio in the output too
|
||||
# Simpler approach: use amix on original speech and ducked music
|
||||
# Reset: use a cleaner approach — amerge the speech mix and ducked music
|
||||
# Actually, let's rebuild. The sidechain approach above uses speech as
|
||||
# the key signal but doesn't consume it from the output chain.
|
||||
# FFmpeg sidechaincompress: input 0 = audio to compress, input 1 = key signal
|
||||
# So music is compressed, speech signal is the key. We need to mix them.
|
||||
# Remove the last filter_part (the acopy that may be empty)
|
||||
if filter_parts and filter_parts[-1] == "":
|
||||
filter_parts.pop()
|
||||
|
||||
# Build speech mix for output separately
|
||||
if len(speech_tracks) > 1:
|
||||
# speech_mix already exists, make a copy for output
|
||||
filter_parts.append(f"{speech_labels}amix=inputs={len(speech_tracks)}:duration=longest[speech_out]")
|
||||
else:
|
||||
filter_parts.append(f"[a{speech_indices[0]}]acopy[speech_out]")
|
||||
|
||||
# Final mix: speech_out + music_out
|
||||
mix_label = "[speech_out][music_out]amix=inputs=2:duration=longest[premix]"
|
||||
|
||||
# Add SFX if present
|
||||
sfx_start = len(speech_tracks) + len(music_tracks)
|
||||
if sfx_tracks:
|
||||
sfx_labels = "".join(f"[a{i}]" for i in range(sfx_start, sfx_start + len(sfx_tracks)))
|
||||
filter_parts.append(mix_label.replace("[premix]", "[pressfx]"))
|
||||
filter_parts.append(
|
||||
f"[pressfx]{sfx_labels}amix=inputs={1 + len(sfx_tracks)}:duration=longest[premix]"
|
||||
)
|
||||
else:
|
||||
filter_parts.append(mix_label)
|
||||
|
||||
else:
|
||||
# No ducking: simple amix of all tracks
|
||||
all_labels = "".join(f"[a{i}]" for i in range(len(all_tracks)))
|
||||
filter_parts.append(
|
||||
f"{all_labels}amix=inputs={len(all_tracks)}:duration=longest:dropout_transition=2[premix]"
|
||||
)
|
||||
|
||||
# Normalize
|
||||
if normalize:
|
||||
filter_parts.append("[premix]loudnorm=I=-16:LRA=11:TP=-1.5[out]")
|
||||
out_label = "[out]"
|
||||
else:
|
||||
out_label = "[premix]"
|
||||
|
||||
filter_complex = ";".join(p for p in filter_parts if p)
|
||||
|
||||
cmd = ["ffmpeg", "-y"]
|
||||
cmd.extend(input_args)
|
||||
cmd.extend(["-filter_complex", filter_complex])
|
||||
cmd.extend(["-map", out_label, str(output_path)])
|
||||
|
||||
self.run_command(cmd)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"operation": "full_mix",
|
||||
"speech_tracks": len(speech_tracks),
|
||||
"music_tracks": len(music_tracks),
|
||||
"sfx_tracks": len(sfx_tracks),
|
||||
"ducking_enabled": duck_enabled,
|
||||
"normalized": normalize,
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
)
|
||||
@@ -0,0 +1,187 @@
|
||||
"""ElevenLabs text-to-speech provider tool."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class ElevenLabsTTS(BaseTool):
|
||||
name = "elevenlabs_tts"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.VOICE
|
||||
capability = "tts"
|
||||
provider = "elevenlabs"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.STOCHASTIC
|
||||
runtime = ToolRuntime.API
|
||||
|
||||
dependencies = []
|
||||
install_instructions = (
|
||||
"Set the ELEVENLABS_API_KEY environment variable:\n"
|
||||
" export ELEVENLABS_API_KEY=your_key_here\n"
|
||||
"Get a key at https://elevenlabs.io"
|
||||
)
|
||||
fallback = "openai_tts"
|
||||
fallback_tools = ["openai_tts", "piper_tts"]
|
||||
agent_skills = ["elevenlabs", "text-to-speech"]
|
||||
|
||||
capabilities = [
|
||||
"text_to_speech",
|
||||
"voice_selection",
|
||||
"ssml_support",
|
||||
"pronunciation_control",
|
||||
]
|
||||
supports = {
|
||||
"voice_cloning": True,
|
||||
"multilingual": True,
|
||||
"offline": False,
|
||||
"native_audio": True,
|
||||
}
|
||||
best_for = [
|
||||
"high-quality narration",
|
||||
"voice-sensitive spokesperson videos",
|
||||
"multilingual spoken delivery",
|
||||
]
|
||||
not_good_for = [
|
||||
"fully offline production",
|
||||
"privacy-constrained local-only workflows",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["text"],
|
||||
"properties": {
|
||||
"text": {"type": "string", "description": "Text to convert to speech"},
|
||||
"voice_id": {
|
||||
"type": "string",
|
||||
"description": "ElevenLabs voice ID (default: Rachel)",
|
||||
},
|
||||
"model_id": {
|
||||
"type": "string",
|
||||
"default": "eleven_multilingual_v2",
|
||||
"description": "TTS model to use",
|
||||
},
|
||||
"stability": {
|
||||
"type": "number",
|
||||
"default": 0.5,
|
||||
"minimum": 0,
|
||||
"maximum": 1,
|
||||
},
|
||||
"similarity_boost": {
|
||||
"type": "number",
|
||||
"default": 0.75,
|
||||
"minimum": 0,
|
||||
"maximum": 1,
|
||||
},
|
||||
"style": {
|
||||
"type": "number",
|
||||
"default": 0.0,
|
||||
"minimum": 0,
|
||||
"maximum": 1,
|
||||
},
|
||||
"output_path": {"type": "string"},
|
||||
"output_format": {
|
||||
"type": "string",
|
||||
"default": "mp3_44100_128",
|
||||
"enum": ["mp3_44100_128", "mp3_44100_192", "pcm_16000", "pcm_24000"],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=1, ram_mb=256, vram_mb=0, disk_mb=50, network_required=True
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=2, retryable_errors=["rate_limit", "timeout"])
|
||||
idempotency_key_fields = ["text", "voice_id", "model_id"]
|
||||
side_effects = ["writes audio file to output_path", "calls ElevenLabs API"]
|
||||
user_visible_verification = ["Listen to generated audio for natural speech quality"]
|
||||
|
||||
DEFAULT_VOICE_ID = "21m00Tcm4TlvDq8ikWAM"
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if os.environ.get("ELEVENLABS_API_KEY"):
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
return round(len(inputs.get("text", "")) * 0.0003, 4)
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
api_key = os.environ.get("ELEVENLABS_API_KEY")
|
||||
if not api_key:
|
||||
return ToolResult(success=False, error="No ElevenLabs API key. " + self.install_instructions)
|
||||
|
||||
start = time.time()
|
||||
try:
|
||||
result = self._generate(inputs, api_key)
|
||||
except Exception as exc:
|
||||
return ToolResult(success=False, error=f"TTS generation failed: {exc}")
|
||||
|
||||
result.duration_seconds = round(time.time() - start, 2)
|
||||
result.cost_usd = self.estimate_cost(inputs)
|
||||
return result
|
||||
|
||||
def _generate(self, inputs: dict[str, Any], api_key: str) -> ToolResult:
|
||||
import requests
|
||||
|
||||
text = inputs["text"]
|
||||
voice_id = inputs.get("voice_id", self.DEFAULT_VOICE_ID)
|
||||
model_id = inputs.get("model_id", "eleven_multilingual_v2")
|
||||
output_format = inputs.get("output_format", "mp3_44100_128")
|
||||
|
||||
response = requests.post(
|
||||
f"https://api.elevenlabs.io/v1/text-to-speech/{voice_id}",
|
||||
headers={
|
||||
"xi-api-key": api_key,
|
||||
"Content-Type": "application/json",
|
||||
"Accept": "audio/mpeg",
|
||||
},
|
||||
json={
|
||||
"text": text,
|
||||
"model_id": model_id,
|
||||
"voice_settings": {
|
||||
"stability": inputs.get("stability", 0.5),
|
||||
"similarity_boost": inputs.get("similarity_boost", 0.75),
|
||||
"style": inputs.get("style", 0.0),
|
||||
},
|
||||
},
|
||||
params={"output_format": output_format},
|
||||
timeout=120,
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
ext = "mp3" if "mp3" in output_format else "wav"
|
||||
output_path = Path(inputs.get("output_path", f"tts_output.{ext}"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(response.content)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": self.provider,
|
||||
"model": model_id,
|
||||
"voice_id": voice_id,
|
||||
"text_length": len(text),
|
||||
"output": str(output_path),
|
||||
"format": output_format,
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
model=model_id,
|
||||
)
|
||||
@@ -0,0 +1,149 @@
|
||||
"""Music generation tool via ElevenLabs Music API.
|
||||
|
||||
Generates background music and sound effects for video production.
|
||||
Reports unavailable when no API key is configured.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class MusicGen(BaseTool):
|
||||
name = "music_gen"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "music_generation"
|
||||
provider = "elevenlabs"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.STOCHASTIC
|
||||
runtime = ToolRuntime.API
|
||||
|
||||
dependencies = [] # checked dynamically via API key
|
||||
install_instructions = (
|
||||
"Set the ELEVENLABS_API_KEY environment variable:\n"
|
||||
" export ELEVENLABS_API_KEY=your_key_here\n"
|
||||
"Get a key at https://elevenlabs.io"
|
||||
)
|
||||
|
||||
agent_skills = ["music", "sound-effects", "elevenlabs"]
|
||||
|
||||
capabilities = [
|
||||
"generate_background_music",
|
||||
"generate_sfx",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {
|
||||
"type": "string",
|
||||
"description": "Music description (mood, genre, instruments, tempo)",
|
||||
},
|
||||
"duration_seconds": {
|
||||
"type": "number",
|
||||
"default": 60,
|
||||
"minimum": 3,
|
||||
"maximum": 600,
|
||||
"description": "Target duration in seconds (API supports 3-600s)",
|
||||
},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=1, ram_mb=256, vram_mb=0, disk_mb=50, network_required=True
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=2, retryable_errors=["rate_limit", "timeout"])
|
||||
idempotency_key_fields = ["prompt", "duration_seconds"]
|
||||
side_effects = ["writes audio file to output_path", "calls ElevenLabs API"]
|
||||
user_visible_verification = [
|
||||
"Listen to generated music for mood and quality",
|
||||
]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if os.environ.get("ELEVENLABS_API_KEY"):
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
# ElevenLabs music generation pricing is per generation
|
||||
duration = inputs.get("duration_seconds", 60)
|
||||
# Approximate: ~$0.05 per 30 seconds
|
||||
return round(duration / 30 * 0.05, 4)
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
api_key = os.environ.get("ELEVENLABS_API_KEY")
|
||||
if not api_key:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="No ElevenLabs API key. " + self.install_instructions,
|
||||
)
|
||||
|
||||
start = time.time()
|
||||
|
||||
try:
|
||||
result = self._generate(inputs, api_key)
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Music generation failed: {e}")
|
||||
|
||||
result.duration_seconds = round(time.time() - start, 2)
|
||||
result.cost_usd = self.estimate_cost(inputs)
|
||||
return result
|
||||
|
||||
def _generate(self, inputs: dict[str, Any], api_key: str) -> ToolResult:
|
||||
import requests
|
||||
|
||||
prompt = inputs["prompt"]
|
||||
duration = inputs.get("duration_seconds", 60)
|
||||
|
||||
url = "https://api.elevenlabs.io/v1/music"
|
||||
|
||||
headers = {
|
||||
"xi-api-key": api_key,
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
payload = {
|
||||
"prompt": prompt,
|
||||
"music_length_ms": int(duration * 1000),
|
||||
}
|
||||
|
||||
response = requests.post(
|
||||
url, headers=headers, json=payload, timeout=180
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
output_path = Path(inputs.get("output_path", "music_output.mp3"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(response.content)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "elevenlabs",
|
||||
"prompt": prompt,
|
||||
"duration_seconds": duration,
|
||||
"output": str(output_path),
|
||||
"format": "mp3",
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
)
|
||||
@@ -0,0 +1,154 @@
|
||||
"""OpenAI text-to-speech provider tool."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class OpenAITTS(BaseTool):
|
||||
name = "openai_tts"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.VOICE
|
||||
capability = "tts"
|
||||
provider = "openai"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.STOCHASTIC
|
||||
runtime = ToolRuntime.API
|
||||
|
||||
dependencies = []
|
||||
install_instructions = (
|
||||
"Set the OPENAI_API_KEY environment variable:\n"
|
||||
" export OPENAI_API_KEY=your_key_here\n"
|
||||
"Get a key at https://platform.openai.com/"
|
||||
)
|
||||
fallback = "piper_tts"
|
||||
fallback_tools = ["piper_tts"]
|
||||
agent_skills = ["openai-docs"]
|
||||
|
||||
capabilities = [
|
||||
"text_to_speech",
|
||||
"voice_selection",
|
||||
]
|
||||
supports = {
|
||||
"voice_cloning": False,
|
||||
"multilingual": True,
|
||||
"offline": False,
|
||||
"native_audio": True,
|
||||
}
|
||||
best_for = [
|
||||
"general narration fallback",
|
||||
"API-based production when ElevenLabs is unavailable",
|
||||
]
|
||||
not_good_for = [
|
||||
"voice clone matching",
|
||||
"fully offline production",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["text"],
|
||||
"properties": {
|
||||
"text": {"type": "string"},
|
||||
"voice": {
|
||||
"type": "string",
|
||||
"default": "alloy",
|
||||
"description": "OpenAI voice name",
|
||||
},
|
||||
"model": {
|
||||
"type": "string",
|
||||
"default": "gpt-4o-mini-tts",
|
||||
"description": "OpenAI speech model",
|
||||
},
|
||||
"format": {
|
||||
"type": "string",
|
||||
"default": "mp3",
|
||||
"enum": ["mp3", "wav", "pcm"],
|
||||
},
|
||||
"instructions": {
|
||||
"type": "string",
|
||||
"description": "Optional delivery instructions for the voice",
|
||||
},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=1, ram_mb=256, vram_mb=0, disk_mb=50, network_required=True
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=2, retryable_errors=["rate_limit", "timeout"])
|
||||
idempotency_key_fields = ["text", "voice", "model", "format"]
|
||||
side_effects = ["writes audio file to output_path", "calls OpenAI API"]
|
||||
user_visible_verification = ["Listen to generated audio for intelligibility and tone"]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if os.environ.get("OPENAI_API_KEY"):
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
return round(len(inputs.get("text", "")) * 0.000015, 4)
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
if not os.environ.get("OPENAI_API_KEY"):
|
||||
return ToolResult(success=False, error="No OpenAI API key. " + self.install_instructions)
|
||||
|
||||
start = time.time()
|
||||
try:
|
||||
result = self._generate(inputs)
|
||||
except Exception as exc:
|
||||
return ToolResult(success=False, error=f"OpenAI TTS failed: {exc}")
|
||||
|
||||
result.duration_seconds = round(time.time() - start, 2)
|
||||
result.cost_usd = self.estimate_cost(inputs)
|
||||
return result
|
||||
|
||||
def _generate(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI()
|
||||
text = inputs["text"]
|
||||
model = inputs.get("model", "gpt-4o-mini-tts")
|
||||
voice = inputs.get("voice", "alloy")
|
||||
fmt = inputs.get("format", "mp3")
|
||||
output_path = Path(inputs.get("output_path", f"openai_tts.{fmt}"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
with client.audio.speech.with_streaming_response.create(
|
||||
model=model,
|
||||
voice=voice,
|
||||
input=text,
|
||||
response_format=fmt,
|
||||
instructions=inputs.get("instructions"),
|
||||
) as response:
|
||||
response.stream_to_file(output_path)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": self.provider,
|
||||
"model": model,
|
||||
"voice": voice,
|
||||
"format": fmt,
|
||||
"text_length": len(text),
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
model=model,
|
||||
)
|
||||
@@ -0,0 +1,159 @@
|
||||
"""Piper local text-to-speech provider tool."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import shutil
|
||||
import subprocess
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class PiperTTS(BaseTool):
|
||||
name = "piper_tts"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.VOICE
|
||||
capability = "tts"
|
||||
provider = "piper"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
runtime = ToolRuntime.LOCAL
|
||||
|
||||
dependencies = ["cmd:piper"]
|
||||
install_instructions = (
|
||||
"Install Piper TTS:\n"
|
||||
" pip install piper-tts\n"
|
||||
"Or download from https://github.com/rhasspy/piper/releases\n"
|
||||
"Then download a voice model:\n"
|
||||
" piper --download-dir ~/.piper/models --model en_US-lessac-medium"
|
||||
)
|
||||
agent_skills = ["text-to-speech"]
|
||||
|
||||
capabilities = [
|
||||
"text_to_speech",
|
||||
"offline_generation",
|
||||
]
|
||||
supports = {
|
||||
"voice_cloning": False,
|
||||
"multilingual": False,
|
||||
"offline": True,
|
||||
"native_audio": True,
|
||||
}
|
||||
best_for = [
|
||||
"offline narration fallback",
|
||||
"privacy-sensitive local-only workflows",
|
||||
]
|
||||
not_good_for = [
|
||||
"best-in-class expressive voice quality",
|
||||
"voice clone matching",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["text"],
|
||||
"properties": {
|
||||
"text": {"type": "string"},
|
||||
"model": {
|
||||
"type": "string",
|
||||
"default": "en_US-lessac-medium",
|
||||
},
|
||||
"speaker_id": {
|
||||
"type": "integer",
|
||||
"default": 0,
|
||||
},
|
||||
"length_scale": {
|
||||
"type": "number",
|
||||
"default": 1.0,
|
||||
},
|
||||
"sentence_silence": {
|
||||
"type": "number",
|
||||
"default": 0.3,
|
||||
},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=2, ram_mb=512, vram_mb=0, disk_mb=200, network_required=False
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=1, retryable_errors=[])
|
||||
idempotency_key_fields = ["text", "model", "speaker_id", "length_scale"]
|
||||
side_effects = ["writes audio file to output_path"]
|
||||
user_visible_verification = ["Listen to generated audio for intelligibility"]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if shutil.which("piper"):
|
||||
return ToolStatus.AVAILABLE
|
||||
try:
|
||||
import piper # noqa: F401
|
||||
return ToolStatus.AVAILABLE
|
||||
except ImportError:
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
return 0.0
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
if self.get_status() != ToolStatus.AVAILABLE:
|
||||
return ToolResult(success=False, error="Piper TTS not available. " + self.install_instructions)
|
||||
|
||||
start = time.time()
|
||||
try:
|
||||
result = self._generate(inputs)
|
||||
except Exception as exc:
|
||||
return ToolResult(success=False, error=f"Local TTS generation failed: {exc}")
|
||||
|
||||
result.duration_seconds = round(time.time() - start, 2)
|
||||
return result
|
||||
|
||||
def _generate(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
output_path = Path(inputs.get("output_path", "tts_output.wav"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
proc = subprocess.run(
|
||||
[
|
||||
"piper",
|
||||
"--model", inputs.get("model", "en_US-lessac-medium"),
|
||||
"--speaker", str(inputs.get("speaker_id", 0)),
|
||||
"--length-scale", str(inputs.get("length_scale", 1.0)),
|
||||
"--sentence-silence", str(inputs.get("sentence_silence", 0.3)),
|
||||
"--output_file", str(output_path),
|
||||
],
|
||||
input=inputs["text"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=300,
|
||||
)
|
||||
|
||||
if proc.returncode != 0:
|
||||
return ToolResult(success=False, error=f"Piper failed (exit {proc.returncode}): {proc.stderr}")
|
||||
if not output_path.exists():
|
||||
return ToolResult(success=False, error=f"Piper output file missing: {output_path}")
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": self.provider,
|
||||
"model": inputs.get("model", "en_US-lessac-medium"),
|
||||
"speaker_id": inputs.get("speaker_id", 0),
|
||||
"text_length": len(inputs["text"]),
|
||||
"output": str(output_path),
|
||||
"format": "wav",
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
model=inputs.get("model", "en_US-lessac-medium"),
|
||||
)
|
||||
@@ -0,0 +1,288 @@
|
||||
"""Suno AI music generation via sunoapi.org REST API.
|
||||
|
||||
Generates full songs, instrumentals, and background music. Async flow:
|
||||
submit a generation request, poll for completion, download the audio file.
|
||||
Each request produces 2 tracks; the tool returns the first by default.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class SunoMusic(BaseTool):
|
||||
name = "suno_music"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "music_generation"
|
||||
provider = "suno"
|
||||
stability = ToolStability.BETA
|
||||
execution_mode = ExecutionMode.ASYNC
|
||||
determinism = Determinism.STOCHASTIC
|
||||
runtime = ToolRuntime.API
|
||||
|
||||
dependencies = [] # checked dynamically via env var
|
||||
install_instructions = (
|
||||
"Set the SUNO_API_KEY environment variable:\n"
|
||||
" export SUNO_API_KEY=your_key_here\n"
|
||||
"Get a key at https://sunoapi.org/api-key"
|
||||
)
|
||||
|
||||
agent_skills = ["music"]
|
||||
|
||||
capabilities = [
|
||||
"generate_background_music",
|
||||
"generate_song",
|
||||
"generate_instrumental",
|
||||
]
|
||||
supports = {
|
||||
"instrumental": True,
|
||||
"vocals": True,
|
||||
"custom_lyrics": True,
|
||||
"style_control": True,
|
||||
"long_form": True,
|
||||
}
|
||||
best_for = [
|
||||
"full song generation with vocals and lyrics",
|
||||
"high-quality instrumental background music",
|
||||
"genre-specific music (any genre)",
|
||||
"longer tracks up to 8 minutes",
|
||||
]
|
||||
not_good_for = [
|
||||
"sound effects (use ElevenLabs SFX instead)",
|
||||
"sub-10-second stingers (minimum ~30s generation)",
|
||||
"offline generation",
|
||||
]
|
||||
|
||||
fallback_tools = ["music_gen"]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {
|
||||
"type": "string",
|
||||
"description": (
|
||||
"In simple mode: a description of desired music (max 500 chars). "
|
||||
"In custom mode: the exact lyrics to sing (max 3000 chars)."
|
||||
),
|
||||
},
|
||||
"style": {
|
||||
"type": "string",
|
||||
"description": "Genre/style description, e.g. 'upbeat electronic pop'. Used in custom mode only (max 200 chars).",
|
||||
},
|
||||
"title": {
|
||||
"type": "string",
|
||||
"description": "Song title. Used in custom mode only (max 80 chars).",
|
||||
},
|
||||
"instrumental": {
|
||||
"type": "boolean",
|
||||
"default": True,
|
||||
"description": "True for instrumental only (no vocals), false for vocals.",
|
||||
},
|
||||
"custom_mode": {
|
||||
"type": "boolean",
|
||||
"default": False,
|
||||
"description": "False = simple mode (prompt is a description, lyrics auto-generated). True = custom mode (prompt is exact lyrics, style/title required).",
|
||||
},
|
||||
"model": {
|
||||
"type": "string",
|
||||
"enum": ["V4", "V4_5", "V5"],
|
||||
"default": "V4",
|
||||
"description": "Suno model version. V4 = 4min max, V4_5/V5 = 8min max.",
|
||||
},
|
||||
"output_path": {"type": "string"},
|
||||
"track_index": {
|
||||
"type": "integer",
|
||||
"default": 0,
|
||||
"enum": [0, 1],
|
||||
"description": "Which of the 2 generated tracks to return (0 or 1).",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=1, ram_mb=256, vram_mb=0, disk_mb=100, network_required=True
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=2, retryable_errors=["rate_limit", "timeout"])
|
||||
idempotency_key_fields = ["prompt", "style", "instrumental", "model"]
|
||||
side_effects = ["writes audio file to output_path", "calls Suno API via sunoapi.org"]
|
||||
user_visible_verification = [
|
||||
"Listen to generated music for mood, genre accuracy, and quality",
|
||||
]
|
||||
|
||||
_BASE_URL = "https://api.sunoapi.org/api/v1"
|
||||
_POLL_INTERVAL = 30 # seconds between status checks
|
||||
_MAX_WAIT = 300 # 5 minutes max wait
|
||||
|
||||
def _get_api_key(self) -> str | None:
|
||||
return os.environ.get("SUNO_API_KEY")
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if self._get_api_key():
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
# Suno credits cost $0.005 each; a generation is roughly 10 credits
|
||||
return 0.05
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
api_key = self._get_api_key()
|
||||
if not api_key:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="No Suno API key. " + self.install_instructions,
|
||||
)
|
||||
|
||||
start = time.time()
|
||||
|
||||
try:
|
||||
# Step 1: Submit generation request
|
||||
task_id = self._submit(inputs, api_key)
|
||||
|
||||
# Step 2: Poll for completion
|
||||
result_data = self._poll(task_id, api_key)
|
||||
|
||||
# Step 3: Download audio
|
||||
track_index = inputs.get("track_index", 0)
|
||||
tracks = result_data.get("data", [])
|
||||
if not tracks:
|
||||
return ToolResult(success=False, error="Suno returned no tracks.")
|
||||
|
||||
track = tracks[min(track_index, len(tracks) - 1)]
|
||||
audio_url = track.get("audio_url")
|
||||
if not audio_url:
|
||||
return ToolResult(success=False, error="No audio_url in Suno response.")
|
||||
|
||||
output_path = self._download(audio_url, inputs, api_key)
|
||||
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Suno generation failed: {e}")
|
||||
|
||||
duration = round(time.time() - start, 2)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "suno",
|
||||
"model": inputs.get("model", "V4"),
|
||||
"prompt": inputs["prompt"],
|
||||
"style": inputs.get("style"),
|
||||
"title": track.get("title", inputs.get("title")),
|
||||
"instrumental": inputs.get("instrumental", True),
|
||||
"duration_seconds": track.get("duration"),
|
||||
"output": str(output_path),
|
||||
"format": "mp3",
|
||||
"track_id": track.get("id"),
|
||||
"tracks_generated": len(tracks),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
cost_usd=self.estimate_cost(inputs),
|
||||
duration_seconds=duration,
|
||||
model=f"suno/{inputs.get('model', 'V4')}",
|
||||
)
|
||||
|
||||
def _submit(self, inputs: dict[str, Any], api_key: str) -> str:
|
||||
"""Submit a generation request and return the taskId."""
|
||||
import requests
|
||||
|
||||
custom_mode = inputs.get("custom_mode", False)
|
||||
instrumental = inputs.get("instrumental", True)
|
||||
model = inputs.get("model", "V4")
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"model": model,
|
||||
"customMode": custom_mode,
|
||||
"instrumental": instrumental,
|
||||
"callBackUrl": "", # no webhook; we poll
|
||||
}
|
||||
|
||||
if custom_mode:
|
||||
payload["prompt"] = inputs["prompt"] # exact lyrics
|
||||
payload["style"] = inputs.get("style", "")
|
||||
payload["title"] = inputs.get("title", "")
|
||||
else:
|
||||
payload["prompt"] = inputs["prompt"][:500] # description, max 500 chars
|
||||
|
||||
response = requests.post(
|
||||
f"{self._BASE_URL}/generate",
|
||||
headers={
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
json=payload,
|
||||
timeout=30,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
task_id = data.get("data", {}).get("taskId") or data.get("taskId")
|
||||
if not task_id:
|
||||
raise RuntimeError(f"No taskId in Suno response: {data}")
|
||||
|
||||
return task_id
|
||||
|
||||
def _poll(self, task_id: str, api_key: str) -> dict:
|
||||
"""Poll for task completion and return the result data."""
|
||||
import requests
|
||||
|
||||
elapsed = 0
|
||||
while elapsed < self._MAX_WAIT:
|
||||
time.sleep(self._POLL_INTERVAL)
|
||||
elapsed += self._POLL_INTERVAL
|
||||
|
||||
response = requests.get(
|
||||
f"{self._BASE_URL}/generate/record-info",
|
||||
params={"taskId": task_id},
|
||||
headers={"Authorization": f"Bearer {api_key}"},
|
||||
timeout=30,
|
||||
)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
status = result.get("data", {}).get("status") or result.get("status", "")
|
||||
|
||||
if status == "SUCCESS":
|
||||
return result.get("data", result)
|
||||
elif status in (
|
||||
"CREATE_TASK_FAILED",
|
||||
"GENERATE_AUDIO_FAILED",
|
||||
"SENSITIVE_WORD_ERROR",
|
||||
):
|
||||
raise RuntimeError(f"Suno generation failed with status: {status}")
|
||||
|
||||
# PENDING, GENERATING, TEXT_SUCCESS, FIRST_SUCCESS — keep polling
|
||||
|
||||
raise TimeoutError(
|
||||
f"Suno generation timed out after {self._MAX_WAIT}s (taskId: {task_id})"
|
||||
)
|
||||
|
||||
def _download(self, audio_url: str, inputs: dict[str, Any], api_key: str) -> Path:
|
||||
"""Download the audio file to the output path."""
|
||||
import requests
|
||||
|
||||
output_path = Path(inputs.get("output_path", "suno_output.mp3"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
response = requests.get(audio_url, timeout=120)
|
||||
response.raise_for_status()
|
||||
output_path.write_bytes(response.content)
|
||||
|
||||
return output_path
|
||||
@@ -0,0 +1,136 @@
|
||||
"""Capability-level text-to-speech selector that chooses among provider tools.
|
||||
|
||||
Provider discovery is automatic — any BaseTool with capability="tts"
|
||||
is picked up from the registry. Adding a new TTS provider requires only creating
|
||||
the tool file in tools/audio/; no changes to this selector are needed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import BaseTool, ToolResult, ToolRuntime, ToolStability, ToolTier, ToolStatus
|
||||
|
||||
|
||||
class TTSSelector(BaseTool):
|
||||
name = "tts_selector"
|
||||
version = "0.2.0"
|
||||
tier = ToolTier.VOICE
|
||||
capability = "tts"
|
||||
provider = "selector"
|
||||
stability = ToolStability.BETA
|
||||
runtime = ToolRuntime.HYBRID
|
||||
agent_skills = ["text-to-speech", "elevenlabs", "openai-docs"]
|
||||
|
||||
capabilities = [
|
||||
"text_to_speech",
|
||||
"provider_selection",
|
||||
]
|
||||
supports = {
|
||||
"user_preference_routing": True,
|
||||
"offline_fallback": True,
|
||||
"multilingual": True,
|
||||
}
|
||||
best_for = [
|
||||
"preflight tool selection",
|
||||
"user-facing recommendation flows",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["text"],
|
||||
"properties": {
|
||||
"text": {"type": "string"},
|
||||
"voice_id": {
|
||||
"type": "string",
|
||||
"description": "Provider-specific voice ID. Passed through to the selected TTS provider.",
|
||||
},
|
||||
"model_id": {
|
||||
"type": "string",
|
||||
"description": "TTS model to use (e.g. eleven_multilingual_v2). Passed through to provider.",
|
||||
},
|
||||
"stability": {
|
||||
"type": "number", "minimum": 0, "maximum": 1,
|
||||
"description": "Voice stability (ElevenLabs). Lower = more expressive.",
|
||||
},
|
||||
"similarity_boost": {
|
||||
"type": "number", "minimum": 0, "maximum": 1,
|
||||
"description": "Voice similarity boost (ElevenLabs).",
|
||||
},
|
||||
"style": {
|
||||
"type": "number", "minimum": 0, "maximum": 1,
|
||||
"description": "Style exaggeration (ElevenLabs). Higher = more expressive.",
|
||||
},
|
||||
"output_format": {
|
||||
"type": "string",
|
||||
"description": "Audio output format (e.g. mp3_44100_128). Passed through to provider.",
|
||||
},
|
||||
"preferred_provider": {
|
||||
"type": "string",
|
||||
"description": "Provider name or 'auto'. Valid values are discovered at runtime from the registry.",
|
||||
"default": "auto",
|
||||
},
|
||||
"allowed_providers": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
def _providers(self) -> list[BaseTool]:
|
||||
"""Auto-discover TTS providers from the registry."""
|
||||
from tools.tool_registry import registry
|
||||
registry.ensure_discovered()
|
||||
return [t for t in registry.get_by_capability("tts")
|
||||
if t.name != self.name]
|
||||
|
||||
@property
|
||||
def fallback_tools(self) -> list[str]:
|
||||
"""Dynamically built from discovered providers."""
|
||||
return [t.name for t in self._providers()]
|
||||
|
||||
@property
|
||||
def provider_matrix(self) -> dict[str, dict[str, str]]:
|
||||
"""Built at runtime from each provider's best_for field."""
|
||||
matrix = {}
|
||||
for tool in self._providers():
|
||||
strength = ", ".join(tool.best_for) if tool.best_for else tool.name
|
||||
matrix[tool.provider] = {"tool": tool.name, "strength": strength}
|
||||
return matrix
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if any(tool.get_status() == ToolStatus.AVAILABLE for tool in self._providers()):
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
tool = self._select_tool(inputs)
|
||||
return tool.estimate_cost(inputs) if tool else 0.0
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
tool = self._select_tool(inputs)
|
||||
if tool is None:
|
||||
return ToolResult(success=False, error="No TTS provider available.")
|
||||
result = tool.execute(inputs)
|
||||
if result.success:
|
||||
result.data.setdefault("selected_tool", tool.name)
|
||||
return result
|
||||
|
||||
def _select_tool(self, inputs: dict[str, Any]) -> BaseTool | None:
|
||||
preferred = inputs.get("preferred_provider", "auto")
|
||||
allowed = set(inputs.get("allowed_providers") or [])
|
||||
candidates = self._providers()
|
||||
if allowed:
|
||||
candidates = [tool for tool in candidates if tool.provider in allowed]
|
||||
|
||||
if preferred != "auto":
|
||||
ordered = [tool for tool in candidates if tool.provider == preferred]
|
||||
ordered.extend([tool for tool in candidates if tool.provider != preferred])
|
||||
else:
|
||||
ordered = candidates
|
||||
|
||||
for tool in ordered:
|
||||
if tool.get_status() == ToolStatus.AVAILABLE:
|
||||
return tool
|
||||
return None
|
||||
@@ -0,0 +1 @@
|
||||
"""Avatar tools for talking head and lip sync generation."""
|
||||
@@ -0,0 +1,231 @@
|
||||
"""Lip-sync tool for dubbing and localization.
|
||||
|
||||
Syncs lip movements in a video to match a different audio track using
|
||||
Wav2Lip or MuseTalk models. Primary use case: replace original speech
|
||||
with translated audio and make the speaker's lips match.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
# Model checkpoint filenames by variant
|
||||
MODEL_CHECKPOINTS = {
|
||||
"wav2lip": "wav2lip.pth",
|
||||
"wav2lip_gan": "wav2lip_gan.pth",
|
||||
}
|
||||
|
||||
|
||||
class LipSync(BaseTool):
|
||||
name = "lip_sync"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "avatar"
|
||||
provider = "wav2lip"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.STOCHASTIC
|
||||
runtime = ToolRuntime.LOCAL_GPU
|
||||
|
||||
dependencies = ["python:torch", "cmd:ffmpeg"]
|
||||
install_instructions = (
|
||||
"Option 1: pip install wav2lip (if available)\n"
|
||||
"Option 2: Clone https://github.com/Rudrabha/Wav2Lip and set WAV2LIP_PATH env var\n"
|
||||
"Requires: PyTorch with CUDA, ffmpeg"
|
||||
)
|
||||
|
||||
agent_skills = ["ffmpeg"]
|
||||
|
||||
capabilities = [
|
||||
"lip_sync",
|
||||
"audio_video_alignment",
|
||||
"dubbing_support",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["video_path", "audio_path"],
|
||||
"properties": {
|
||||
"video_path": {
|
||||
"type": "string",
|
||||
"description": "Path to source video with face",
|
||||
},
|
||||
"audio_path": {
|
||||
"type": "string",
|
||||
"description": "Path to audio track to sync lips to",
|
||||
},
|
||||
"output_path": {
|
||||
"type": "string",
|
||||
"description": "Output video path (defaults to {stem}_lipsync.mp4)",
|
||||
},
|
||||
"model": {
|
||||
"type": "string",
|
||||
"enum": ["wav2lip", "wav2lip_gan"],
|
||||
"default": "wav2lip",
|
||||
"description": "Model variant (gan = higher quality but slower)",
|
||||
},
|
||||
"face_padding": {
|
||||
"type": "array",
|
||||
"items": {"type": "integer"},
|
||||
"minItems": 4,
|
||||
"maxItems": 4,
|
||||
"default": [0, 10, 0, 0],
|
||||
"description": "Padding around face crop: [top, bottom, left, right]",
|
||||
},
|
||||
"resize_factor": {
|
||||
"type": "integer",
|
||||
"default": 1,
|
||||
"description": "Downscale factor for faster processing",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=2, ram_mb=4096, vram_mb=4096, disk_mb=2000
|
||||
)
|
||||
idempotency_key_fields = ["video_path", "audio_path", "model", "face_padding", "resize_factor"]
|
||||
side_effects = ["writes lip-synced video to output_path"]
|
||||
user_visible_verification = [
|
||||
"Watch output video to verify lip movements match the new audio",
|
||||
"Check face region for visual artifacts or jitter",
|
||||
]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
"""Check Wav2Lip availability via env var or Python import."""
|
||||
# Check WAV2LIP_PATH environment variable
|
||||
wav2lip_path = os.environ.get("WAV2LIP_PATH")
|
||||
if wav2lip_path and Path(wav2lip_path).is_dir():
|
||||
return ToolStatus.AVAILABLE
|
||||
|
||||
# Fallback: try importing wav2lip as a Python package
|
||||
try:
|
||||
import wav2lip # noqa: F401
|
||||
return ToolStatus.AVAILABLE
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
return 0.0 # local GPU, free
|
||||
|
||||
def _resolve_wav2lip_dir(self) -> Path | None:
|
||||
"""Locate the Wav2Lip installation directory."""
|
||||
wav2lip_path = os.environ.get("WAV2LIP_PATH")
|
||||
if wav2lip_path:
|
||||
p = Path(wav2lip_path)
|
||||
if p.is_dir():
|
||||
return p
|
||||
|
||||
# Fallback: check if wav2lip is importable and find its location
|
||||
try:
|
||||
import wav2lip
|
||||
return Path(wav2lip.__file__).parent
|
||||
except (ImportError, AttributeError):
|
||||
pass
|
||||
|
||||
return None
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
if self.get_status() != ToolStatus.AVAILABLE:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="Wav2Lip not available. " + self.install_instructions,
|
||||
)
|
||||
|
||||
video_path = Path(inputs["video_path"])
|
||||
audio_path = Path(inputs["audio_path"])
|
||||
|
||||
if not video_path.exists():
|
||||
return ToolResult(success=False, error=f"Video not found: {video_path}")
|
||||
if not audio_path.exists():
|
||||
return ToolResult(success=False, error=f"Audio not found: {audio_path}")
|
||||
|
||||
output_path = Path(
|
||||
inputs.get("output_path", str(video_path.with_stem(f"{video_path.stem}_lipsync")))
|
||||
)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
model_variant = inputs.get("model", "wav2lip")
|
||||
face_padding = inputs.get("face_padding", [0, 10, 0, 0])
|
||||
resize_factor = inputs.get("resize_factor", 1)
|
||||
|
||||
wav2lip_dir = self._resolve_wav2lip_dir()
|
||||
if wav2lip_dir is None:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="Could not locate Wav2Lip directory. " + self.install_instructions,
|
||||
)
|
||||
|
||||
checkpoint = wav2lip_dir / "checkpoints" / MODEL_CHECKPOINTS[model_variant]
|
||||
if not checkpoint.exists():
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"Model checkpoint not found: {checkpoint}",
|
||||
)
|
||||
|
||||
inference_script = wav2lip_dir / "inference.py"
|
||||
if not inference_script.exists():
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"Inference script not found: {inference_script}",
|
||||
)
|
||||
|
||||
start = time.time()
|
||||
|
||||
cmd = [
|
||||
"python", str(inference_script),
|
||||
"--checkpoint_path", str(checkpoint),
|
||||
"--face", str(video_path),
|
||||
"--audio", str(audio_path),
|
||||
"--outfile", str(output_path),
|
||||
"--pads", *[str(p) for p in face_padding],
|
||||
"--resize_factor", str(resize_factor),
|
||||
]
|
||||
|
||||
try:
|
||||
self.run_command(cmd, timeout=600, cwd=wav2lip_dir)
|
||||
except subprocess.TimeoutExpired:
|
||||
return ToolResult(success=False, error="Lip-sync timed out after 600 seconds")
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Wav2Lip inference failed: {e}")
|
||||
|
||||
if not output_path.exists():
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"Inference completed but output file missing: {output_path}",
|
||||
)
|
||||
|
||||
elapsed = time.time() - start
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"video_input": str(video_path),
|
||||
"audio_input": str(audio_path),
|
||||
"output": str(output_path),
|
||||
"model": model_variant,
|
||||
"resize_factor": resize_factor,
|
||||
"face_padding": face_padding,
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
duration_seconds=round(elapsed, 2),
|
||||
model=model_variant,
|
||||
)
|
||||
@@ -0,0 +1,258 @@
|
||||
"""Photo-to-talking-head video generation tool.
|
||||
|
||||
Animates a still face photo to appear as if speaking provided audio.
|
||||
Uses SadTalker or MuseTalk models for audio-driven face animation.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import glob
|
||||
import os
|
||||
import shutil
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class TalkingHead(BaseTool):
|
||||
name = "talking_head"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "avatar"
|
||||
provider = "sadtalker"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.STOCHASTIC
|
||||
runtime = ToolRuntime.LOCAL_GPU
|
||||
|
||||
dependencies = [] # checked dynamically via get_status()
|
||||
install_instructions = (
|
||||
"Clone https://github.com/OpenTalker/SadTalker and set SADTALKER_PATH env var\n"
|
||||
"Requires: PyTorch with CUDA, ffmpeg\n"
|
||||
"pip install sadtalker # or clone the repo"
|
||||
)
|
||||
|
||||
agent_skills = ["ffmpeg"]
|
||||
fallback = "lip_sync"
|
||||
|
||||
capabilities = [
|
||||
"photo_to_video",
|
||||
"face_animation",
|
||||
"audio_driven_animation",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["image_path", "audio_path"],
|
||||
"properties": {
|
||||
"image_path": {
|
||||
"type": "string",
|
||||
"description": "Path to source face photo",
|
||||
},
|
||||
"audio_path": {
|
||||
"type": "string",
|
||||
"description": "Path to driving audio file",
|
||||
},
|
||||
"output_path": {
|
||||
"type": "string",
|
||||
"description": "Output video path (default: {stem}_talking.mp4)",
|
||||
},
|
||||
"model": {
|
||||
"type": "string",
|
||||
"enum": ["sadtalker", "musetalk"],
|
||||
"default": "sadtalker",
|
||||
"description": "Model to use for face animation",
|
||||
},
|
||||
"expression_scale": {
|
||||
"type": "number",
|
||||
"default": 1.0,
|
||||
"description": "Expression intensity multiplier",
|
||||
},
|
||||
"still_mode": {
|
||||
"type": "boolean",
|
||||
"default": False,
|
||||
"description": "Only animate mouth, keep head still",
|
||||
},
|
||||
"preprocess": {
|
||||
"type": "string",
|
||||
"enum": ["crop", "resize", "full"],
|
||||
"default": "crop",
|
||||
"description": "Face preprocessing mode",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=2, ram_mb=4096, vram_mb=4096, disk_mb=2000
|
||||
)
|
||||
idempotency_key_fields = ["image_path", "audio_path", "model", "expression_scale", "still_mode"]
|
||||
side_effects = ["writes video file to output_path"]
|
||||
user_visible_verification = [
|
||||
"Watch generated video for lip-sync accuracy",
|
||||
"Check for face distortion or unnatural artifacts",
|
||||
]
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Status
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
"""Check for SadTalker availability via env var or Python import."""
|
||||
# 1. SADTALKER_PATH env var pointing to cloned repo
|
||||
sadtalker_path = os.environ.get("SADTALKER_PATH", "")
|
||||
if sadtalker_path and Path(sadtalker_path).is_dir():
|
||||
return ToolStatus.AVAILABLE
|
||||
|
||||
# 2. Installed as a Python package
|
||||
try:
|
||||
import sadtalker # noqa: F401
|
||||
return ToolStatus.AVAILABLE
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Cost & runtime estimates
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
return 0.0 # local GPU, no API cost
|
||||
|
||||
def estimate_runtime(self, inputs: dict[str, Any]) -> float:
|
||||
# SadTalker typically takes 30-120s depending on audio length
|
||||
return 60.0
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Execution
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
image_path = Path(inputs["image_path"])
|
||||
audio_path = Path(inputs["audio_path"])
|
||||
|
||||
if not image_path.exists():
|
||||
return ToolResult(success=False, error=f"Image not found: {image_path}")
|
||||
if not audio_path.exists():
|
||||
return ToolResult(success=False, error=f"Audio not found: {audio_path}")
|
||||
|
||||
model = inputs.get("model", "sadtalker")
|
||||
output_path = Path(
|
||||
inputs.get("output_path", str(image_path.with_stem(f"{image_path.stem}_talking").with_suffix(".mp4")))
|
||||
)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
start = time.time()
|
||||
|
||||
try:
|
||||
if model == "sadtalker":
|
||||
result = self._run_sadtalker(inputs, image_path, audio_path, output_path)
|
||||
elif model == "musetalk":
|
||||
result = self._run_musetalk(inputs, image_path, audio_path, output_path)
|
||||
else:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"Unknown model: {model}. Supported: sadtalker, musetalk",
|
||||
)
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Talking head generation failed: {e}")
|
||||
|
||||
result.duration_seconds = round(time.time() - start, 2)
|
||||
return result
|
||||
|
||||
def _run_sadtalker(
|
||||
self,
|
||||
inputs: dict[str, Any],
|
||||
image_path: Path,
|
||||
audio_path: Path,
|
||||
output_path: Path,
|
||||
) -> ToolResult:
|
||||
"""Run SadTalker inference via subprocess."""
|
||||
sadtalker_path = os.environ.get("SADTALKER_PATH", "")
|
||||
if not sadtalker_path or not Path(sadtalker_path).is_dir():
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="SADTALKER_PATH not set or directory does not exist.\n" + self.install_instructions,
|
||||
)
|
||||
|
||||
sadtalker_dir = Path(sadtalker_path)
|
||||
result_dir = output_path.parent / "sadtalker_results"
|
||||
result_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
expression_scale = inputs.get("expression_scale", 1.0)
|
||||
still_mode = inputs.get("still_mode", False)
|
||||
preprocess = inputs.get("preprocess", "crop")
|
||||
|
||||
# Build SadTalker inference command
|
||||
cmd = [
|
||||
"python", str(sadtalker_dir / "inference.py"),
|
||||
"--driven_audio", str(audio_path),
|
||||
"--source_image", str(image_path),
|
||||
"--result_dir", str(result_dir),
|
||||
"--expression_scale", str(expression_scale),
|
||||
"--preprocess", preprocess,
|
||||
]
|
||||
|
||||
if still_mode:
|
||||
cmd.append("--still")
|
||||
|
||||
self.run_command(cmd, cwd=sadtalker_dir, timeout=600)
|
||||
|
||||
# Find the output video in result_dir (SadTalker names it automatically)
|
||||
generated = glob.glob(str(result_dir / "**" / "*.mp4"), recursive=True)
|
||||
if not generated:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"No output video found in {result_dir}",
|
||||
)
|
||||
|
||||
# Use the most recently created file
|
||||
generated.sort(key=os.path.getmtime, reverse=True)
|
||||
generated_path = Path(generated[0])
|
||||
|
||||
# Move to the desired output path
|
||||
shutil.move(str(generated_path), str(output_path))
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"model": "sadtalker",
|
||||
"image": str(image_path),
|
||||
"audio": str(audio_path),
|
||||
"output": str(output_path),
|
||||
"expression_scale": expression_scale,
|
||||
"still_mode": still_mode,
|
||||
"preprocess": preprocess,
|
||||
"format": "mp4",
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
model="sadtalker",
|
||||
)
|
||||
|
||||
def _run_musetalk(
|
||||
self,
|
||||
inputs: dict[str, Any],
|
||||
image_path: Path,
|
||||
audio_path: Path,
|
||||
output_path: Path,
|
||||
) -> ToolResult:
|
||||
"""MuseTalk support — placeholder for future implementation."""
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=(
|
||||
"MuseTalk support is not yet implemented. "
|
||||
"Use model='sadtalker' instead."
|
||||
),
|
||||
)
|
||||
@@ -0,0 +1,285 @@
|
||||
"""Base tool class implementing the expanded ToolContract.
|
||||
|
||||
Every tool in OpenMontage inherits from BaseTool. This enforces a uniform
|
||||
interface for discovery, execution, cost estimation, and health reporting.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import inspect
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import shutil
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Optional
|
||||
|
||||
|
||||
class ToolTier(str, Enum):
|
||||
CORE = "core"
|
||||
VOICE = "voice"
|
||||
ENHANCE = "enhance"
|
||||
GENERATE = "generate"
|
||||
SOURCE = "source"
|
||||
ANALYZE = "analyze"
|
||||
PUBLISH = "publish"
|
||||
|
||||
|
||||
class ToolStability(str, Enum):
|
||||
EXPERIMENTAL = "experimental"
|
||||
BETA = "beta"
|
||||
PRODUCTION = "production"
|
||||
|
||||
|
||||
class ToolStatus(str, Enum):
|
||||
AVAILABLE = "available"
|
||||
UNAVAILABLE = "unavailable"
|
||||
DEGRADED = "degraded"
|
||||
|
||||
|
||||
class ToolRuntime(str, Enum):
|
||||
"""Where and how a tool executes."""
|
||||
LOCAL = "local" # Runs entirely on-device, free, no network
|
||||
LOCAL_GPU = "local_gpu" # Runs on-device but needs GPU (VRAM)
|
||||
API = "api" # Calls an external API, requires API key, costs money
|
||||
HYBRID = "hybrid" # Can run locally OR via API (e.g., image_selector)
|
||||
|
||||
|
||||
class ExecutionMode(str, Enum):
|
||||
SYNC = "sync"
|
||||
ASYNC = "async"
|
||||
|
||||
|
||||
class Determinism(str, Enum):
|
||||
DETERMINISTIC = "deterministic"
|
||||
SEEDED = "seeded"
|
||||
STOCHASTIC = "stochastic"
|
||||
|
||||
|
||||
class ResumeSupport(str, Enum):
|
||||
NONE = "none"
|
||||
FROM_START = "from_start"
|
||||
FROM_CHECKPOINT = "from_checkpoint"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResourceProfile:
|
||||
"""Hardware resource envelope for a tool."""
|
||||
cpu_cores: int = 1
|
||||
ram_mb: int = 512
|
||||
vram_mb: int = 0
|
||||
disk_mb: int = 100
|
||||
network_required: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetryPolicy:
|
||||
"""Safe retry behavior for a tool."""
|
||||
max_retries: int = 0
|
||||
backoff_seconds: float = 1.0
|
||||
retryable_errors: list[str] = field(default_factory=list)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolResult:
|
||||
"""Standard result returned by tool execution."""
|
||||
success: bool
|
||||
data: dict[str, Any] = field(default_factory=dict)
|
||||
artifacts: list[str] = field(default_factory=list)
|
||||
error: Optional[str] = None
|
||||
cost_usd: float = 0.0
|
||||
duration_seconds: float = 0.0
|
||||
seed: Optional[int] = None
|
||||
model: Optional[str] = None
|
||||
|
||||
|
||||
class BaseTool(ABC):
|
||||
"""Abstract base class for all OpenMontage tools."""
|
||||
|
||||
# --- Identity (override in subclasses) ---
|
||||
name: str = ""
|
||||
version: str = "0.1.0"
|
||||
tier: ToolTier = ToolTier.CORE
|
||||
stability: ToolStability = ToolStability.EXPERIMENTAL
|
||||
execution_mode: ExecutionMode = ExecutionMode.SYNC
|
||||
determinism: Determinism = Determinism.DETERMINISTIC
|
||||
runtime: ToolRuntime = ToolRuntime.LOCAL
|
||||
|
||||
# --- Dependencies ---
|
||||
# For API tools, add "env:ENVVAR_NAME" to signal required API keys
|
||||
dependencies: list[str] = []
|
||||
install_instructions: str = ""
|
||||
|
||||
# --- Capabilities ---
|
||||
capability: str = "generic"
|
||||
provider: str = "openmontage"
|
||||
capabilities: list[str] = []
|
||||
input_schema: dict = {}
|
||||
output_schema: dict = {}
|
||||
artifact_schema: dict = {}
|
||||
progress_schema: Optional[dict] = None
|
||||
supports: dict[str, Any] = {}
|
||||
best_for: list[str] = []
|
||||
not_good_for: list[str] = []
|
||||
provider_matrix: dict[str, Any] = {}
|
||||
|
||||
# --- Resource & retry ---
|
||||
resource_profile: ResourceProfile = ResourceProfile()
|
||||
retry_policy: RetryPolicy = RetryPolicy()
|
||||
|
||||
# --- Resume & idempotency ---
|
||||
resume_support: ResumeSupport = ResumeSupport.NONE
|
||||
idempotency_key_fields: list[str] = []
|
||||
|
||||
# --- Side effects & fallback ---
|
||||
side_effects: list[str] = []
|
||||
fallback: Optional[str] = None
|
||||
fallback_tools: list[str] = []
|
||||
|
||||
# --- Agent skills (Layer 3 references) ---
|
||||
# Names of installed agent skills in .agents/skills/ that teach the
|
||||
# underlying technology. The orchestrator uses these to load relevant
|
||||
# API knowledge when planning tool usage.
|
||||
agent_skills: list[str] = []
|
||||
|
||||
# --- Verification ---
|
||||
user_visible_verification: list[str] = []
|
||||
|
||||
# ---- Status reporting ----
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
"""Check if this tool's dependencies are satisfied."""
|
||||
try:
|
||||
self.check_dependencies()
|
||||
return ToolStatus.AVAILABLE
|
||||
except DependencyError:
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def check_dependencies(self) -> None:
|
||||
"""Verify all dependencies are installed. Raises DependencyError if not."""
|
||||
for dep in self.dependencies:
|
||||
if dep.startswith("cmd:"):
|
||||
cmd_name = dep[4:]
|
||||
if shutil.which(cmd_name) is None:
|
||||
raise DependencyError(
|
||||
f"Command {cmd_name!r} not found. {self.install_instructions}"
|
||||
)
|
||||
elif dep.startswith("env:"):
|
||||
env_name = dep[4:]
|
||||
if not os.environ.get(env_name):
|
||||
raise DependencyError(
|
||||
f"Environment variable {env_name!r} not set. {self.install_instructions}"
|
||||
)
|
||||
elif dep.startswith("python:"):
|
||||
module_name = dep[7:]
|
||||
try:
|
||||
__import__(module_name)
|
||||
except ImportError:
|
||||
raise DependencyError(
|
||||
f"Python module {module_name!r} not installed. {self.install_instructions}"
|
||||
)
|
||||
|
||||
def get_info(self) -> dict[str, Any]:
|
||||
"""Return full tool contract info for registry/discovery."""
|
||||
usage_location = inspect.getfile(self.__class__)
|
||||
return {
|
||||
"name": self.name,
|
||||
"version": self.version,
|
||||
"tier": self.tier.value,
|
||||
"capability": self.capability,
|
||||
"provider": self.provider,
|
||||
"stability": self.stability.value,
|
||||
"status": self.get_status().value,
|
||||
"execution_mode": self.execution_mode.value,
|
||||
"determinism": self.determinism.value,
|
||||
"runtime": self.runtime.value,
|
||||
"module_path": self.__class__.__module__,
|
||||
"usage_location": usage_location,
|
||||
"dependencies": self.dependencies,
|
||||
"install_instructions": self.install_instructions,
|
||||
"capabilities": self.capabilities,
|
||||
"input_schema": self.input_schema,
|
||||
"output_schema": self.output_schema,
|
||||
"artifact_schema": self.artifact_schema,
|
||||
"supports": self.supports,
|
||||
"best_for": self.best_for,
|
||||
"not_good_for": self.not_good_for,
|
||||
"provider_matrix": self.provider_matrix,
|
||||
"resource_profile": {
|
||||
"cpu_cores": self.resource_profile.cpu_cores,
|
||||
"ram_mb": self.resource_profile.ram_mb,
|
||||
"vram_mb": self.resource_profile.vram_mb,
|
||||
"disk_mb": self.resource_profile.disk_mb,
|
||||
"network_required": self.resource_profile.network_required,
|
||||
},
|
||||
"resume_support": self.resume_support.value,
|
||||
"side_effects": self.side_effects,
|
||||
"fallback": self.fallback,
|
||||
"fallback_tools": self.fallback_tools or ([self.fallback] if self.fallback else []),
|
||||
"agent_skills": self.agent_skills,
|
||||
"related_skills": self.agent_skills,
|
||||
"user_visible_verification": self.user_visible_verification,
|
||||
}
|
||||
|
||||
# ---- Cost estimation ----
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
"""Estimate cost in USD for the given inputs. Override for paid tools."""
|
||||
return 0.0
|
||||
|
||||
def estimate_runtime(self, inputs: dict[str, Any]) -> float:
|
||||
"""Estimate runtime in seconds. Override for long-running tools."""
|
||||
return 0.0
|
||||
|
||||
# ---- Idempotency ----
|
||||
|
||||
def idempotency_key(self, inputs: dict[str, Any]) -> str:
|
||||
"""Compute a cache key from idempotency fields."""
|
||||
key_data = {k: inputs.get(k) for k in self.idempotency_key_fields}
|
||||
raw = json.dumps(key_data, sort_keys=True)
|
||||
return hashlib.sha256(raw.encode()).hexdigest()[:16]
|
||||
|
||||
# ---- Execution ----
|
||||
|
||||
@abstractmethod
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
"""Run the tool. Subclasses must implement this."""
|
||||
...
|
||||
|
||||
def dry_run(self, inputs: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Preflight check without side effects. Override for paid/publishing tools."""
|
||||
return {
|
||||
"tool": self.name,
|
||||
"estimated_cost_usd": self.estimate_cost(inputs),
|
||||
"estimated_runtime_seconds": self.estimate_runtime(inputs),
|
||||
"status": self.get_status().value,
|
||||
"would_execute": True,
|
||||
}
|
||||
|
||||
# ---- CLI helper ----
|
||||
|
||||
def run_command(
|
||||
self,
|
||||
cmd: list[str],
|
||||
*,
|
||||
timeout: Optional[int] = None,
|
||||
cwd: Optional[Path] = None,
|
||||
) -> subprocess.CompletedProcess:
|
||||
"""Run a subprocess command with standard error handling."""
|
||||
return subprocess.run(
|
||||
cmd,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=timeout,
|
||||
cwd=cwd,
|
||||
check=True,
|
||||
)
|
||||
|
||||
|
||||
class DependencyError(Exception):
|
||||
"""Raised when a tool's dependency is not satisfied."""
|
||||
pass
|
||||
@@ -0,0 +1,212 @@
|
||||
"""Cost tracker core: estimate, reserve, reconcile, and persist to cost_log.json.
|
||||
|
||||
Implements the budget governance rules from the spec:
|
||||
- Every paid operation produces a preflight estimate
|
||||
- The orchestrator reserves estimated budget before execution
|
||||
- Budget overruns trigger pauses (in warn/cap mode)
|
||||
- Actual spend is reconciled when the tool finishes or fails
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
|
||||
from lib.config_model import BudgetMode
|
||||
|
||||
|
||||
class EntryStatus(str, Enum):
|
||||
ESTIMATED = "estimated"
|
||||
RESERVED = "reserved"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
REFUNDED = "refunded"
|
||||
|
||||
|
||||
class BudgetExceededError(Exception):
|
||||
"""Raised when an operation would exceed the budget in cap mode."""
|
||||
pass
|
||||
|
||||
|
||||
class ApprovalRequiredError(Exception):
|
||||
"""Raised when an operation needs user approval before proceeding."""
|
||||
pass
|
||||
|
||||
|
||||
class CostTracker:
|
||||
"""Tracks estimated, reserved, and actual costs for a pipeline project."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
budget_total_usd: float = 10.0,
|
||||
reserve_pct: float = 0.10,
|
||||
single_action_approval_usd: float = 0.50,
|
||||
require_approval_for_new_paid_tool: bool = True,
|
||||
mode: BudgetMode = BudgetMode.WARN,
|
||||
cost_log_path: Optional[Path] = None,
|
||||
) -> None:
|
||||
self.budget_total_usd = budget_total_usd
|
||||
self.reserve_pct = reserve_pct
|
||||
self.single_action_approval_usd = single_action_approval_usd
|
||||
self.require_approval_for_new_paid_tool = require_approval_for_new_paid_tool
|
||||
self.mode = mode
|
||||
self.cost_log_path = cost_log_path
|
||||
self.entries: list[dict[str, Any]] = []
|
||||
self._approved_tools: set[str] = set()
|
||||
|
||||
if cost_log_path and cost_log_path.exists():
|
||||
self._load()
|
||||
|
||||
# ---- Budget calculations ----
|
||||
|
||||
@property
|
||||
def budget_reserved_usd(self) -> float:
|
||||
return sum(
|
||||
e.get("reserved_usd", 0.0)
|
||||
for e in self.entries
|
||||
if e["status"] == EntryStatus.RESERVED.value
|
||||
)
|
||||
|
||||
@property
|
||||
def budget_spent_usd(self) -> float:
|
||||
return sum(
|
||||
e.get("actual_usd", 0.0)
|
||||
for e in self.entries
|
||||
if e["status"] in (EntryStatus.COMPLETED.value, EntryStatus.FAILED.value)
|
||||
)
|
||||
|
||||
@property
|
||||
def budget_remaining_usd(self) -> float:
|
||||
return self.budget_total_usd - self.budget_spent_usd - self.budget_reserved_usd
|
||||
|
||||
@property
|
||||
def usable_budget_usd(self) -> float:
|
||||
"""Budget minus the reserve holdback."""
|
||||
holdback = self.budget_total_usd * self.reserve_pct
|
||||
return max(0.0, self.budget_remaining_usd - holdback)
|
||||
|
||||
def cost_snapshot(self) -> dict[str, float]:
|
||||
return {
|
||||
"total_spent_usd": round(self.budget_spent_usd, 4),
|
||||
"total_reserved_usd": round(self.budget_reserved_usd, 4),
|
||||
"budget_remaining_usd": round(self.budget_remaining_usd, 4),
|
||||
}
|
||||
|
||||
# ---- Core operations ----
|
||||
|
||||
def estimate(self, tool: str, operation: str, estimated_usd: float) -> str:
|
||||
"""Record an estimate. Returns entry ID."""
|
||||
entry_id = self._new_id()
|
||||
self.entries.append({
|
||||
"id": entry_id,
|
||||
"tool": tool,
|
||||
"operation": operation,
|
||||
"status": EntryStatus.ESTIMATED.value,
|
||||
"estimated_usd": round(estimated_usd, 4),
|
||||
"reserved_usd": 0.0,
|
||||
"actual_usd": 0.0,
|
||||
"timestamp": self._now(),
|
||||
})
|
||||
self._save()
|
||||
return entry_id
|
||||
|
||||
def reserve(self, entry_id: str) -> None:
|
||||
"""Reserve budget for an estimated entry.
|
||||
|
||||
Raises BudgetExceededError in cap mode, or ApprovalRequiredError
|
||||
when the action exceeds the single-action approval threshold.
|
||||
"""
|
||||
entry = self._find(entry_id)
|
||||
estimated = entry["estimated_usd"]
|
||||
|
||||
# Check single-action approval threshold
|
||||
if estimated > self.single_action_approval_usd:
|
||||
if self.mode != BudgetMode.OBSERVE:
|
||||
raise ApprovalRequiredError(
|
||||
f"Action costs ${estimated:.2f}, exceeds "
|
||||
f"single-action threshold ${self.single_action_approval_usd:.2f}"
|
||||
)
|
||||
|
||||
# Check new paid tool approval
|
||||
if self.require_approval_for_new_paid_tool and estimated > 0:
|
||||
if entry["tool"] not in self._approved_tools:
|
||||
if self.mode != BudgetMode.OBSERVE:
|
||||
raise ApprovalRequiredError(
|
||||
f"First paid use of tool {entry['tool']!r} requires approval"
|
||||
)
|
||||
|
||||
# Check budget
|
||||
if estimated > self.usable_budget_usd:
|
||||
if self.mode == BudgetMode.CAP:
|
||||
raise BudgetExceededError(
|
||||
f"Reservation of ${estimated:.2f} exceeds usable budget "
|
||||
f"${self.usable_budget_usd:.2f}"
|
||||
)
|
||||
|
||||
entry["status"] = EntryStatus.RESERVED.value
|
||||
entry["reserved_usd"] = estimated
|
||||
entry["timestamp"] = self._now()
|
||||
self._save()
|
||||
|
||||
def approve_tool(self, tool: str) -> None:
|
||||
"""Mark a tool as approved for paid operations."""
|
||||
self._approved_tools.add(tool)
|
||||
|
||||
def reconcile(self, entry_id: str, actual_usd: float, success: bool = True) -> None:
|
||||
"""Reconcile actual spend after tool execution."""
|
||||
entry = self._find(entry_id)
|
||||
entry["status"] = EntryStatus.COMPLETED.value if success else EntryStatus.FAILED.value
|
||||
entry["actual_usd"] = round(actual_usd, 4)
|
||||
entry["reserved_usd"] = 0.0
|
||||
entry["timestamp"] = self._now()
|
||||
self._save()
|
||||
|
||||
def refund(self, entry_id: str) -> None:
|
||||
"""Cancel a reservation without executing."""
|
||||
entry = self._find(entry_id)
|
||||
entry["status"] = EntryStatus.REFUNDED.value
|
||||
entry["reserved_usd"] = 0.0
|
||||
entry["timestamp"] = self._now()
|
||||
self._save()
|
||||
|
||||
# ---- Persistence ----
|
||||
|
||||
def _save(self) -> None:
|
||||
if self.cost_log_path is None:
|
||||
return
|
||||
data = {
|
||||
"version": "1.0",
|
||||
"budget_total_usd": self.budget_total_usd,
|
||||
"budget_reserved_usd": round(self.budget_reserved_usd, 4),
|
||||
"budget_spent_usd": round(self.budget_spent_usd, 4),
|
||||
"entries": self.entries,
|
||||
}
|
||||
self.cost_log_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(self.cost_log_path, "w") as f:
|
||||
json.dump(data, f, indent=2)
|
||||
|
||||
def _load(self) -> None:
|
||||
with open(self.cost_log_path) as f: # type: ignore[arg-type]
|
||||
data = json.load(f)
|
||||
self.entries = data.get("entries", [])
|
||||
self.budget_total_usd = data.get("budget_total_usd", self.budget_total_usd)
|
||||
|
||||
# ---- Helpers ----
|
||||
|
||||
def _find(self, entry_id: str) -> dict[str, Any]:
|
||||
for entry in self.entries:
|
||||
if entry["id"] == entry_id:
|
||||
return entry
|
||||
raise KeyError(f"Cost entry {entry_id!r} not found")
|
||||
|
||||
@staticmethod
|
||||
def _new_id() -> str:
|
||||
return uuid.uuid4().hex[:12]
|
||||
|
||||
@staticmethod
|
||||
def _now() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
@@ -0,0 +1 @@
|
||||
"""Enhancement tools for image and video quality improvement."""
|
||||
@@ -0,0 +1,161 @@
|
||||
"""Background removal tool wrapping rembg.
|
||||
|
||||
Removes backgrounds from images using the rembg library (U2Net models).
|
||||
Outputs transparent PNGs or composites onto a custom background color.
|
||||
Supports local execution via rembg and optionally cloud APIs.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class BgRemove(BaseTool):
|
||||
name = "bg_remove"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.ENHANCE
|
||||
capability = "enhancement"
|
||||
provider = "rembg"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
runtime = ToolRuntime.HYBRID
|
||||
|
||||
dependencies = ["python:rembg", "python:PIL"]
|
||||
install_instructions = (
|
||||
"pip install rembg # CPU mode\n"
|
||||
"pip install rembg[gpu] # GPU mode (requires CUDA + onnxruntime-gpu)"
|
||||
)
|
||||
agent_skills = ["ffmpeg"]
|
||||
|
||||
capabilities = [
|
||||
"background_removal",
|
||||
"alpha_matte",
|
||||
"batch_processing",
|
||||
"custom_background",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["input_path"],
|
||||
"properties": {
|
||||
"input_path": {
|
||||
"type": "string",
|
||||
"description": "Path to image or video frame",
|
||||
},
|
||||
"output_path": {
|
||||
"type": "string",
|
||||
"description": "Output path; defaults to {stem}_nobg.png",
|
||||
},
|
||||
"model": {
|
||||
"type": "string",
|
||||
"enum": ["u2net", "u2net_human_seg", "isnet-general-use"],
|
||||
"default": "u2net",
|
||||
},
|
||||
"bg_color": {
|
||||
"type": "string",
|
||||
"description": "Replacement background color hex (e.g. #00FF00). Transparent if not set.",
|
||||
},
|
||||
"alpha_matting": {
|
||||
"type": "boolean",
|
||||
"default": False,
|
||||
"description": "Use alpha matting for finer edges",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=2, ram_mb=2048, vram_mb=0, disk_mb=500
|
||||
)
|
||||
|
||||
idempotency_key_fields = ["input_path", "model", "bg_color", "alpha_matting"]
|
||||
side_effects = ["writes background-removed image to output_path"]
|
||||
user_visible_verification = [
|
||||
"Inspect output for clean edges around the subject",
|
||||
"Verify transparency or background color is applied correctly",
|
||||
]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
try:
|
||||
import rembg # noqa: F401
|
||||
return ToolStatus.AVAILABLE
|
||||
except ImportError:
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
input_path = Path(inputs["input_path"])
|
||||
if not input_path.exists():
|
||||
return ToolResult(success=False, error=f"Input not found: {input_path}")
|
||||
|
||||
output_path = Path(
|
||||
inputs.get("output_path", str(input_path.with_stem(f"{input_path.stem}_nobg").with_suffix(".png")))
|
||||
)
|
||||
model_name = inputs.get("model", "u2net")
|
||||
bg_color = inputs.get("bg_color")
|
||||
alpha_matting = inputs.get("alpha_matting", False)
|
||||
|
||||
try:
|
||||
import rembg
|
||||
except ImportError:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="rembg is not installed. Run: pip install rembg",
|
||||
)
|
||||
|
||||
try:
|
||||
from PIL import Image
|
||||
except ImportError:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="Pillow is not installed. Run: pip install Pillow",
|
||||
)
|
||||
|
||||
start = time.time()
|
||||
|
||||
input_image = Image.open(input_path)
|
||||
|
||||
result_image = rembg.remove(
|
||||
input_image,
|
||||
model_name=model_name,
|
||||
alpha_matting=alpha_matting,
|
||||
)
|
||||
|
||||
# Composite onto a colored background if requested
|
||||
if bg_color:
|
||||
hex_color = bg_color.lstrip("#")
|
||||
r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
|
||||
background = Image.new("RGBA", result_image.size, (r, g, b, 255))
|
||||
background.paste(result_image, mask=result_image.split()[3])
|
||||
result_image = background.convert("RGB")
|
||||
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
result_image.save(str(output_path))
|
||||
|
||||
elapsed = time.time() - start
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"input": str(input_path),
|
||||
"output": str(output_path),
|
||||
"model": model_name,
|
||||
"alpha_matting": alpha_matting,
|
||||
"bg_color": bg_color,
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
duration_seconds=round(elapsed, 2),
|
||||
)
|
||||
@@ -0,0 +1,212 @@
|
||||
"""Color grading tool wrapping FFmpeg LUT and filter chains.
|
||||
|
||||
Applies cinematic color grading profiles to video. Supports both
|
||||
built-in profile presets and external .cube LUT files.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
ToolResult,
|
||||
ToolStability,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
# Built-in grading profiles using FFmpeg colorbalance/curves/eq filters
|
||||
PROFILES = {
|
||||
"cinematic_warm": {
|
||||
"description": "Warm cinematic look with lifted shadows and orange highlights",
|
||||
"vf": (
|
||||
"colorbalance=rs=0.08:gs=0.02:bs=-0.05:rh=0.06:gh=0.02:bh=-0.04,"
|
||||
"curves=all='0/0.03 0.25/0.22 0.5/0.50 0.75/0.78 1/0.97',"
|
||||
"eq=contrast=1.05:saturation=1.1"
|
||||
),
|
||||
},
|
||||
"cinematic_cool": {
|
||||
"description": "Cool teal-and-orange cinematic grade",
|
||||
"vf": (
|
||||
"colorbalance=rs=-0.02:gs=-0.03:bs=0.08:rh=0.06:gh=-0.02:bh=-0.06,"
|
||||
"curves=all='0/0.02 0.25/0.20 0.5/0.48 0.75/0.78 1/0.98',"
|
||||
"eq=contrast=1.08:saturation=1.05"
|
||||
),
|
||||
},
|
||||
"moody_dark": {
|
||||
"description": "Crushed blacks, desaturated midtones, dark atmosphere",
|
||||
"vf": (
|
||||
"curves=all='0/0.05 0.15/0.12 0.5/0.45 0.85/0.82 1/0.95',"
|
||||
"eq=contrast=1.12:saturation=0.8:brightness=-0.03"
|
||||
),
|
||||
},
|
||||
"bright_clean": {
|
||||
"description": "Bright, clean look with lifted shadows and vivid color",
|
||||
"vf": (
|
||||
"curves=all='0/0.05 0.25/0.30 0.5/0.55 0.75/0.80 1/1.0',"
|
||||
"eq=contrast=1.0:saturation=1.15:brightness=0.02"
|
||||
),
|
||||
},
|
||||
"vintage_film": {
|
||||
"description": "Faded film look with grain texture and warm tint",
|
||||
"vf": (
|
||||
"colorbalance=rs=0.06:gs=0.03:bs=-0.03:ms=0.03:mh=-0.02,"
|
||||
"curves=all='0/0.06 0.25/0.25 0.5/0.50 0.75/0.74 1/0.94',"
|
||||
"eq=saturation=0.85:contrast=0.95"
|
||||
),
|
||||
},
|
||||
"high_contrast": {
|
||||
"description": "Punchy high-contrast grade for dynamic content",
|
||||
"vf": (
|
||||
"curves=all='0/0 0.20/0.12 0.5/0.50 0.80/0.88 1/1',"
|
||||
"eq=contrast=1.2:saturation=1.1"
|
||||
),
|
||||
},
|
||||
"neutral": {
|
||||
"description": "Minimal correction — normalize levels and light contrast",
|
||||
"vf": "eq=contrast=1.02:saturation=1.02:brightness=0.01",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class ColorGrade(BaseTool):
|
||||
name = "color_grade"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.CORE
|
||||
capability = "enhancement"
|
||||
provider = "ffmpeg"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
|
||||
dependencies = ["cmd:ffmpeg"]
|
||||
install_instructions = "Install FFmpeg: https://ffmpeg.org/download.html"
|
||||
agent_skills = ["ffmpeg"]
|
||||
|
||||
capabilities = [
|
||||
"grade_preset",
|
||||
"grade_lut",
|
||||
"grade_custom",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["input_path"],
|
||||
"properties": {
|
||||
"input_path": {"type": "string"},
|
||||
"output_path": {"type": "string"},
|
||||
"profile": {
|
||||
"type": "string",
|
||||
"enum": list(PROFILES.keys()),
|
||||
"default": "cinematic_warm",
|
||||
},
|
||||
"lut_path": {
|
||||
"type": "string",
|
||||
"description": "Path to external .cube LUT file",
|
||||
},
|
||||
"intensity": {
|
||||
"type": "number",
|
||||
"minimum": 0.0,
|
||||
"maximum": 1.0,
|
||||
"default": 1.0,
|
||||
"description": "Blend intensity: 0 = original, 1 = full grade",
|
||||
},
|
||||
"custom_vf": {"type": "string"},
|
||||
"codec": {"type": "string", "default": "libx264"},
|
||||
"crf": {"type": "integer", "default": 20},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(cpu_cores=2, ram_mb=1024, vram_mb=0, disk_mb=2000)
|
||||
idempotency_key_fields = ["input_path", "profile", "lut_path", "intensity"]
|
||||
side_effects = ["writes graded video to output_path"]
|
||||
user_visible_verification = [
|
||||
"Compare graded output with original for color accuracy",
|
||||
"Verify skin tones look natural, not oversaturated",
|
||||
]
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
input_path = Path(inputs["input_path"])
|
||||
if not input_path.exists():
|
||||
return ToolResult(success=False, error=f"Input not found: {input_path}")
|
||||
|
||||
output_path = Path(
|
||||
inputs.get("output_path", str(input_path.with_stem(f"{input_path.stem}_graded")))
|
||||
)
|
||||
codec = inputs.get("codec", "libx264")
|
||||
crf = inputs.get("crf", 20)
|
||||
|
||||
vf = self._build_filter(inputs)
|
||||
if not vf:
|
||||
return ToolResult(success=False, error="No profile, lut_path, or custom_vf specified")
|
||||
|
||||
start = time.time()
|
||||
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", str(input_path),
|
||||
"-vf", vf,
|
||||
"-c:v", codec, "-crf", str(crf),
|
||||
"-c:a", "copy",
|
||||
str(output_path),
|
||||
]
|
||||
|
||||
try:
|
||||
self.run_command(cmd)
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"FFmpeg failed: {e}")
|
||||
|
||||
elapsed = time.time() - start
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"input": str(input_path),
|
||||
"output": str(output_path),
|
||||
"profile": inputs.get("profile"),
|
||||
"lut": inputs.get("lut_path"),
|
||||
"intensity": inputs.get("intensity", 1.0),
|
||||
"filter": vf,
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
duration_seconds=round(elapsed, 2),
|
||||
)
|
||||
|
||||
def _build_filter(self, inputs: dict[str, Any]) -> str:
|
||||
if "custom_vf" in inputs:
|
||||
return inputs["custom_vf"]
|
||||
|
||||
lut_path = inputs.get("lut_path")
|
||||
if lut_path and Path(lut_path).exists():
|
||||
safe_path = str(Path(lut_path).resolve()).replace("\\", "/").replace(":", "\\:")
|
||||
return f"lut3d='{safe_path}'"
|
||||
|
||||
profile_name = inputs.get("profile", "cinematic_warm")
|
||||
profile = PROFILES.get(profile_name)
|
||||
if not profile:
|
||||
return ""
|
||||
|
||||
vf = profile["vf"]
|
||||
|
||||
# Apply intensity blending if < 1.0
|
||||
intensity = inputs.get("intensity", 1.0)
|
||||
if 0 < intensity < 1.0:
|
||||
# Use split + overlay approach: blend graded with original
|
||||
vf = (
|
||||
f"split[original][tograde];"
|
||||
f"[tograde]{vf}[graded];"
|
||||
f"[original][graded]blend=all_mode=normal:all_opacity={intensity}"
|
||||
)
|
||||
|
||||
return vf
|
||||
|
||||
@staticmethod
|
||||
def list_profiles() -> dict[str, str]:
|
||||
"""Return available profiles and their descriptions."""
|
||||
return {name: p["description"] for name, p in PROFILES.items()}
|
||||
@@ -0,0 +1,193 @@
|
||||
"""Face enhancement tool wrapping FFmpeg filters.
|
||||
|
||||
Applies skin smoothing, sharpening, and lighting correction presets
|
||||
to talking-head footage. All presets are FFmpeg filter chains — no GPU
|
||||
or external models required.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
ToolResult,
|
||||
ToolStability,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
# Named presets mapping to FFmpeg filter chains
|
||||
PRESETS = {
|
||||
"soft_skin": {
|
||||
"description": "Gentle skin smoothing while preserving edges",
|
||||
"vf": "smartblur=lr=1.0:ls=-0.5:lt=-3.0:cr=0.5:cs=-0.5:ct=-3.0",
|
||||
},
|
||||
"sharpen": {
|
||||
"description": "Edge sharpening for crisp detail",
|
||||
"vf": "unsharp=5:5:1.0:5:5:0.0",
|
||||
},
|
||||
"sharpen_light": {
|
||||
"description": "Subtle sharpening for soft cameras",
|
||||
"vf": "unsharp=3:3:0.5:3:3:0.0",
|
||||
},
|
||||
"brighten": {
|
||||
"description": "Lift shadows and midtones for poorly lit footage",
|
||||
"vf": "curves=all='0/0 0.25/0.35 0.5/0.55 0.75/0.8 1/1'",
|
||||
},
|
||||
"contrast_boost": {
|
||||
"description": "Add punch with an S-curve contrast adjustment",
|
||||
"vf": "curves=all='0/0 0.25/0.20 0.5/0.5 0.75/0.80 1/1'",
|
||||
},
|
||||
"warm": {
|
||||
"description": "Warm skin tones — slight orange shift",
|
||||
"vf": "colorbalance=rs=0.05:gs=0.0:bs=-0.05:rm=0.05:gm=0.0:bm=-0.03",
|
||||
},
|
||||
"cool": {
|
||||
"description": "Cool tones — slight blue shift",
|
||||
"vf": "colorbalance=rs=-0.03:gs=0.0:bs=0.05:rm=-0.02:gm=0.0:bm=0.03",
|
||||
},
|
||||
"denoise": {
|
||||
"description": "Temporal noise reduction for grainy footage",
|
||||
"vf": "hqdn3d=4:3:6:4",
|
||||
},
|
||||
"talking_head_standard": {
|
||||
"description": "Combined preset: skin smoothing + sharpen edges + warm skin tones",
|
||||
"vf": (
|
||||
"smartblur=lr=1.0:ls=-0.5:lt=-3.0:cr=0.5:cs=-0.5:ct=-3.0,"
|
||||
"unsharp=5:5:0.6:5:5:0.0,"
|
||||
"colorbalance=rs=0.06:gs=0.01:bs=-0.04:rm=0.04:gm=0.01:bm=-0.03"
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class FaceEnhance(BaseTool):
|
||||
name = "face_enhance"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.CORE
|
||||
capability = "enhancement"
|
||||
provider = "ffmpeg"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
|
||||
dependencies = ["cmd:ffmpeg"]
|
||||
install_instructions = "Install FFmpeg: https://ffmpeg.org/download.html"
|
||||
agent_skills = ["ffmpeg"]
|
||||
|
||||
capabilities = [
|
||||
"skin_smoothing",
|
||||
"sharpening",
|
||||
"lighting_correction",
|
||||
"color_balance",
|
||||
"denoise",
|
||||
"preset_chain",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["input_path"],
|
||||
"properties": {
|
||||
"input_path": {"type": "string"},
|
||||
"output_path": {"type": "string"},
|
||||
"preset": {
|
||||
"type": "string",
|
||||
"enum": list(PRESETS.keys()),
|
||||
"default": "talking_head_standard",
|
||||
},
|
||||
"presets": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
"description": "Apply multiple presets in sequence",
|
||||
},
|
||||
"custom_vf": {
|
||||
"type": "string",
|
||||
"description": "Custom FFmpeg video filter string (advanced)",
|
||||
},
|
||||
"codec": {"type": "string", "default": "libx264"},
|
||||
"crf": {"type": "integer", "default": 20},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(cpu_cores=2, ram_mb=1024, vram_mb=0, disk_mb=2000)
|
||||
idempotency_key_fields = ["input_path", "preset", "presets", "custom_vf"]
|
||||
side_effects = ["writes enhanced video to output_path"]
|
||||
user_visible_verification = [
|
||||
"Compare enhanced output with original side-by-side",
|
||||
"Verify skin texture is natural, not plastic",
|
||||
]
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
input_path = Path(inputs["input_path"])
|
||||
if not input_path.exists():
|
||||
return ToolResult(success=False, error=f"Input not found: {input_path}")
|
||||
|
||||
output_path = Path(
|
||||
inputs.get("output_path", str(input_path.with_stem(f"{input_path.stem}_enhanced")))
|
||||
)
|
||||
codec = inputs.get("codec", "libx264")
|
||||
crf = inputs.get("crf", 20)
|
||||
|
||||
# Build filter chain
|
||||
vf = self._build_filter(inputs)
|
||||
if not vf:
|
||||
return ToolResult(success=False, error="No preset, presets, or custom_vf specified")
|
||||
|
||||
start = time.time()
|
||||
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", str(input_path),
|
||||
"-vf", vf,
|
||||
"-c:v", codec, "-crf", str(crf),
|
||||
"-c:a", "copy",
|
||||
str(output_path),
|
||||
]
|
||||
|
||||
try:
|
||||
self.run_command(cmd)
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"FFmpeg failed: {e}")
|
||||
|
||||
elapsed = time.time() - start
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"input": str(input_path),
|
||||
"output": str(output_path),
|
||||
"filter": vf,
|
||||
"preset": inputs.get("preset"),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
duration_seconds=round(elapsed, 2),
|
||||
)
|
||||
|
||||
def _build_filter(self, inputs: dict[str, Any]) -> str:
|
||||
if "custom_vf" in inputs:
|
||||
return inputs["custom_vf"]
|
||||
|
||||
if "presets" in inputs:
|
||||
chains = []
|
||||
for name in inputs["presets"]:
|
||||
if name not in PRESETS:
|
||||
continue
|
||||
chains.append(PRESETS[name]["vf"])
|
||||
return ",".join(chains)
|
||||
|
||||
preset_name = inputs.get("preset", "talking_head_standard")
|
||||
preset = PRESETS.get(preset_name)
|
||||
if preset:
|
||||
return preset["vf"]
|
||||
return ""
|
||||
|
||||
@staticmethod
|
||||
def list_presets() -> dict[str, str]:
|
||||
"""Return available presets and their descriptions."""
|
||||
return {name: p["description"] for name, p in PRESETS.items()}
|
||||
@@ -0,0 +1,230 @@
|
||||
"""Face restoration tool wrapping CodeFormer / GFPGAN.
|
||||
|
||||
Restores degraded or low-quality faces in images and video frames.
|
||||
Fixes blur, compression artifacts, and low resolution specifically on
|
||||
face regions while preserving the rest of the image.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class FaceRestore(BaseTool):
|
||||
name = "face_restore"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.ENHANCE
|
||||
capability = "enhancement"
|
||||
provider = "codeformer"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
runtime = ToolRuntime.LOCAL_GPU
|
||||
|
||||
dependencies = ["python:gfpgan", "python:torch"]
|
||||
install_instructions = (
|
||||
"pip install gfpgan # Includes CodeFormer support. Requires PyTorch."
|
||||
)
|
||||
agent_skills = ["ffmpeg"]
|
||||
fallback = None
|
||||
|
||||
capabilities = [
|
||||
"face_restoration",
|
||||
"face_detection",
|
||||
"quality_enhancement",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["input_path"],
|
||||
"properties": {
|
||||
"input_path": {
|
||||
"type": "string",
|
||||
"description": "Path to image or video frame",
|
||||
},
|
||||
"output_path": {
|
||||
"type": "string",
|
||||
"description": "Output path (defaults to {stem}_restored.{ext})",
|
||||
},
|
||||
"model": {
|
||||
"type": "string",
|
||||
"enum": ["CodeFormer", "GFPGAN"],
|
||||
"default": "CodeFormer",
|
||||
"description": "Restoration model to use",
|
||||
},
|
||||
"fidelity": {
|
||||
"type": "number",
|
||||
"default": 0.5,
|
||||
"description": (
|
||||
"0 = max quality, 1 = max fidelity to input (CodeFormer only)"
|
||||
),
|
||||
},
|
||||
"upscale": {
|
||||
"type": "integer",
|
||||
"default": 2,
|
||||
"description": "Face upscale factor",
|
||||
},
|
||||
"bg_upsampler": {
|
||||
"type": "boolean",
|
||||
"default": False,
|
||||
"description": "Also upscale background with Real-ESRGAN",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=2, ram_mb=2048, vram_mb=2048, disk_mb=1000
|
||||
)
|
||||
idempotency_key_fields = ["input_path", "model", "fidelity", "upscale"]
|
||||
side_effects = ["writes restored image to output_path"]
|
||||
user_visible_verification = [
|
||||
"Compare restored face with original for natural appearance",
|
||||
"Verify face identity is preserved after restoration",
|
||||
]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
try:
|
||||
import gfpgan # noqa: F401
|
||||
return ToolStatus.AVAILABLE
|
||||
except ImportError:
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
input_path = Path(inputs["input_path"])
|
||||
if not input_path.exists():
|
||||
return ToolResult(success=False, error=f"Input not found: {input_path}")
|
||||
|
||||
output_path = Path(
|
||||
inputs.get(
|
||||
"output_path",
|
||||
str(input_path.with_stem(f"{input_path.stem}_restored")),
|
||||
)
|
||||
)
|
||||
model_name = inputs.get("model", "CodeFormer")
|
||||
fidelity = inputs.get("fidelity", 0.5)
|
||||
upscale = inputs.get("upscale", 2)
|
||||
bg_upsampler_flag = inputs.get("bg_upsampler", False)
|
||||
|
||||
try:
|
||||
import cv2
|
||||
from gfpgan import GFPGANer
|
||||
except ImportError as e:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"Missing dependency: {e}. Run: pip install gfpgan",
|
||||
)
|
||||
|
||||
start = time.time()
|
||||
|
||||
# Optional background upsampler
|
||||
bg_upsampler = None
|
||||
if bg_upsampler_flag:
|
||||
try:
|
||||
from basicsr.archs.rrdbnet_arch import RRDBNet
|
||||
from realesrgan import RealESRGANer
|
||||
|
||||
realesrgan_model = RRDBNet(
|
||||
num_in_ch=3, num_out_ch=3, num_feat=64,
|
||||
num_block=23, num_grow_ch=32, scale=2,
|
||||
)
|
||||
bg_upsampler = RealESRGANer(
|
||||
scale=2,
|
||||
model_path=(
|
||||
"https://github.com/xinntao/Real-ESRGAN/releases/download/"
|
||||
"v0.2.1/RealESRGAN_x2plus.pth"
|
||||
),
|
||||
model=realesrgan_model,
|
||||
tile=400,
|
||||
tile_pad=10,
|
||||
pre_pad=0,
|
||||
half=True,
|
||||
)
|
||||
except ImportError:
|
||||
bg_upsampler = None
|
||||
|
||||
# Select model path based on model choice
|
||||
if model_name == "CodeFormer":
|
||||
model_path = (
|
||||
"https://github.com/sczhou/CodeFormer/releases/download/"
|
||||
"v0.1.0/codeformer.pth"
|
||||
)
|
||||
arch = "CodeFormer"
|
||||
else:
|
||||
model_path = (
|
||||
"https://github.com/TencentARC/GFPGAN/releases/download/"
|
||||
"v1.3.0/GFPGANv1.3.pth"
|
||||
)
|
||||
arch = "clean"
|
||||
|
||||
# Instantiate restorer
|
||||
try:
|
||||
restorer = GFPGANer(
|
||||
model_path=model_path,
|
||||
upscale=upscale,
|
||||
arch=arch,
|
||||
bg_upsampler=bg_upsampler,
|
||||
)
|
||||
except Exception as e:
|
||||
return ToolResult(
|
||||
success=False, error=f"Failed to load {model_name} model: {e}"
|
||||
)
|
||||
|
||||
# Read input image
|
||||
input_img = cv2.imread(str(input_path), cv2.IMREAD_COLOR)
|
||||
if input_img is None:
|
||||
return ToolResult(
|
||||
success=False, error=f"Failed to read image: {input_path}"
|
||||
)
|
||||
|
||||
# Run restoration
|
||||
try:
|
||||
_, restored_faces, restored_img = restorer.enhance(
|
||||
input_img,
|
||||
has_aligned=False,
|
||||
only_center_face=False,
|
||||
paste_back=True,
|
||||
weight=fidelity if model_name == "CodeFormer" else None,
|
||||
)
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Restoration failed: {e}")
|
||||
|
||||
if restored_img is None:
|
||||
return ToolResult(
|
||||
success=False, error="Restoration produced no output"
|
||||
)
|
||||
|
||||
# Save restored output
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
cv2.imwrite(str(output_path), restored_img)
|
||||
|
||||
elapsed = time.time() - start
|
||||
faces_detected = len(restored_faces) if restored_faces else 0
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"input": str(input_path),
|
||||
"output": str(output_path),
|
||||
"model": model_name,
|
||||
"faces_detected": faces_detected,
|
||||
"upscale": upscale,
|
||||
"fidelity": fidelity if model_name == "CodeFormer" else None,
|
||||
"bg_upsampler": bg_upsampler_flag,
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
duration_seconds=round(elapsed, 2),
|
||||
)
|
||||
@@ -0,0 +1,339 @@
|
||||
"""Image and video upscaling tool using Real-ESRGAN.
|
||||
|
||||
Takes low-resolution images or video and produces higher-resolution output
|
||||
(2x or 4x). For video, frames are extracted via FFmpeg, upscaled individually,
|
||||
and reassembled into the output file.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import shutil
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
VIDEO_EXTENSIONS = {".mp4", ".mov", ".avi"}
|
||||
|
||||
MODELS = {
|
||||
"RealESRGAN_x4plus": {
|
||||
"description": "General-purpose photo/video upscaler (default)",
|
||||
"scale": 4,
|
||||
},
|
||||
"RealESRGAN_x4plus_anime_6B": {
|
||||
"description": "Optimised for anime/illustration content",
|
||||
"scale": 4,
|
||||
},
|
||||
"RealESRNet_x4plus": {
|
||||
"description": "Lighter network, faster but lower quality",
|
||||
"scale": 4,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class Upscale(BaseTool):
|
||||
name = "upscale"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.ENHANCE
|
||||
capability = "enhancement"
|
||||
provider = "realesrgan"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
runtime = ToolRuntime.LOCAL_GPU
|
||||
|
||||
dependencies = ["python:realesrgan", "python:torch", "cmd:ffmpeg"]
|
||||
install_instructions = "pip install realesrgan # Requires PyTorch with CUDA"
|
||||
agent_skills = ["ffmpeg"]
|
||||
|
||||
capabilities = [
|
||||
"image_upscale",
|
||||
"video_upscale",
|
||||
"face_aware_upscale",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["input_path"],
|
||||
"properties": {
|
||||
"input_path": {"type": "string"},
|
||||
"output_path": {"type": "string"},
|
||||
"scale": {
|
||||
"type": "integer",
|
||||
"enum": [2, 4],
|
||||
"default": 4,
|
||||
},
|
||||
"model": {
|
||||
"type": "string",
|
||||
"enum": list(MODELS.keys()),
|
||||
"default": "RealESRGAN_x4plus",
|
||||
},
|
||||
"face_enhance": {
|
||||
"type": "boolean",
|
||||
"default": False,
|
||||
"description": "Use GFPGAN for face regions",
|
||||
},
|
||||
"denoise_strength": {
|
||||
"type": "number",
|
||||
"minimum": 0.0,
|
||||
"maximum": 1.0,
|
||||
"default": 0.5,
|
||||
"description": "Denoising strength (0 = no denoise, 1 = full)",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(cpu_cores=2, ram_mb=4096, vram_mb=2048, disk_mb=2000)
|
||||
idempotency_key_fields = ["input_path", "scale", "model", "face_enhance", "denoise_strength"]
|
||||
side_effects = ["writes upscaled file to output_path"]
|
||||
user_visible_verification = [
|
||||
"Compare upscaled output with original for detail and artifact quality",
|
||||
"Verify faces look natural if face_enhance was enabled",
|
||||
]
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Status
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
try:
|
||||
import realesrgan # noqa: F401
|
||||
return ToolStatus.AVAILABLE
|
||||
except ImportError:
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Execution
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
input_path = Path(inputs["input_path"])
|
||||
if not input_path.exists():
|
||||
return ToolResult(success=False, error=f"Input not found: {input_path}")
|
||||
|
||||
is_video = input_path.suffix.lower() in VIDEO_EXTENSIONS
|
||||
|
||||
default_output = str(input_path.with_stem(f"{input_path.stem}_upscaled"))
|
||||
output_path = Path(inputs.get("output_path", default_output))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
scale = inputs.get("scale", 4)
|
||||
model_name = inputs.get("model", "RealESRGAN_x4plus")
|
||||
face_enhance = inputs.get("face_enhance", False)
|
||||
denoise_strength = inputs.get("denoise_strength", 0.5)
|
||||
|
||||
start = time.time()
|
||||
|
||||
try:
|
||||
if is_video:
|
||||
result = self._upscale_video(
|
||||
input_path, output_path, scale, model_name,
|
||||
face_enhance, denoise_strength,
|
||||
)
|
||||
else:
|
||||
result = self._upscale_image(
|
||||
input_path, output_path, scale, model_name,
|
||||
face_enhance, denoise_strength,
|
||||
)
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Upscale failed: {e}")
|
||||
|
||||
elapsed = time.time() - start
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"input": str(input_path),
|
||||
"output": str(output_path),
|
||||
"scale": scale,
|
||||
"model": model_name,
|
||||
"face_enhance": face_enhance,
|
||||
"type": "video" if is_video else "image",
|
||||
**result,
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
duration_seconds=round(elapsed, 2),
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Image upscaling
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _upscale_image(
|
||||
self,
|
||||
input_path: Path,
|
||||
output_path: Path,
|
||||
scale: int,
|
||||
model_name: str,
|
||||
face_enhance: bool,
|
||||
denoise_strength: float,
|
||||
) -> dict[str, Any]:
|
||||
import cv2
|
||||
|
||||
upsampler = self._build_upsampler(scale, model_name, denoise_strength, face_enhance)
|
||||
|
||||
img = cv2.imread(str(input_path), cv2.IMREAD_UNCHANGED)
|
||||
if img is None:
|
||||
raise ValueError(f"Could not read image: {input_path}")
|
||||
|
||||
output, _ = upsampler.enhance(img, outscale=scale)
|
||||
cv2.imwrite(str(output_path), output)
|
||||
|
||||
h, w = output.shape[:2]
|
||||
return {"output_width": w, "output_height": h}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Video upscaling
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _upscale_video(
|
||||
self,
|
||||
input_path: Path,
|
||||
output_path: Path,
|
||||
scale: int,
|
||||
model_name: str,
|
||||
face_enhance: bool,
|
||||
denoise_strength: float,
|
||||
) -> dict[str, Any]:
|
||||
import cv2
|
||||
|
||||
upsampler = self._build_upsampler(scale, model_name, denoise_strength, face_enhance)
|
||||
|
||||
# Get source frame rate
|
||||
fps = self._get_video_fps(input_path)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
frames_dir = Path(tmpdir) / "frames"
|
||||
upscaled_dir = Path(tmpdir) / "upscaled"
|
||||
frames_dir.mkdir()
|
||||
upscaled_dir.mkdir()
|
||||
|
||||
# Extract frames
|
||||
self.run_command([
|
||||
"ffmpeg", "-y",
|
||||
"-i", str(input_path),
|
||||
str(frames_dir / "frame_%06d.png"),
|
||||
])
|
||||
|
||||
# Upscale each frame
|
||||
frame_files = sorted(frames_dir.glob("*.png"))
|
||||
total_frames = len(frame_files)
|
||||
|
||||
for frame_file in frame_files:
|
||||
img = cv2.imread(str(frame_file), cv2.IMREAD_UNCHANGED)
|
||||
output, _ = upsampler.enhance(img, outscale=scale)
|
||||
cv2.imwrite(str(upscaled_dir / frame_file.name), output)
|
||||
|
||||
# Reassemble with ffmpeg, copy audio from original
|
||||
reassemble_cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-framerate", str(fps),
|
||||
"-i", str(upscaled_dir / "frame_%06d.png"),
|
||||
"-i", str(input_path),
|
||||
"-map", "0:v",
|
||||
"-map", "1:a?",
|
||||
"-c:v", "libx264", "-crf", "18",
|
||||
"-c:a", "copy",
|
||||
"-pix_fmt", "yuv420p",
|
||||
str(output_path),
|
||||
]
|
||||
self.run_command(reassemble_cmd)
|
||||
|
||||
return {"total_frames": total_frames, "fps": fps}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _build_upsampler(
|
||||
self,
|
||||
scale: int,
|
||||
model_name: str,
|
||||
denoise_strength: float,
|
||||
face_enhance: bool,
|
||||
):
|
||||
"""Build and return a RealESRGANer instance."""
|
||||
import torch
|
||||
from basicsr.archs.rrdbnet_arch import RRDBNet
|
||||
from realesrgan import RealESRGANer
|
||||
|
||||
# Select architecture based on model
|
||||
if model_name == "RealESRGAN_x4plus_anime_6B":
|
||||
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
|
||||
else:
|
||||
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
|
||||
|
||||
# Resolve model path — realesrgan ships weights or downloads them
|
||||
model_url = f"https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/{model_name}.pth"
|
||||
if model_name == "RealESRGAN_x4plus_anime_6B":
|
||||
model_url = f"https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/{model_name}.pth"
|
||||
|
||||
half = torch.cuda.is_available()
|
||||
|
||||
upsampler = RealESRGANer(
|
||||
scale=4,
|
||||
model_path=model_url,
|
||||
model=model,
|
||||
dni_weight=denoise_strength,
|
||||
half=half,
|
||||
)
|
||||
|
||||
if face_enhance:
|
||||
from gfpgan import GFPGANer
|
||||
face_enhancer = GFPGANer(
|
||||
model_path="https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth",
|
||||
upscale=scale,
|
||||
arch="clean",
|
||||
channel_multiplier=2,
|
||||
bg_upsampler=upsampler,
|
||||
)
|
||||
# Monkey-patch so the caller can use the same interface
|
||||
original_enhance = upsampler.enhance
|
||||
|
||||
def enhance_with_face(img, outscale=scale):
|
||||
_, _, output = face_enhancer.enhance(
|
||||
img, has_aligned=False, only_center_face=False, paste_back=True,
|
||||
)
|
||||
return output, None
|
||||
|
||||
upsampler.enhance = enhance_with_face
|
||||
|
||||
return upsampler
|
||||
|
||||
def _get_video_fps(self, video_path: Path) -> float:
|
||||
"""Extract frame rate from video using ffprobe."""
|
||||
import json
|
||||
|
||||
if not shutil.which("ffprobe"):
|
||||
return 30.0 # safe default
|
||||
|
||||
try:
|
||||
proc = self.run_command([
|
||||
"ffprobe", "-v", "quiet",
|
||||
"-print_format", "json",
|
||||
"-show_streams",
|
||||
str(video_path),
|
||||
])
|
||||
probe = json.loads(proc.stdout)
|
||||
for stream in probe.get("streams", []):
|
||||
if stream.get("codec_type") == "video":
|
||||
r_frame_rate = stream.get("r_frame_rate", "30/1")
|
||||
num, den = r_frame_rate.split("/")
|
||||
return round(int(num) / int(den), 3)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return 30.0
|
||||
@@ -0,0 +1 @@
|
||||
"""Graphics tools for image, diagram, and animation generation."""
|
||||
@@ -0,0 +1,232 @@
|
||||
"""Code snippet renderer for overlay images.
|
||||
|
||||
Generates styled code screenshots using Pygments for syntax
|
||||
highlighting and Pillow for rendering. No external services required.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
ToolResult,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
# Theme presets mapping to Pygments styles and background colors
|
||||
THEMES = {
|
||||
"monokai": {
|
||||
"pygments_style": "monokai",
|
||||
"bg_color": "#272822",
|
||||
"text_color": "#f8f8f2",
|
||||
"border_color": "#3e3d32",
|
||||
},
|
||||
"github_dark": {
|
||||
"pygments_style": "github-dark",
|
||||
"bg_color": "#0d1117",
|
||||
"text_color": "#c9d1d9",
|
||||
"border_color": "#30363d",
|
||||
},
|
||||
"dracula": {
|
||||
"pygments_style": "dracula",
|
||||
"bg_color": "#282a36",
|
||||
"text_color": "#f8f8f2",
|
||||
"border_color": "#44475a",
|
||||
},
|
||||
"one_dark": {
|
||||
"pygments_style": "one-dark",
|
||||
"bg_color": "#282c34",
|
||||
"text_color": "#abb2bf",
|
||||
"border_color": "#3e4452",
|
||||
},
|
||||
"solarized_dark": {
|
||||
"pygments_style": "solarized-dark",
|
||||
"bg_color": "#002b36",
|
||||
"text_color": "#839496",
|
||||
"border_color": "#073642",
|
||||
},
|
||||
"light": {
|
||||
"pygments_style": "default",
|
||||
"bg_color": "#ffffff",
|
||||
"text_color": "#333333",
|
||||
"border_color": "#e1e4e8",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class CodeSnippet(BaseTool):
|
||||
name = "code_snippet"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.CORE
|
||||
capability = "graphics"
|
||||
provider = "pygments"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
|
||||
dependencies = ["python:pygments", "python:PIL"]
|
||||
install_instructions = "pip install Pygments Pillow"
|
||||
agent_skills = []
|
||||
|
||||
capabilities = [
|
||||
"render_code_image",
|
||||
"syntax_highlight",
|
||||
"themed_code_card",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["code"],
|
||||
"properties": {
|
||||
"code": {"type": "string"},
|
||||
"language": {"type": "string", "default": "python"},
|
||||
"theme": {
|
||||
"type": "string",
|
||||
"enum": list(THEMES.keys()),
|
||||
"default": "monokai",
|
||||
},
|
||||
"font_size": {"type": "integer", "default": 20},
|
||||
"padding": {"type": "integer", "default": 40},
|
||||
"border_radius": {"type": "integer", "default": 12},
|
||||
"line_numbers": {"type": "boolean", "default": True},
|
||||
"title": {"type": "string", "description": "Optional title bar text"},
|
||||
"output_path": {"type": "string"},
|
||||
"width": {"type": "integer", "description": "Force specific width"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(cpu_cores=1, ram_mb=256, vram_mb=0, disk_mb=50)
|
||||
idempotency_key_fields = ["code", "language", "theme", "font_size"]
|
||||
side_effects = ["writes image to output_path"]
|
||||
user_visible_verification = [
|
||||
"Verify code is readable and syntax highlighting is correct",
|
||||
]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
try:
|
||||
import pygments # noqa: F401
|
||||
from PIL import Image # noqa: F401
|
||||
return ToolStatus.AVAILABLE
|
||||
except ImportError:
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
try:
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
from pygments import highlight
|
||||
from pygments.lexers import get_lexer_by_name, guess_lexer
|
||||
from pygments.formatters import ImageFormatter
|
||||
except ImportError:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="Pygments and Pillow required. Run: pip install Pygments Pillow",
|
||||
)
|
||||
|
||||
start = time.time()
|
||||
|
||||
code = inputs["code"]
|
||||
language = inputs.get("language", "python")
|
||||
theme_name = inputs.get("theme", "monokai")
|
||||
font_size = inputs.get("font_size", 20)
|
||||
padding = inputs.get("padding", 40)
|
||||
line_numbers = inputs.get("line_numbers", True)
|
||||
title = inputs.get("title")
|
||||
output_path = Path(inputs.get("output_path", "code_snippet.png"))
|
||||
|
||||
theme = THEMES.get(theme_name, THEMES["monokai"])
|
||||
|
||||
try:
|
||||
lexer = get_lexer_by_name(language)
|
||||
except Exception:
|
||||
lexer = guess_lexer(code)
|
||||
|
||||
# Use Pygments ImageFormatter for rendering
|
||||
formatter = ImageFormatter(
|
||||
style=theme["pygments_style"],
|
||||
font_size=font_size,
|
||||
line_numbers=line_numbers,
|
||||
image_pad=padding,
|
||||
line_number_bg=theme["bg_color"],
|
||||
line_number_fg="#6272a4",
|
||||
)
|
||||
|
||||
# Render to bytes
|
||||
image_bytes = highlight(code, lexer, formatter)
|
||||
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(image_bytes)
|
||||
|
||||
# Add title bar if requested
|
||||
if title:
|
||||
self._add_title_bar(output_path, title, theme, font_size)
|
||||
|
||||
elapsed = time.time() - start
|
||||
img = Image.open(output_path)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"output": str(output_path),
|
||||
"language": language,
|
||||
"theme": theme_name,
|
||||
"width": img.width,
|
||||
"height": img.height,
|
||||
"line_count": code.count("\n") + 1,
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
duration_seconds=round(elapsed, 2),
|
||||
)
|
||||
|
||||
def _add_title_bar(
|
||||
self, image_path: Path, title: str, theme: dict, font_size: int
|
||||
) -> None:
|
||||
"""Add a title bar to the top of the code image."""
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
|
||||
img = Image.open(image_path)
|
||||
bar_height = font_size + 20
|
||||
|
||||
new_img = Image.new("RGB", (img.width, img.height + bar_height), theme["bg_color"])
|
||||
|
||||
# Draw title bar
|
||||
draw = ImageDraw.Draw(new_img)
|
||||
draw.rectangle(
|
||||
[(0, 0), (img.width, bar_height)],
|
||||
fill=theme["border_color"],
|
||||
)
|
||||
|
||||
# Draw window dots
|
||||
dot_y = bar_height // 2
|
||||
for i, color in enumerate(["#ff5f56", "#ffbd2e", "#27c93f"]):
|
||||
draw.ellipse(
|
||||
[(15 + i * 22, dot_y - 6), (15 + i * 22 + 12, dot_y + 6)],
|
||||
fill=color,
|
||||
)
|
||||
|
||||
# Draw title text
|
||||
try:
|
||||
font = ImageFont.truetype("arial.ttf", font_size - 4)
|
||||
except (IOError, OSError):
|
||||
font = ImageFont.load_default()
|
||||
|
||||
bbox = draw.textbbox((0, 0), title, font=font)
|
||||
text_width = bbox[2] - bbox[0]
|
||||
text_x = (img.width - text_width) // 2
|
||||
draw.text((text_x, 8), title, fill=theme["text_color"], font=font)
|
||||
|
||||
# Paste original image below title bar
|
||||
new_img.paste(img, (0, bar_height))
|
||||
new_img.save(image_path)
|
||||
|
||||
@staticmethod
|
||||
def list_themes() -> dict[str, str]:
|
||||
return {name: f"Background: {t['bg_color']}" for name, t in THEMES.items()}
|
||||
@@ -0,0 +1,351 @@
|
||||
"""Diagram generation tool using Mermaid CLI, Cairo/Pillow, or Graphviz.
|
||||
|
||||
Generates technical diagrams from text descriptions. Supports Mermaid
|
||||
syntax (flowcharts, sequence diagrams, etc.) and simple box/arrow
|
||||
diagrams via Pillow as fallback.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import shutil
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
ToolResult,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class DiagramGen(BaseTool):
|
||||
name = "diagram_gen"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.CORE
|
||||
capability = "graphics"
|
||||
provider = "mermaid"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
|
||||
dependencies = [] # checked dynamically
|
||||
install_instructions = (
|
||||
"For Mermaid diagrams:\n"
|
||||
" npm install -g @mermaid-js/mermaid-cli\n"
|
||||
"For Pillow-based diagrams (fallback):\n"
|
||||
" pip install Pillow"
|
||||
)
|
||||
agent_skills = ["beautiful-mermaid", "d3-viz"]
|
||||
|
||||
capabilities = [
|
||||
"generate_mermaid",
|
||||
"generate_flowchart",
|
||||
"generate_box_diagram",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["diagram_type"],
|
||||
"properties": {
|
||||
"diagram_type": {
|
||||
"type": "string",
|
||||
"enum": ["mermaid", "flowchart", "boxes"],
|
||||
},
|
||||
"definition": {
|
||||
"type": "string",
|
||||
"description": "Mermaid syntax or diagram description",
|
||||
},
|
||||
"boxes": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"label": {"type": "string"},
|
||||
"color": {"type": "string"},
|
||||
},
|
||||
},
|
||||
"description": "Box definitions for box diagram type",
|
||||
},
|
||||
"connections": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"from": {"type": "integer"},
|
||||
"to": {"type": "integer"},
|
||||
"label": {"type": "string"},
|
||||
},
|
||||
},
|
||||
},
|
||||
"title": {"type": "string"},
|
||||
"theme": {
|
||||
"type": "string",
|
||||
"enum": ["dark", "light", "neutral"],
|
||||
"default": "dark",
|
||||
},
|
||||
"width": {"type": "integer", "default": 1200},
|
||||
"height": {"type": "integer", "default": 800},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(cpu_cores=1, ram_mb=256, vram_mb=0, disk_mb=50)
|
||||
idempotency_key_fields = ["diagram_type", "definition", "boxes"]
|
||||
side_effects = ["writes diagram image to output_path"]
|
||||
user_visible_verification = [
|
||||
"Verify diagram accurately represents the described structure",
|
||||
]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if self._has_mermaid() or self._has_pillow():
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def _has_mermaid(self) -> bool:
|
||||
return shutil.which("mmdc") is not None
|
||||
|
||||
def _has_pillow(self) -> bool:
|
||||
try:
|
||||
from PIL import Image # noqa: F401
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
diagram_type = inputs["diagram_type"]
|
||||
start = time.time()
|
||||
|
||||
try:
|
||||
if diagram_type == "mermaid":
|
||||
result = self._render_mermaid(inputs)
|
||||
elif diagram_type in ("flowchart", "boxes"):
|
||||
result = self._render_boxes(inputs)
|
||||
else:
|
||||
return ToolResult(success=False, error=f"Unknown diagram type: {diagram_type}")
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Diagram generation failed: {e}")
|
||||
|
||||
result.duration_seconds = round(time.time() - start, 2)
|
||||
return result
|
||||
|
||||
def _render_mermaid(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
definition = inputs.get("definition", "")
|
||||
if not definition:
|
||||
return ToolResult(success=False, error="Mermaid definition required")
|
||||
|
||||
output_path = Path(inputs.get("output_path", "diagram.png"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
theme = inputs.get("theme", "dark")
|
||||
|
||||
if self._has_mermaid():
|
||||
# Write temp mermaid file
|
||||
temp_mmd = output_path.with_suffix(".mmd")
|
||||
temp_mmd.write_text(definition, encoding="utf-8")
|
||||
|
||||
mermaid_config = {"theme": theme}
|
||||
config_path = output_path.with_suffix(".mermaid.json")
|
||||
config_path.write_text(json.dumps(mermaid_config), encoding="utf-8")
|
||||
|
||||
cmd = [
|
||||
"mmdc",
|
||||
"-i", str(temp_mmd),
|
||||
"-o", str(output_path),
|
||||
"-c", str(config_path),
|
||||
"-b", "transparent",
|
||||
"-w", str(inputs.get("width", 1200)),
|
||||
]
|
||||
|
||||
try:
|
||||
self.run_command(cmd, timeout=30)
|
||||
finally:
|
||||
temp_mmd.unlink(missing_ok=True)
|
||||
config_path.unlink(missing_ok=True)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"method": "mermaid-cli",
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
)
|
||||
else:
|
||||
# Fallback: render mermaid text as a styled text card
|
||||
return self._render_text_card(definition, inputs)
|
||||
|
||||
def _render_boxes(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
"""Render a box-and-arrow diagram using Pillow."""
|
||||
if not self._has_pillow():
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="Pillow required for box diagrams. Run: pip install Pillow",
|
||||
)
|
||||
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
|
||||
boxes = inputs.get("boxes", [])
|
||||
connections = inputs.get("connections", [])
|
||||
title = inputs.get("title", "")
|
||||
theme = inputs.get("theme", "dark")
|
||||
width = inputs.get("width", 1200)
|
||||
height = inputs.get("height", 800)
|
||||
output_path = Path(inputs.get("output_path", "diagram.png"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Theme colors
|
||||
if theme == "dark":
|
||||
bg, text_color, box_default, line_color = "#1e1e2e", "#cdd6f4", "#45475a", "#89b4fa"
|
||||
elif theme == "light":
|
||||
bg, text_color, box_default, line_color = "#ffffff", "#333333", "#e1e4e8", "#0366d6"
|
||||
else:
|
||||
bg, text_color, box_default, line_color = "#2d2d2d", "#d4d4d4", "#404040", "#569cd6"
|
||||
|
||||
img = Image.new("RGB", (width, height), bg)
|
||||
draw = ImageDraw.Draw(img)
|
||||
|
||||
try:
|
||||
font = ImageFont.truetype("arial.ttf", 18)
|
||||
title_font = ImageFont.truetype("arial.ttf", 24)
|
||||
except (IOError, OSError):
|
||||
font = ImageFont.load_default()
|
||||
title_font = font
|
||||
|
||||
# Draw title
|
||||
y_offset = 20
|
||||
if title:
|
||||
bbox = draw.textbbox((0, 0), title, font=title_font)
|
||||
tw = bbox[2] - bbox[0]
|
||||
draw.text(((width - tw) // 2, y_offset), title, fill=text_color, font=title_font)
|
||||
y_offset += 50
|
||||
|
||||
# Layout boxes in a grid
|
||||
if not boxes:
|
||||
boxes = [{"label": "Empty"}]
|
||||
|
||||
cols = min(len(boxes), 4)
|
||||
rows = (len(boxes) + cols - 1) // cols
|
||||
box_w = min(200, (width - 80) // cols - 20)
|
||||
box_h = 60
|
||||
x_gap = (width - cols * box_w) // (cols + 1)
|
||||
y_gap = max(40, (height - y_offset - rows * box_h) // (rows + 1))
|
||||
|
||||
box_positions = []
|
||||
for i, box in enumerate(boxes):
|
||||
col = i % cols
|
||||
row = i // cols
|
||||
x = x_gap + col * (box_w + x_gap)
|
||||
y = y_offset + y_gap + row * (box_h + y_gap)
|
||||
|
||||
fill = box.get("color", box_default)
|
||||
draw.rounded_rectangle(
|
||||
[(x, y), (x + box_w, y + box_h)],
|
||||
radius=8,
|
||||
fill=fill,
|
||||
outline=line_color,
|
||||
width=2,
|
||||
)
|
||||
|
||||
label = box.get("label", f"Box {i}")
|
||||
bbox = draw.textbbox((0, 0), label, font=font)
|
||||
lw = bbox[2] - bbox[0]
|
||||
lh = bbox[3] - bbox[1]
|
||||
draw.text(
|
||||
(x + (box_w - lw) // 2, y + (box_h - lh) // 2),
|
||||
label, fill=text_color, font=font,
|
||||
)
|
||||
|
||||
box_positions.append((x, y, x + box_w, y + box_h))
|
||||
|
||||
# Draw connections
|
||||
for conn in connections:
|
||||
fi = conn.get("from", 0)
|
||||
ti = conn.get("to", 0)
|
||||
if fi >= len(box_positions) or ti >= len(box_positions):
|
||||
continue
|
||||
|
||||
fx1, fy1, fx2, fy2 = box_positions[fi]
|
||||
tx1, ty1, tx2, ty2 = box_positions[ti]
|
||||
|
||||
start_x = (fx1 + fx2) // 2
|
||||
start_y = fy2
|
||||
end_x = (tx1 + tx2) // 2
|
||||
end_y = ty1
|
||||
|
||||
draw.line([(start_x, start_y), (end_x, end_y)], fill=line_color, width=2)
|
||||
|
||||
# Arrow head
|
||||
arrow_size = 8
|
||||
draw.polygon(
|
||||
[(end_x, end_y), (end_x - arrow_size, end_y - arrow_size * 2), (end_x + arrow_size, end_y - arrow_size * 2)],
|
||||
fill=line_color,
|
||||
)
|
||||
|
||||
# Connection label
|
||||
conn_label = conn.get("label")
|
||||
if conn_label:
|
||||
mid_x = (start_x + end_x) // 2
|
||||
mid_y = (start_y + end_y) // 2
|
||||
draw.text((mid_x + 5, mid_y - 10), conn_label, fill=text_color, font=font)
|
||||
|
||||
img.save(output_path)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"method": "pillow",
|
||||
"output": str(output_path),
|
||||
"box_count": len(boxes),
|
||||
"connection_count": len(connections),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
)
|
||||
|
||||
def _render_text_card(self, text: str, inputs: dict[str, Any]) -> ToolResult:
|
||||
"""Fallback: render text as a styled card image."""
|
||||
if not self._has_pillow():
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="Pillow required. Run: pip install Pillow",
|
||||
)
|
||||
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
|
||||
output_path = Path(inputs.get("output_path", "diagram.png"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
width = inputs.get("width", 800)
|
||||
|
||||
try:
|
||||
font = ImageFont.truetype("consola.ttf", 16)
|
||||
except (IOError, OSError):
|
||||
font = ImageFont.load_default()
|
||||
|
||||
# Calculate needed height
|
||||
lines = text.split("\n")
|
||||
line_height = 22
|
||||
height = max(200, len(lines) * line_height + 80)
|
||||
|
||||
img = Image.new("RGB", (width, height), "#1e1e2e")
|
||||
draw = ImageDraw.Draw(img)
|
||||
|
||||
y = 40
|
||||
for line in lines:
|
||||
draw.text((40, y), line, fill="#cdd6f4", font=font)
|
||||
y += line_height
|
||||
|
||||
img.save(output_path)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"method": "text_card",
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
)
|
||||
@@ -0,0 +1,164 @@
|
||||
"""FLUX image generation via fal.ai API."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class FluxImage(BaseTool):
|
||||
name = "flux_image"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "image_generation"
|
||||
provider = "flux"
|
||||
stability = ToolStability.BETA
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.SEEDED
|
||||
runtime = ToolRuntime.API
|
||||
|
||||
dependencies = [] # checked dynamically via env var
|
||||
install_instructions = (
|
||||
"Set FAL_KEY to your fal.ai API key.\n"
|
||||
" Get one at https://fal.ai/dashboard/keys"
|
||||
)
|
||||
agent_skills = ["flux-best-practices", "bfl-api"]
|
||||
|
||||
capabilities = ["generate_image", "generate_illustration", "text_to_image"]
|
||||
supports = {
|
||||
"negative_prompt": True,
|
||||
"seed": True,
|
||||
"custom_size": True,
|
||||
}
|
||||
best_for = [
|
||||
"photorealistic images",
|
||||
"general-purpose image generation",
|
||||
"high quality at low cost (~$0.03/image)",
|
||||
]
|
||||
not_good_for = ["text rendering in images", "offline generation"]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {"type": "string"},
|
||||
"negative_prompt": {"type": "string", "default": ""},
|
||||
"width": {"type": "integer", "default": 1024},
|
||||
"height": {"type": "integer", "default": 1024},
|
||||
"model": {
|
||||
"type": "string",
|
||||
"enum": ["flux-pro/v1.1", "flux/dev", "flux-pro"],
|
||||
"default": "flux-pro/v1.1",
|
||||
},
|
||||
"seed": {"type": "integer"},
|
||||
"num_inference_steps": {"type": "integer"},
|
||||
"guidance_scale": {"type": "number"},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=1, ram_mb=512, vram_mb=0, disk_mb=100, network_required=True
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=2, retryable_errors=["rate_limit", "timeout"])
|
||||
idempotency_key_fields = ["prompt", "width", "height", "seed", "model"]
|
||||
side_effects = ["writes image file to output_path", "calls fal.ai API"]
|
||||
user_visible_verification = ["Inspect generated image for relevance and quality"]
|
||||
|
||||
def _get_api_key(self) -> str | None:
|
||||
return os.environ.get("FAL_KEY") or os.environ.get("FAL_AI_API_KEY")
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if self._get_api_key():
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
model = inputs.get("model", "flux-pro/v1.1")
|
||||
if "pro" in model:
|
||||
return 0.05
|
||||
return 0.03 # dev tier
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
api_key = self._get_api_key()
|
||||
if not api_key:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="No fal.ai API key found. " + self.install_instructions,
|
||||
)
|
||||
|
||||
import requests
|
||||
|
||||
start = time.time()
|
||||
model = inputs.get("model", "flux-pro/v1.1")
|
||||
prompt = inputs["prompt"]
|
||||
width = inputs.get("width", 1024)
|
||||
height = inputs.get("height", 1024)
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"prompt": prompt,
|
||||
"image_size": {"width": width, "height": height},
|
||||
}
|
||||
if inputs.get("seed") is not None:
|
||||
payload["seed"] = inputs["seed"]
|
||||
if inputs.get("num_inference_steps"):
|
||||
payload["num_inference_steps"] = inputs["num_inference_steps"]
|
||||
if inputs.get("guidance_scale"):
|
||||
payload["guidance_scale"] = inputs["guidance_scale"]
|
||||
if inputs.get("negative_prompt"):
|
||||
payload["negative_prompt"] = inputs["negative_prompt"]
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
f"https://fal.run/fal-ai/{model}",
|
||||
headers={
|
||||
"Authorization": f"Key {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
json=payload,
|
||||
timeout=120,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
image_url = data["images"][0]["url"]
|
||||
image_response = requests.get(image_url, timeout=60)
|
||||
image_response.raise_for_status()
|
||||
|
||||
output_path = Path(inputs.get("output_path", "generated_image.png"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(image_response.content)
|
||||
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"FLUX generation failed: {e}")
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "flux",
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"output": str(output_path),
|
||||
"seed": data.get("seed"),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
cost_usd=self.estimate_cost(inputs),
|
||||
duration_seconds=round(time.time() - start, 2),
|
||||
seed=data.get("seed"),
|
||||
model=f"fal-ai/{model}",
|
||||
)
|
||||
@@ -0,0 +1,269 @@
|
||||
"""Image generation tool for diagrams, overlays, and illustrations.
|
||||
|
||||
.. deprecated::
|
||||
Use ``image_selector`` instead. This monolithic tool has been replaced by
|
||||
the selector/provider pattern: ``image_selector`` routes to per-provider
|
||||
tools (flux_image, openai_image, recraft_image, local_diffusion,
|
||||
pexels_image, pixabay_image). This file is kept for backwards
|
||||
compatibility and will be removed in a future release.
|
||||
|
||||
Supports cloud API providers (FLUX via fal.ai/Replicate, OpenAI DALL-E)
|
||||
and local Stable Diffusion via diffusers. Reports unavailable with
|
||||
install instructions when no provider is configured.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class ImageGen(BaseTool):
|
||||
name = "image_gen"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.CORE
|
||||
capability = "image_generation"
|
||||
provider = "multi"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.SEEDED
|
||||
runtime = ToolRuntime.HYBRID # API (DALL-E/FLUX) or local (diffusers)
|
||||
|
||||
dependencies = [] # checked dynamically based on provider
|
||||
install_instructions = (
|
||||
"Set one of these environment variables:\n"
|
||||
" OPENAI_API_KEY — for DALL-E 3\n"
|
||||
" FAL_KEY — for FLUX via fal.ai\n"
|
||||
"Or install diffusers for local generation:\n"
|
||||
" pip install diffusers transformers accelerate torch"
|
||||
)
|
||||
agent_skills = ["flux-best-practices", "bfl-api"]
|
||||
|
||||
capabilities = [
|
||||
"generate_image",
|
||||
"generate_diagram_overlay",
|
||||
"generate_illustration",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {"type": "string"},
|
||||
"negative_prompt": {"type": "string", "default": ""},
|
||||
"width": {"type": "integer", "default": 1024},
|
||||
"height": {"type": "integer", "default": 1024},
|
||||
"provider": {
|
||||
"type": "string",
|
||||
"enum": ["openai", "flux", "local"],
|
||||
"description": "Auto-detected if not specified",
|
||||
},
|
||||
"model": {"type": "string"},
|
||||
"seed": {"type": "integer"},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=1, ram_mb=512, vram_mb=0, disk_mb=100, network_required=True
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=2, retryable_errors=["rate_limit", "timeout"])
|
||||
idempotency_key_fields = ["prompt", "width", "height", "seed"]
|
||||
side_effects = ["writes image file to output_path", "calls external API"]
|
||||
user_visible_verification = [
|
||||
"Inspect generated image for relevance and quality",
|
||||
]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
provider = self._detect_provider()
|
||||
if provider:
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def _detect_provider(self) -> Optional[str]:
|
||||
if os.environ.get("OPENAI_API_KEY"):
|
||||
return "openai"
|
||||
if os.environ.get("FAL_KEY") or os.environ.get("FAL_AI_API_KEY"):
|
||||
return "flux"
|
||||
try:
|
||||
import diffusers # noqa: F401
|
||||
return "local"
|
||||
except ImportError:
|
||||
pass
|
||||
return None
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
provider = inputs.get("provider") or self._detect_provider()
|
||||
if provider == "openai":
|
||||
return 0.04 # DALL-E 3 standard
|
||||
if provider == "flux":
|
||||
return 0.03
|
||||
return 0.0 # local
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
provider = inputs.get("provider") or self._detect_provider()
|
||||
if not provider:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="No image generation provider available. " + self.install_instructions,
|
||||
)
|
||||
|
||||
start = time.time()
|
||||
|
||||
try:
|
||||
if provider == "openai":
|
||||
result = self._generate_openai(inputs)
|
||||
elif provider == "flux":
|
||||
result = self._generate_flux(inputs)
|
||||
elif provider == "local":
|
||||
result = self._generate_local(inputs)
|
||||
else:
|
||||
return ToolResult(success=False, error=f"Unknown provider: {provider}")
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Generation failed: {e}")
|
||||
|
||||
result.duration_seconds = round(time.time() - start, 2)
|
||||
result.cost_usd = self.estimate_cost(inputs)
|
||||
return result
|
||||
|
||||
def _generate_openai(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
from openai import OpenAI
|
||||
import base64
|
||||
|
||||
client = OpenAI()
|
||||
prompt = inputs["prompt"]
|
||||
size = f"{inputs.get('width', 1024)}x{inputs.get('height', 1024)}"
|
||||
model = inputs.get("model", "dall-e-3")
|
||||
|
||||
response = client.images.generate(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
size=size,
|
||||
n=1,
|
||||
response_format="b64_json",
|
||||
)
|
||||
|
||||
image_data = base64.b64decode(response.data[0].b64_json)
|
||||
output_path = Path(inputs.get("output_path", "generated_image.png"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(image_data)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "openai",
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
model=model,
|
||||
)
|
||||
|
||||
def _generate_flux(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
import requests
|
||||
|
||||
api_key = os.environ.get("FAL_KEY") or os.environ["FAL_AI_API_KEY"]
|
||||
prompt = inputs["prompt"]
|
||||
width = inputs.get("width", 1024)
|
||||
height = inputs.get("height", 1024)
|
||||
seed = inputs.get("seed")
|
||||
|
||||
payload = {
|
||||
"prompt": prompt,
|
||||
"image_size": {"width": width, "height": height},
|
||||
}
|
||||
if seed is not None:
|
||||
payload["seed"] = seed
|
||||
|
||||
response = requests.post(
|
||||
"https://fal.run/fal-ai/flux/dev",
|
||||
headers={"Authorization": f"Key {api_key}", "Content-Type": "application/json"},
|
||||
json=payload,
|
||||
timeout=120,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
image_url = data["images"][0]["url"]
|
||||
image_response = requests.get(image_url, timeout=60)
|
||||
image_response.raise_for_status()
|
||||
|
||||
output_path = Path(inputs.get("output_path", "generated_image.png"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(image_response.content)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "flux",
|
||||
"prompt": prompt,
|
||||
"output": str(output_path),
|
||||
"seed": data.get("seed"),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
seed=data.get("seed"),
|
||||
model="flux-dev",
|
||||
)
|
||||
|
||||
def _generate_local(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
import torch
|
||||
from diffusers import StableDiffusionPipeline
|
||||
|
||||
prompt = inputs["prompt"]
|
||||
negative = inputs.get("negative_prompt", "")
|
||||
width = inputs.get("width", 512)
|
||||
height = inputs.get("height", 512)
|
||||
seed = inputs.get("seed")
|
||||
model_id = inputs.get("model", "stabilityai/stable-diffusion-2-1-base")
|
||||
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
dtype = torch.float16 if device == "cuda" else torch.float32
|
||||
|
||||
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=dtype)
|
||||
pipe = pipe.to(device)
|
||||
|
||||
generator = None
|
||||
if seed is not None:
|
||||
generator = torch.Generator(device=device).manual_seed(seed)
|
||||
|
||||
image = pipe(
|
||||
prompt,
|
||||
negative_prompt=negative,
|
||||
width=width,
|
||||
height=height,
|
||||
generator=generator,
|
||||
).images[0]
|
||||
|
||||
output_path = Path(inputs.get("output_path", "generated_image.png"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
image.save(str(output_path))
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "local",
|
||||
"model": model_id,
|
||||
"prompt": prompt,
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
seed=seed,
|
||||
model=model_id,
|
||||
)
|
||||
@@ -0,0 +1,142 @@
|
||||
"""Capability-level image selector that routes between generation and stock providers.
|
||||
|
||||
Provider discovery is automatic — any BaseTool with capability="image_generation"
|
||||
is picked up from the registry. Adding a new image provider requires only creating
|
||||
the tool file in tools/graphics/; no changes to this selector are needed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import BaseTool, ToolResult, ToolRuntime, ToolStability, ToolStatus, ToolTier
|
||||
|
||||
|
||||
class ImageSelector(BaseTool):
|
||||
name = "image_selector"
|
||||
version = "0.2.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "image_generation"
|
||||
provider = "selector"
|
||||
stability = ToolStability.BETA
|
||||
runtime = ToolRuntime.HYBRID
|
||||
agent_skills = ["flux-best-practices", "bfl-api"]
|
||||
|
||||
capabilities = [
|
||||
"generate_image", "search_image", "download_image",
|
||||
"provider_selection", "text_to_image", "stock_image",
|
||||
]
|
||||
supports = {
|
||||
"user_preference_routing": True,
|
||||
"offline_fallback": True,
|
||||
"stock_fallback": True,
|
||||
}
|
||||
best_for = [
|
||||
"preflight routing — pick the best image provider for the task",
|
||||
"switching between generated and stock images",
|
||||
"automatic fallback when preferred provider is unavailable",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {
|
||||
"type": "string",
|
||||
"description": "Image description (used as prompt for generation or query for stock)",
|
||||
},
|
||||
"negative_prompt": {
|
||||
"type": "string",
|
||||
"description": "What to avoid in the generated image. Passed to providers that support it.",
|
||||
},
|
||||
"width": {"type": "integer", "description": "Image width in pixels"},
|
||||
"height": {"type": "integer", "description": "Image height in pixels"},
|
||||
"seed": {"type": "integer", "description": "Random seed for reproducibility (generation providers only)"},
|
||||
"preferred_provider": {
|
||||
"type": "string",
|
||||
"description": "Provider name or 'auto'. Valid values are discovered at runtime from the registry.",
|
||||
"default": "auto",
|
||||
},
|
||||
"allowed_providers": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
def _providers(self) -> list[BaseTool]:
|
||||
"""Auto-discover image generation providers from the registry."""
|
||||
from tools.tool_registry import registry
|
||||
registry.ensure_discovered()
|
||||
return [t for t in registry.get_by_capability("image_generation")
|
||||
if t.name != self.name]
|
||||
|
||||
@property
|
||||
def fallback_tools(self) -> list[str]:
|
||||
"""Dynamically built from discovered providers."""
|
||||
return [t.name for t in self._providers()]
|
||||
|
||||
@property
|
||||
def provider_matrix(self) -> dict[str, dict[str, str]]:
|
||||
"""Built at runtime from each provider's best_for field."""
|
||||
matrix = {}
|
||||
for tool in self._providers():
|
||||
strength = ", ".join(tool.best_for) if tool.best_for else tool.name
|
||||
matrix[tool.provider] = {"tool": tool.name, "strength": strength}
|
||||
return matrix
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if any(tool.get_status() == ToolStatus.AVAILABLE for tool in self._providers()):
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
tool = self._select_tool(inputs)
|
||||
return tool.estimate_cost(inputs) if tool else 0.0
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
tool = self._select_tool(inputs)
|
||||
if tool is None:
|
||||
return ToolResult(success=False, error="No image provider available.")
|
||||
|
||||
# Adapt input keys: stock tools use 'query' while generators use 'prompt'
|
||||
adapted = dict(inputs)
|
||||
if hasattr(tool, 'input_schema'):
|
||||
props = tool.input_schema.get("properties", {})
|
||||
if "query" in props and "query" not in adapted:
|
||||
adapted["query"] = adapted.get("prompt", "")
|
||||
|
||||
# Strip selector-only keys that downstream tools don't understand
|
||||
adapted.pop("preferred_provider", None)
|
||||
adapted.pop("allowed_providers", None)
|
||||
|
||||
# Pass through generation params only to tools that accept them
|
||||
if hasattr(tool, 'input_schema'):
|
||||
props = tool.input_schema.get("properties", {})
|
||||
for passthrough_key in ("negative_prompt", "width", "height", "seed"):
|
||||
if passthrough_key in adapted and passthrough_key not in props:
|
||||
adapted.pop(passthrough_key)
|
||||
|
||||
result = tool.execute(adapted)
|
||||
if result.success:
|
||||
result.data.setdefault("selected_tool", tool.name)
|
||||
return result
|
||||
|
||||
def _select_tool(self, inputs: dict[str, Any]) -> BaseTool | None:
|
||||
preferred = inputs.get("preferred_provider", "auto")
|
||||
allowed = set(inputs.get("allowed_providers") or [])
|
||||
candidates = self._providers()
|
||||
if allowed:
|
||||
candidates = [tool for tool in candidates if tool.provider in allowed]
|
||||
|
||||
if preferred != "auto":
|
||||
ordered = [tool for tool in candidates if tool.provider == preferred]
|
||||
ordered.extend([tool for tool in candidates if tool.provider != preferred])
|
||||
else:
|
||||
ordered = candidates
|
||||
|
||||
for tool in ordered:
|
||||
if tool.get_status() == ToolStatus.AVAILABLE:
|
||||
return tool
|
||||
return None
|
||||
@@ -0,0 +1,159 @@
|
||||
"""Local Stable Diffusion image generation via diffusers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class LocalDiffusion(BaseTool):
|
||||
name = "local_diffusion"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "image_generation"
|
||||
provider = "local_diffusion"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.SEEDED
|
||||
runtime = ToolRuntime.LOCAL_GPU
|
||||
|
||||
dependencies = [] # checked dynamically
|
||||
install_instructions = (
|
||||
"Install diffusers for local Stable Diffusion:\n"
|
||||
" pip install diffusers transformers accelerate torch"
|
||||
)
|
||||
agent_skills = []
|
||||
|
||||
capabilities = ["generate_image", "generate_illustration", "text_to_image"]
|
||||
supports = {
|
||||
"negative_prompt": True,
|
||||
"seed": True,
|
||||
"offline": True,
|
||||
"custom_size": True,
|
||||
}
|
||||
best_for = [
|
||||
"offline/air-gapped generation",
|
||||
"free image generation (no API cost)",
|
||||
"privacy-sensitive workflows",
|
||||
]
|
||||
not_good_for = [
|
||||
"CPU-only machines (very slow)",
|
||||
"highest quality output (API models are better)",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {"type": "string"},
|
||||
"negative_prompt": {"type": "string", "default": ""},
|
||||
"width": {"type": "integer", "default": 512},
|
||||
"height": {"type": "integer", "default": 512},
|
||||
"model": {
|
||||
"type": "string",
|
||||
"default": "stabilityai/stable-diffusion-2-1-base",
|
||||
},
|
||||
"seed": {"type": "integer"},
|
||||
"num_inference_steps": {"type": "integer", "default": 30},
|
||||
"guidance_scale": {"type": "number", "default": 7.5},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=2, ram_mb=8000, vram_mb=4000, disk_mb=5000, network_required=False
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=1)
|
||||
idempotency_key_fields = ["prompt", "width", "height", "seed", "model"]
|
||||
side_effects = ["writes image file to output_path", "may download model weights on first run"]
|
||||
user_visible_verification = ["Inspect generated image for relevance and quality"]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
try:
|
||||
import diffusers # noqa: F401
|
||||
return ToolStatus.AVAILABLE
|
||||
except ImportError:
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
return 0.0
|
||||
|
||||
def estimate_runtime(self, inputs: dict[str, Any]) -> float:
|
||||
return 30.0 # ~30s on a mid-range GPU
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
if self.get_status() != ToolStatus.AVAILABLE:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="diffusers not installed. " + self.install_instructions,
|
||||
)
|
||||
|
||||
import torch
|
||||
from diffusers import StableDiffusionPipeline
|
||||
|
||||
start = time.time()
|
||||
prompt = inputs["prompt"]
|
||||
negative = inputs.get("negative_prompt", "")
|
||||
width = inputs.get("width", 512)
|
||||
height = inputs.get("height", 512)
|
||||
seed = inputs.get("seed")
|
||||
model_id = inputs.get("model", "stabilityai/stable-diffusion-2-1-base")
|
||||
steps = inputs.get("num_inference_steps", 30)
|
||||
guidance = inputs.get("guidance_scale", 7.5)
|
||||
|
||||
try:
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
dtype = torch.float16 if device == "cuda" else torch.float32
|
||||
|
||||
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=dtype)
|
||||
pipe = pipe.to(device)
|
||||
|
||||
generator = None
|
||||
if seed is not None:
|
||||
generator = torch.Generator(device=device).manual_seed(seed)
|
||||
|
||||
image = pipe(
|
||||
prompt,
|
||||
negative_prompt=negative,
|
||||
width=width,
|
||||
height=height,
|
||||
num_inference_steps=steps,
|
||||
guidance_scale=guidance,
|
||||
generator=generator,
|
||||
).images[0]
|
||||
|
||||
output_path = Path(inputs.get("output_path", "generated_image.png"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
image.save(str(output_path))
|
||||
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Local diffusion generation failed: {e}")
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "local_diffusion",
|
||||
"model": model_id,
|
||||
"prompt": prompt,
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
cost_usd=0.0,
|
||||
duration_seconds=round(time.time() - start, 2),
|
||||
seed=seed,
|
||||
model=model_id,
|
||||
)
|
||||
@@ -0,0 +1,373 @@
|
||||
"""Mathematical animation tool via ManimCE.
|
||||
|
||||
Generates animated math/science/explainer videos from Python scene code
|
||||
using the Manim Community Edition engine. Free, local, no API key required.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
# Quality presets mapping to Manim CLI flags
|
||||
QUALITY_PRESETS = {
|
||||
"low": {"flag": "-ql", "resolution": "854x480", "fps": 15},
|
||||
"medium": {"flag": "-qm", "resolution": "1280x720", "fps": 30},
|
||||
"high": {"flag": "-qh", "resolution": "1920x1080", "fps": 60},
|
||||
"4k": {"flag": "-qk", "resolution": "3840x2160", "fps": 60},
|
||||
"preview": {"flag": "-ql --format gif", "resolution": "854x480", "fps": 15},
|
||||
}
|
||||
|
||||
|
||||
class MathAnimate(BaseTool):
|
||||
name = "math_animate"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "graphics"
|
||||
provider = "manim"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
runtime = ToolRuntime.LOCAL
|
||||
|
||||
dependencies = ["cmd:manim"]
|
||||
install_instructions = (
|
||||
"Install ManimCE:\n"
|
||||
" pip install manim\n"
|
||||
" manim checkhealth\n"
|
||||
"Requires: Python 3.8+, FFmpeg, LaTeX (optional, for math formulas)\n"
|
||||
" Windows: choco install miktex ffmpeg\n"
|
||||
" macOS: brew install mactex ffmpeg\n"
|
||||
" Linux: sudo apt install texlive-full ffmpeg"
|
||||
)
|
||||
agent_skills = ["manimce-best-practices", "manim-composer"]
|
||||
|
||||
capabilities = [
|
||||
"render_scene",
|
||||
"render_from_code",
|
||||
"render_from_template",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["scene_code"],
|
||||
"properties": {
|
||||
"scene_code": {
|
||||
"type": "string",
|
||||
"description": (
|
||||
"Python code defining a Manim scene. Must contain a class "
|
||||
"inheriting from Scene with a construct() method. "
|
||||
"Import 'from manim import *' is auto-added if missing."
|
||||
),
|
||||
},
|
||||
"scene_name": {
|
||||
"type": "string",
|
||||
"description": "Name of the Scene class to render. Auto-detected if only one scene.",
|
||||
},
|
||||
"quality": {
|
||||
"type": "string",
|
||||
"enum": list(QUALITY_PRESETS.keys()),
|
||||
"default": "medium",
|
||||
"description": "Render quality preset",
|
||||
},
|
||||
"format": {
|
||||
"type": "string",
|
||||
"enum": ["mp4", "gif", "png", "webm"],
|
||||
"default": "mp4",
|
||||
},
|
||||
"output_path": {"type": "string"},
|
||||
"transparent": {
|
||||
"type": "boolean",
|
||||
"default": False,
|
||||
"description": "Render with transparent background (PNG sequence or WebM)",
|
||||
},
|
||||
"background_color": {
|
||||
"type": "string",
|
||||
"description": "Background hex color (e.g., '#1a1a2e'). Default: Manim default (black).",
|
||||
},
|
||||
"extra_args": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
"description": "Additional Manim CLI arguments",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=2, ram_mb=1024, vram_mb=0, disk_mb=500, network_required=False
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=1, retryable_errors=["timeout"])
|
||||
idempotency_key_fields = ["scene_code", "scene_name", "quality"]
|
||||
side_effects = ["writes video/image file to output_path", "creates temp files"]
|
||||
user_visible_verification = [
|
||||
"Watch the animation for correctness and visual quality",
|
||||
"Verify math formulas render correctly (requires LaTeX)",
|
||||
]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if shutil.which("manim"):
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
return 0.0 # local, free
|
||||
|
||||
def estimate_runtime(self, inputs: dict[str, Any]) -> float:
|
||||
quality = inputs.get("quality", "medium")
|
||||
# Rough estimates based on scene complexity (assuming ~10s scene)
|
||||
estimates = {
|
||||
"low": 5.0,
|
||||
"medium": 15.0,
|
||||
"high": 45.0,
|
||||
"4k": 120.0,
|
||||
"preview": 3.0,
|
||||
}
|
||||
return estimates.get(quality, 15.0)
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
if not shutil.which("manim"):
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="Manim not found. " + self.install_instructions,
|
||||
)
|
||||
|
||||
start = time.time()
|
||||
|
||||
try:
|
||||
result = self._render(inputs)
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Manim render failed: {e}")
|
||||
|
||||
result.duration_seconds = round(time.time() - start, 2)
|
||||
return result
|
||||
|
||||
def _render(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
scene_code = inputs["scene_code"]
|
||||
scene_name = inputs.get("scene_name")
|
||||
quality = inputs.get("quality", "medium")
|
||||
output_format = inputs.get("format", "mp4")
|
||||
output_path = inputs.get("output_path")
|
||||
transparent = inputs.get("transparent", False)
|
||||
bg_color = inputs.get("background_color")
|
||||
extra_args = inputs.get("extra_args", [])
|
||||
|
||||
# Ensure import statement
|
||||
if "from manim import" not in scene_code:
|
||||
scene_code = "from manim import *\n\n" + scene_code
|
||||
|
||||
# Auto-detect scene name if not provided
|
||||
if not scene_name:
|
||||
scene_name = self._detect_scene_name(scene_code)
|
||||
if not scene_name:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="Could not detect Scene class name. Provide scene_name explicitly.",
|
||||
)
|
||||
|
||||
# Write scene code to temp file
|
||||
work_dir = Path(tempfile.mkdtemp(prefix="manim_"))
|
||||
scene_file = work_dir / "scene.py"
|
||||
scene_file.write_text(scene_code, encoding="utf-8")
|
||||
|
||||
# Build Manim CLI command
|
||||
cmd = ["manim"]
|
||||
|
||||
# Quality flag
|
||||
preset = QUALITY_PRESETS.get(quality, QUALITY_PRESETS["medium"])
|
||||
for flag_part in preset["flag"].split():
|
||||
cmd.append(flag_part)
|
||||
|
||||
# Format
|
||||
if output_format == "gif":
|
||||
cmd.append("--format")
|
||||
cmd.append("gif")
|
||||
elif output_format == "webm":
|
||||
cmd.append("--format")
|
||||
cmd.append("webm")
|
||||
elif output_format == "png":
|
||||
cmd.append("-s") # save last frame as PNG
|
||||
|
||||
# Transparent background
|
||||
if transparent:
|
||||
cmd.append("--transparent")
|
||||
|
||||
# Background color
|
||||
if bg_color:
|
||||
cmd.extend(["--background_color", bg_color])
|
||||
|
||||
# Disable window preview (headless rendering)
|
||||
cmd.append("--disable_caching")
|
||||
|
||||
# Extra args
|
||||
cmd.extend(extra_args)
|
||||
|
||||
# Scene file and class name
|
||||
cmd.append(str(scene_file))
|
||||
cmd.append(scene_name)
|
||||
|
||||
# Execute Manim
|
||||
try:
|
||||
proc = subprocess.run(
|
||||
cmd,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=300, # 5 min timeout
|
||||
cwd=str(work_dir),
|
||||
)
|
||||
except subprocess.TimeoutExpired:
|
||||
self._cleanup(work_dir)
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"Manim render timed out after 300s. Try 'low' or 'preview' quality.",
|
||||
)
|
||||
|
||||
if proc.returncode != 0:
|
||||
error_msg = proc.stderr or proc.stdout or "Unknown error"
|
||||
# Extract the most useful part of the error
|
||||
lines = error_msg.strip().split("\n")
|
||||
# Look for the actual error (skip Manim header/progress)
|
||||
error_lines = [l for l in lines if "Error" in l or "error" in l or "Traceback" in l]
|
||||
if error_lines:
|
||||
error_msg = "\n".join(lines[lines.index(error_lines[0]):])
|
||||
self._cleanup(work_dir)
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"Manim render failed:\n{error_msg}",
|
||||
data={"full_stderr": proc.stderr, "full_stdout": proc.stdout},
|
||||
)
|
||||
|
||||
# Find the output file
|
||||
rendered_file = self._find_output(work_dir, scene_name, output_format)
|
||||
if not rendered_file:
|
||||
self._cleanup(work_dir)
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"Render succeeded but output file not found. Manim output:\n{proc.stdout}",
|
||||
)
|
||||
|
||||
# Move to desired output path
|
||||
if output_path:
|
||||
final_path = Path(output_path)
|
||||
else:
|
||||
ext = rendered_file.suffix
|
||||
final_path = Path(f"manim_{scene_name}{ext}")
|
||||
|
||||
final_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copy2(str(rendered_file), str(final_path))
|
||||
|
||||
# Get video info
|
||||
video_info = self._probe_output(final_path)
|
||||
|
||||
# Cleanup temp directory
|
||||
self._cleanup(work_dir)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"scene_name": scene_name,
|
||||
"quality": quality,
|
||||
"format": output_format,
|
||||
"output": str(final_path),
|
||||
"resolution": preset["resolution"],
|
||||
"fps": preset["fps"],
|
||||
**video_info,
|
||||
},
|
||||
artifacts=[str(final_path)],
|
||||
)
|
||||
|
||||
def _detect_scene_name(self, code: str) -> Optional[str]:
|
||||
"""Extract Scene subclass name from code."""
|
||||
import re
|
||||
|
||||
# Match class definitions that inherit from Scene or its variants
|
||||
pattern = r"class\s+(\w+)\s*\(\s*(?:Scene|ThreeDScene|MovingCameraScene|ZoomedScene)\s*\)"
|
||||
matches = re.findall(pattern, code)
|
||||
if len(matches) == 1:
|
||||
return matches[0]
|
||||
if len(matches) > 1:
|
||||
# Return the last one (convention: main scene is last)
|
||||
return matches[-1]
|
||||
return None
|
||||
|
||||
def _find_output(self, work_dir: Path, scene_name: str, fmt: str) -> Optional[Path]:
|
||||
"""Find Manim's output file in the media directory."""
|
||||
media_dir = work_dir / "media"
|
||||
if not media_dir.exists():
|
||||
return None
|
||||
|
||||
# Manim outputs to media/videos/<scene_file>/<quality>/<SceneName>.<ext>
|
||||
# or media/images/<scene_file>/<SceneName>.<ext> for -s flag
|
||||
ext_map = {"mp4": ".mp4", "gif": ".gif", "webm": ".webm", "png": ".png"}
|
||||
target_ext = ext_map.get(fmt, ".mp4")
|
||||
|
||||
# Search recursively for the output file
|
||||
for path in media_dir.rglob(f"{scene_name}{target_ext}"):
|
||||
return path
|
||||
|
||||
# Fallback: any file with the right extension
|
||||
for path in media_dir.rglob(f"*{target_ext}"):
|
||||
return path
|
||||
|
||||
return None
|
||||
|
||||
def _probe_output(self, path: Path) -> dict[str, Any]:
|
||||
"""Get basic info about the rendered file."""
|
||||
info: dict[str, Any] = {"file_size_bytes": path.stat().st_size}
|
||||
|
||||
if not shutil.which("ffprobe"):
|
||||
return info
|
||||
|
||||
try:
|
||||
proc = subprocess.run(
|
||||
[
|
||||
"ffprobe", "-v", "quiet",
|
||||
"-print_format", "json",
|
||||
"-show_format", "-show_streams",
|
||||
str(path),
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=10,
|
||||
)
|
||||
if proc.returncode == 0:
|
||||
import json
|
||||
probe = json.loads(proc.stdout)
|
||||
fmt = probe.get("format", {})
|
||||
info["duration_seconds"] = float(fmt.get("duration", 0))
|
||||
info["file_size_mb"] = round(path.stat().st_size / (1024 * 1024), 2)
|
||||
for stream in probe.get("streams", []):
|
||||
if stream.get("codec_type") == "video":
|
||||
info["video_width"] = int(stream.get("width", 0))
|
||||
info["video_height"] = int(stream.get("height", 0))
|
||||
info["video_codec"] = stream.get("codec_name", "")
|
||||
break
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return info
|
||||
|
||||
def _cleanup(self, work_dir: Path) -> None:
|
||||
"""Remove temp working directory."""
|
||||
try:
|
||||
shutil.rmtree(str(work_dir), ignore_errors=True)
|
||||
except Exception:
|
||||
pass
|
||||
@@ -0,0 +1,176 @@
|
||||
"""OpenAI GPT Image generation (gpt-image-1 / DALL-E 3)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class OpenAIImage(BaseTool):
|
||||
name = "openai_image"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "image_generation"
|
||||
provider = "openai"
|
||||
stability = ToolStability.BETA
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.STOCHASTIC
|
||||
runtime = ToolRuntime.API
|
||||
|
||||
dependencies = [] # checked dynamically
|
||||
install_instructions = (
|
||||
"Set OPENAI_API_KEY to your OpenAI API key.\n"
|
||||
" pip install openai"
|
||||
)
|
||||
agent_skills = ["flux-best-practices"] # general image gen knowledge
|
||||
|
||||
capabilities = ["generate_image", "generate_illustration", "text_to_image"]
|
||||
supports = {
|
||||
"complex_instructions": True,
|
||||
"text_in_image": True,
|
||||
"multiple_outputs": True,
|
||||
}
|
||||
best_for = [
|
||||
"complex multi-element compositions",
|
||||
"images with text/labels",
|
||||
"following detailed instructions accurately",
|
||||
]
|
||||
not_good_for = ["offline generation", "budget-constrained projects at high quality"]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {"type": "string"},
|
||||
"model": {
|
||||
"type": "string",
|
||||
"enum": ["gpt-image-1", "dall-e-3"],
|
||||
"default": "gpt-image-1",
|
||||
},
|
||||
"size": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"1024x1024", "1536x1024", "1024x1536", "auto",
|
||||
"1024x1792", "1792x1024", # dall-e-3 only
|
||||
],
|
||||
"default": "1024x1024",
|
||||
},
|
||||
"quality": {
|
||||
"type": "string",
|
||||
"enum": ["low", "medium", "high", "auto", "standard", "hd"],
|
||||
"default": "high",
|
||||
},
|
||||
"output_format": {
|
||||
"type": "string",
|
||||
"enum": ["png", "jpeg", "webp"],
|
||||
"default": "png",
|
||||
},
|
||||
"n": {"type": "integer", "default": 1, "minimum": 1, "maximum": 4},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=1, ram_mb=512, vram_mb=0, disk_mb=100, network_required=True
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=2, retryable_errors=["rate_limit", "timeout"])
|
||||
idempotency_key_fields = ["prompt", "size", "quality", "model"]
|
||||
side_effects = ["writes image file to output_path", "calls OpenAI API"]
|
||||
user_visible_verification = ["Inspect generated image for relevance and quality"]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if os.environ.get("OPENAI_API_KEY"):
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
model = inputs.get("model", "gpt-image-1")
|
||||
quality = inputs.get("quality", "high")
|
||||
n = inputs.get("n", 1)
|
||||
if model == "gpt-image-1":
|
||||
cost_map = {"low": 0.011, "medium": 0.042, "high": 0.167, "auto": 0.042}
|
||||
return cost_map.get(quality, 0.042) * n
|
||||
# dall-e-3 fallback pricing
|
||||
quality_map = {"standard": 0.04, "hd": 0.08}
|
||||
return quality_map.get(quality, 0.04) * n
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
if not os.environ.get("OPENAI_API_KEY"):
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="OPENAI_API_KEY not set. " + self.install_instructions,
|
||||
)
|
||||
|
||||
from openai import OpenAI
|
||||
|
||||
start = time.time()
|
||||
client = OpenAI()
|
||||
model = inputs.get("model", "gpt-image-1")
|
||||
prompt = inputs["prompt"]
|
||||
size = inputs.get("size", "1024x1024")
|
||||
n = inputs.get("n", 1)
|
||||
|
||||
try:
|
||||
if model == "gpt-image-1":
|
||||
quality = inputs.get("quality", "high")
|
||||
output_format = inputs.get("output_format", "png")
|
||||
response = client.images.generate(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
size=size,
|
||||
quality=quality,
|
||||
output_format=output_format,
|
||||
n=n,
|
||||
)
|
||||
else:
|
||||
# dall-e-3 path
|
||||
quality = inputs.get("quality", "standard")
|
||||
if quality in ("low", "medium", "high", "auto"):
|
||||
quality = "standard" # map to dall-e-3 quality options
|
||||
response = client.images.generate(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
size=size,
|
||||
quality=quality,
|
||||
n=1, # dall-e-3 only supports n=1
|
||||
response_format="b64_json",
|
||||
)
|
||||
|
||||
image_data = base64.b64decode(response.data[0].b64_json)
|
||||
ext = inputs.get("output_format", "png")
|
||||
output_path = Path(inputs.get("output_path", f"generated_image.{ext}"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(image_data)
|
||||
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"OpenAI image generation failed: {e}")
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "openai",
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
cost_usd=self.estimate_cost(inputs),
|
||||
duration_seconds=round(time.time() - start, 2),
|
||||
model=model,
|
||||
)
|
||||
@@ -0,0 +1,184 @@
|
||||
"""Stock image acquisition from Pexels API (free)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class PexelsImage(BaseTool):
|
||||
name = "pexels_image"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.SOURCE
|
||||
capability = "image_generation"
|
||||
provider = "pexels"
|
||||
stability = ToolStability.BETA
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
runtime = ToolRuntime.API
|
||||
|
||||
dependencies = []
|
||||
install_instructions = (
|
||||
"Set PEXELS_API_KEY to your Pexels API key.\n"
|
||||
" Get one free at https://www.pexels.com/api/"
|
||||
)
|
||||
agent_skills = []
|
||||
|
||||
capabilities = ["search_image", "download_image", "stock_image"]
|
||||
supports = {
|
||||
"orientation_filter": True,
|
||||
"size_filter": True,
|
||||
"color_filter": True,
|
||||
"locale": True,
|
||||
"free_commercial_use": True,
|
||||
}
|
||||
best_for = [
|
||||
"real-world photography (cities, nature, people, objects)",
|
||||
"establishing shots and B-roll stills",
|
||||
"free stock images — no cost, no attribution required",
|
||||
]
|
||||
not_good_for = [
|
||||
"custom/specific compositions",
|
||||
"abstract or stylized graphics",
|
||||
"offline use",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["query"],
|
||||
"properties": {
|
||||
"query": {"type": "string", "description": "Search term"},
|
||||
"orientation": {
|
||||
"type": "string",
|
||||
"enum": ["landscape", "portrait", "square"],
|
||||
},
|
||||
"size": {
|
||||
"type": "string",
|
||||
"enum": ["large", "medium", "small"],
|
||||
"description": "large=24MP+, medium=12MP+, small=4MP+",
|
||||
},
|
||||
"color": {
|
||||
"type": "string",
|
||||
"description": "Hex without # (e.g. FF0000) or color name (red, blue, etc.)",
|
||||
},
|
||||
"per_page": {"type": "integer", "default": 5, "minimum": 1, "maximum": 80},
|
||||
"page": {"type": "integer", "default": 1},
|
||||
"download_size": {
|
||||
"type": "string",
|
||||
"enum": ["original", "large2x", "large", "medium"],
|
||||
"default": "large2x",
|
||||
},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=1, ram_mb=256, vram_mb=0, disk_mb=50, network_required=True
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=2, retryable_errors=["rate_limit", "timeout"])
|
||||
idempotency_key_fields = ["query", "orientation", "size", "color", "page"]
|
||||
side_effects = ["writes image file to output_path", "calls Pexels API"]
|
||||
user_visible_verification = ["Check that downloaded image matches the intended scene"]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if os.environ.get("PEXELS_API_KEY"):
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
return 0.0 # Pexels is free
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
api_key = os.environ.get("PEXELS_API_KEY")
|
||||
if not api_key:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="PEXELS_API_KEY not set. " + self.install_instructions,
|
||||
)
|
||||
|
||||
import requests
|
||||
|
||||
start = time.time()
|
||||
query = inputs["query"]
|
||||
|
||||
params: dict[str, Any] = {
|
||||
"query": query,
|
||||
"per_page": inputs.get("per_page", 5),
|
||||
"page": inputs.get("page", 1),
|
||||
}
|
||||
if inputs.get("orientation"):
|
||||
params["orientation"] = inputs["orientation"]
|
||||
if inputs.get("size"):
|
||||
params["size"] = inputs["size"]
|
||||
if inputs.get("color"):
|
||||
params["color"] = inputs["color"]
|
||||
|
||||
try:
|
||||
search_response = requests.get(
|
||||
"https://api.pexels.com/v1/search",
|
||||
headers={"Authorization": api_key},
|
||||
params=params,
|
||||
timeout=30,
|
||||
)
|
||||
search_response.raise_for_status()
|
||||
data = search_response.json()
|
||||
|
||||
photos = data.get("photos", [])
|
||||
if not photos:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"No images found for query: {query}",
|
||||
data={"total_results": data.get("total_results", 0)},
|
||||
)
|
||||
|
||||
# Pick the first result (agent can refine query if needed)
|
||||
photo = photos[0]
|
||||
download_size = inputs.get("download_size", "large2x")
|
||||
image_url = photo["src"].get(download_size, photo["src"]["large2x"])
|
||||
|
||||
image_response = requests.get(image_url, timeout=60)
|
||||
image_response.raise_for_status()
|
||||
|
||||
output_path = Path(inputs.get("output_path", f"pexels_{photo['id']}.jpg"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(image_response.content)
|
||||
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Pexels image search failed: {e}")
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "pexels",
|
||||
"photo_id": photo["id"],
|
||||
"photographer": photo.get("photographer", "Unknown"),
|
||||
"photographer_url": photo.get("photographer_url", ""),
|
||||
"alt": photo.get("alt", ""),
|
||||
"width": photo.get("width"),
|
||||
"height": photo.get("height"),
|
||||
"query": query,
|
||||
"output": str(output_path),
|
||||
"total_results": data.get("total_results", 0),
|
||||
"results_returned": len(photos),
|
||||
"license": "Pexels License (free, no attribution required)",
|
||||
"pexels_url": photo.get("url", ""),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
cost_usd=0.0,
|
||||
duration_seconds=round(time.time() - start, 2),
|
||||
)
|
||||
@@ -0,0 +1,196 @@
|
||||
"""Stock image acquisition from Pixabay API (free)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class PixabayImage(BaseTool):
|
||||
name = "pixabay_image"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.SOURCE
|
||||
capability = "image_generation"
|
||||
provider = "pixabay"
|
||||
stability = ToolStability.BETA
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
runtime = ToolRuntime.API
|
||||
|
||||
dependencies = []
|
||||
install_instructions = (
|
||||
"Set PIXABAY_API_KEY to your Pixabay API key.\n"
|
||||
" Get one free at https://pixabay.com/api/docs/"
|
||||
)
|
||||
agent_skills = []
|
||||
|
||||
capabilities = ["search_image", "download_image", "stock_image"]
|
||||
supports = {
|
||||
"orientation_filter": True,
|
||||
"category_filter": True,
|
||||
"color_filter": True,
|
||||
"image_type_filter": True,
|
||||
"editors_choice": True,
|
||||
"free_commercial_use": True,
|
||||
}
|
||||
best_for = [
|
||||
"large royalty-free library (5M+ images)",
|
||||
"category-based filtering (nature, business, science, etc.)",
|
||||
"free stock images — no cost, no attribution required",
|
||||
]
|
||||
not_good_for = [
|
||||
"full-resolution originals (standard API limited to 1280px)",
|
||||
"custom compositions",
|
||||
"offline use",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["query"],
|
||||
"properties": {
|
||||
"query": {"type": "string", "description": "Search term (max 100 chars)"},
|
||||
"image_type": {
|
||||
"type": "string",
|
||||
"enum": ["all", "photo", "illustration", "vector"],
|
||||
"default": "all",
|
||||
},
|
||||
"orientation": {
|
||||
"type": "string",
|
||||
"enum": ["all", "horizontal", "vertical"],
|
||||
"default": "all",
|
||||
},
|
||||
"category": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"backgrounds", "fashion", "nature", "science", "education",
|
||||
"feelings", "health", "people", "religion", "places",
|
||||
"animals", "industry", "computer", "food", "sports",
|
||||
"transportation", "travel", "buildings", "business", "music",
|
||||
],
|
||||
},
|
||||
"colors": {
|
||||
"type": "string",
|
||||
"description": "Comma-separated: grayscale, transparent, red, orange, yellow, green, turquoise, blue, lilac, pink, white, gray, black, brown",
|
||||
},
|
||||
"editors_choice": {"type": "boolean", "default": False},
|
||||
"safesearch": {"type": "boolean", "default": True},
|
||||
"per_page": {"type": "integer", "default": 5, "minimum": 3, "maximum": 200},
|
||||
"page": {"type": "integer", "default": 1},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=1, ram_mb=256, vram_mb=0, disk_mb=50, network_required=True
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=2, retryable_errors=["rate_limit", "timeout"])
|
||||
idempotency_key_fields = ["query", "image_type", "orientation", "category", "page"]
|
||||
side_effects = ["writes image file to output_path", "calls Pixabay API"]
|
||||
user_visible_verification = ["Check that downloaded image matches the intended scene"]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if os.environ.get("PIXABAY_API_KEY"):
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
return 0.0 # Pixabay is free
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
api_key = os.environ.get("PIXABAY_API_KEY")
|
||||
if not api_key:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="PIXABAY_API_KEY not set. " + self.install_instructions,
|
||||
)
|
||||
|
||||
import requests
|
||||
|
||||
start = time.time()
|
||||
query = inputs["query"]
|
||||
|
||||
params: dict[str, Any] = {
|
||||
"key": api_key,
|
||||
"q": query,
|
||||
"per_page": inputs.get("per_page", 5),
|
||||
"page": inputs.get("page", 1),
|
||||
"safesearch": str(inputs.get("safesearch", True)).lower(),
|
||||
}
|
||||
if inputs.get("image_type") and inputs["image_type"] != "all":
|
||||
params["image_type"] = inputs["image_type"]
|
||||
if inputs.get("orientation") and inputs["orientation"] != "all":
|
||||
params["orientation"] = inputs["orientation"]
|
||||
if inputs.get("category"):
|
||||
params["category"] = inputs["category"]
|
||||
if inputs.get("colors"):
|
||||
params["colors"] = inputs["colors"]
|
||||
if inputs.get("editors_choice"):
|
||||
params["editors_choice"] = "true"
|
||||
|
||||
try:
|
||||
search_response = requests.get(
|
||||
"https://pixabay.com/api/",
|
||||
params=params,
|
||||
timeout=30,
|
||||
)
|
||||
search_response.raise_for_status()
|
||||
data = search_response.json()
|
||||
|
||||
hits = data.get("hits", [])
|
||||
if not hits:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"No images found for query: {query}",
|
||||
data={"total_results": data.get("total", 0)},
|
||||
)
|
||||
|
||||
hit = hits[0]
|
||||
# largeImageURL is the best available at standard API tier (1280px)
|
||||
image_url = hit.get("largeImageURL", hit.get("webformatURL"))
|
||||
|
||||
# Download immediately — Pixabay URLs contain embedded tokens that expire
|
||||
image_response = requests.get(image_url, timeout=60)
|
||||
image_response.raise_for_status()
|
||||
|
||||
output_path = Path(inputs.get("output_path", f"pixabay_{hit['id']}.jpg"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(image_response.content)
|
||||
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Pixabay image search failed: {e}")
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "pixabay",
|
||||
"image_id": hit["id"],
|
||||
"user": hit.get("user", "Unknown"),
|
||||
"tags": hit.get("tags", ""),
|
||||
"image_width": hit.get("imageWidth"),
|
||||
"image_height": hit.get("imageHeight"),
|
||||
"query": query,
|
||||
"output": str(output_path),
|
||||
"total_results": data.get("total", 0),
|
||||
"results_returned": len(hits),
|
||||
"license": "Pixabay Content License (free, no attribution required)",
|
||||
"page_url": hit.get("pageURL", ""),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
cost_usd=0.0,
|
||||
duration_seconds=round(time.time() - start, 2),
|
||||
)
|
||||
@@ -0,0 +1,188 @@
|
||||
"""Recraft V4 image generation via fal.ai API.
|
||||
|
||||
Best for logos, brand assets, SVG vectors, and images with accurate text rendering.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class RecraftImage(BaseTool):
|
||||
name = "recraft_image"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "image_generation"
|
||||
provider = "recraft"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.STOCHASTIC
|
||||
runtime = ToolRuntime.API
|
||||
|
||||
dependencies = []
|
||||
install_instructions = (
|
||||
"Set FAL_KEY to your fal.ai API key.\n"
|
||||
" Get one at https://fal.ai/dashboard/keys"
|
||||
)
|
||||
agent_skills = []
|
||||
|
||||
capabilities = [
|
||||
"generate_image",
|
||||
"generate_logo",
|
||||
"generate_vector",
|
||||
"text_to_image",
|
||||
]
|
||||
supports = {
|
||||
"svg_output": True,
|
||||
"text_rendering": True,
|
||||
"color_palette": True,
|
||||
"custom_size": True,
|
||||
}
|
||||
best_for = [
|
||||
"logos and brand assets",
|
||||
"SVG vector output",
|
||||
"images with accurate text rendering",
|
||||
"clean professional graphics",
|
||||
]
|
||||
not_good_for = ["photorealistic images", "offline generation"]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {"type": "string"},
|
||||
"model": {
|
||||
"type": "string",
|
||||
"enum": ["v4", "v4-pro"],
|
||||
"default": "v4",
|
||||
},
|
||||
"image_size": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"square", "square_hd",
|
||||
"landscape_4_3", "landscape_16_9",
|
||||
"portrait_4_3", "portrait_16_9",
|
||||
],
|
||||
"default": "square_hd",
|
||||
},
|
||||
"style": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"any", "realistic_image", "digital_illustration",
|
||||
"vector_illustration", "icon",
|
||||
],
|
||||
"default": "any",
|
||||
},
|
||||
"colors": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
"description": "Color palette as hex strings, e.g. ['#FF5733', '#2E86C1']",
|
||||
},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=1, ram_mb=512, vram_mb=0, disk_mb=100, network_required=True
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=2, retryable_errors=["rate_limit", "timeout"])
|
||||
idempotency_key_fields = ["prompt", "model", "style", "image_size"]
|
||||
side_effects = ["writes image file to output_path", "calls fal.ai API"]
|
||||
user_visible_verification = ["Inspect generated image for brand accuracy and text readability"]
|
||||
|
||||
def _get_api_key(self) -> str | None:
|
||||
return os.environ.get("FAL_KEY") or os.environ.get("FAL_AI_API_KEY")
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if self._get_api_key():
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
model = inputs.get("model", "v4")
|
||||
if model == "v4-pro":
|
||||
return 0.25
|
||||
return 0.04
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
api_key = self._get_api_key()
|
||||
if not api_key:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="FAL_KEY not set. " + self.install_instructions,
|
||||
)
|
||||
|
||||
import requests
|
||||
|
||||
start = time.time()
|
||||
model = inputs.get("model", "v4")
|
||||
prompt = inputs["prompt"]
|
||||
|
||||
model_path = f"recraft/{model}/text-to-image"
|
||||
if model == "v4-pro":
|
||||
model_path = "recraft/v4/pro/text-to-image"
|
||||
elif model == "v4":
|
||||
model_path = "recraft/v4/text-to-image"
|
||||
|
||||
payload: dict[str, Any] = {"prompt": prompt}
|
||||
if inputs.get("image_size"):
|
||||
payload["image_size"] = inputs["image_size"]
|
||||
if inputs.get("style"):
|
||||
payload["style"] = inputs["style"]
|
||||
if inputs.get("colors"):
|
||||
payload["colors"] = inputs["colors"]
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
f"https://fal.run/fal-ai/{model_path}",
|
||||
headers={
|
||||
"Authorization": f"Key {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
json=payload,
|
||||
timeout=120,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
image_url = data["images"][0]["url"]
|
||||
image_response = requests.get(image_url, timeout=60)
|
||||
image_response.raise_for_status()
|
||||
|
||||
ext = "svg" if inputs.get("style") == "vector_illustration" else "png"
|
||||
output_path = Path(inputs.get("output_path", f"generated_image.{ext}"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(image_response.content)
|
||||
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Recraft generation failed: {e}")
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "recraft",
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
cost_usd=self.estimate_cost(inputs),
|
||||
duration_seconds=round(time.time() - start, 2),
|
||||
model=f"fal-ai/{model_path}",
|
||||
)
|
||||
@@ -0,0 +1 @@
|
||||
"""Subtitle generation and formatting tools."""
|
||||
@@ -0,0 +1,276 @@
|
||||
"""Subtitle generation tool.
|
||||
|
||||
Converts word-level timestamps from the transcriber into SRT, VTT,
|
||||
or caption JSON formats. Pure Python — no external dependencies beyond
|
||||
the standard library.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
ToolResult,
|
||||
ToolStability,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class SubtitleGen(BaseTool):
|
||||
name = "subtitle_gen"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.CORE
|
||||
capability = "subtitle"
|
||||
provider = "openmontage"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
|
||||
dependencies = [] # pure Python
|
||||
install_instructions = "No external dependencies required."
|
||||
agent_skills = ["remotion-best-practices"]
|
||||
|
||||
capabilities = ["generate_srt", "generate_vtt", "generate_caption_json"]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["segments"],
|
||||
"properties": {
|
||||
"segments": {
|
||||
"type": "array",
|
||||
"description": "Transcript segments from transcriber (with words and timestamps)",
|
||||
},
|
||||
"format": {
|
||||
"type": "string",
|
||||
"enum": ["srt", "vtt", "json"],
|
||||
"default": "srt",
|
||||
},
|
||||
"output_path": {"type": "string"},
|
||||
"max_chars_per_line": {"type": "integer", "default": 42},
|
||||
"max_words_per_cue": {"type": "integer", "default": 8},
|
||||
"highlight_style": {
|
||||
"type": "string",
|
||||
"enum": ["none", "word_by_word", "karaoke"],
|
||||
"default": "none",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(cpu_cores=1, ram_mb=128, vram_mb=0, disk_mb=10)
|
||||
idempotency_key_fields = ["segments", "format", "max_words_per_cue"]
|
||||
side_effects = ["writes subtitle file to output_path"]
|
||||
user_visible_verification = [
|
||||
"Play video with generated subtitles and verify timing",
|
||||
]
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
segments = inputs["segments"]
|
||||
fmt = inputs.get("format", "srt")
|
||||
max_words = inputs.get("max_words_per_cue", 8)
|
||||
max_chars = inputs.get("max_chars_per_line", 42)
|
||||
highlight_style = inputs.get("highlight_style", "none")
|
||||
output_path = inputs.get("output_path")
|
||||
|
||||
start = time.time()
|
||||
|
||||
# Build cues from word-level timestamps
|
||||
cues = self._build_cues(segments, max_words, max_chars)
|
||||
|
||||
if fmt == "srt":
|
||||
content = self._render_srt(cues, highlight_style)
|
||||
ext = ".srt"
|
||||
elif fmt == "vtt":
|
||||
content = self._render_vtt(cues, highlight_style)
|
||||
ext = ".vtt"
|
||||
elif fmt == "json":
|
||||
content = json.dumps({"cues": cues, "highlight_style": highlight_style}, indent=2)
|
||||
ext = ".caption.json"
|
||||
else:
|
||||
return ToolResult(success=False, error=f"Unknown format: {fmt}")
|
||||
|
||||
if output_path is None:
|
||||
output_path = f"subtitles{ext}"
|
||||
out = Path(output_path)
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
out.write_text(content, encoding="utf-8")
|
||||
|
||||
elapsed = time.time() - start
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"format": fmt,
|
||||
"cue_count": len(cues),
|
||||
"output": str(out),
|
||||
},
|
||||
artifacts=[str(out)],
|
||||
duration_seconds=round(elapsed, 2),
|
||||
)
|
||||
|
||||
def _build_cues(
|
||||
self, segments: list[dict], max_words: int, max_chars: int
|
||||
) -> list[dict]:
|
||||
"""Group words into display cues respecting max_words and max_chars."""
|
||||
# Collect all words with timestamps
|
||||
all_words = []
|
||||
for seg in segments:
|
||||
words = seg.get("words", [])
|
||||
if words:
|
||||
all_words.extend(words)
|
||||
elif "text" in seg:
|
||||
# Fallback: segment-level only (no word timestamps)
|
||||
all_words.append({
|
||||
"word": seg["text"],
|
||||
"start": seg["start"],
|
||||
"end": seg["end"],
|
||||
})
|
||||
|
||||
if not all_words:
|
||||
return []
|
||||
|
||||
cues = []
|
||||
buf: list[dict] = []
|
||||
buf_text = ""
|
||||
|
||||
for w in all_words:
|
||||
word_text = w["word"].strip()
|
||||
candidate = f"{buf_text} {word_text}".strip() if buf_text else word_text
|
||||
|
||||
if buf and (len(buf) >= max_words or len(candidate) > max_chars):
|
||||
cues.append({
|
||||
"index": len(cues) + 1,
|
||||
"start": buf[0]["start"],
|
||||
"end": buf[-1]["end"],
|
||||
"text": buf_text,
|
||||
"words": [
|
||||
{"word": b["word"].strip(), "start": b["start"], "end": b["end"]}
|
||||
for b in buf
|
||||
],
|
||||
})
|
||||
buf = []
|
||||
buf_text = ""
|
||||
|
||||
buf.append(w)
|
||||
buf_text = f"{buf_text} {word_text}".strip() if buf_text else word_text
|
||||
|
||||
# Flush remaining
|
||||
if buf:
|
||||
cues.append({
|
||||
"index": len(cues) + 1,
|
||||
"start": buf[0]["start"],
|
||||
"end": buf[-1]["end"],
|
||||
"text": buf_text,
|
||||
"words": [
|
||||
{"word": b["word"].strip(), "start": b["start"], "end": b["end"]}
|
||||
for b in buf
|
||||
],
|
||||
})
|
||||
|
||||
return cues
|
||||
|
||||
def _render_srt(self, cues: list[dict], highlight_style: str = "none") -> str:
|
||||
lines = []
|
||||
if highlight_style == "word_by_word":
|
||||
# Emit one cue per word for word-by-word reveal
|
||||
idx = 1
|
||||
for cue in cues:
|
||||
for word_info in cue.get("words", []):
|
||||
lines.append(str(idx))
|
||||
lines.append(
|
||||
f"{self._ts_srt(word_info['start'])} --> {self._ts_srt(word_info['end'])}"
|
||||
)
|
||||
lines.append(word_info["word"])
|
||||
lines.append("")
|
||||
idx += 1
|
||||
elif highlight_style == "karaoke":
|
||||
# Show full cue text but bold the active word using SRT HTML tags
|
||||
for cue in cues:
|
||||
words = cue.get("words", [])
|
||||
if not words:
|
||||
lines.append(str(cue["index"]))
|
||||
lines.append(f"{self._ts_srt(cue['start'])} --> {self._ts_srt(cue['end'])}")
|
||||
lines.append(cue["text"])
|
||||
lines.append("")
|
||||
continue
|
||||
for wi, word_info in enumerate(words):
|
||||
lines.append(str(cue["index"] * 100 + wi))
|
||||
lines.append(
|
||||
f"{self._ts_srt(word_info['start'])} --> {self._ts_srt(word_info['end'])}"
|
||||
)
|
||||
parts = []
|
||||
for wj, w in enumerate(words):
|
||||
if wj == wi:
|
||||
parts.append(f"<b>{w['word']}</b>")
|
||||
else:
|
||||
parts.append(w["word"])
|
||||
lines.append(" ".join(parts))
|
||||
lines.append("")
|
||||
else:
|
||||
for cue in cues:
|
||||
lines.append(str(cue["index"]))
|
||||
lines.append(f"{self._ts_srt(cue['start'])} --> {self._ts_srt(cue['end'])}")
|
||||
lines.append(cue["text"])
|
||||
lines.append("")
|
||||
return "\n".join(lines)
|
||||
|
||||
def _render_vtt(self, cues: list[dict], highlight_style: str = "none") -> str:
|
||||
lines = ["WEBVTT", ""]
|
||||
if highlight_style == "word_by_word":
|
||||
for cue in cues:
|
||||
for word_info in cue.get("words", []):
|
||||
lines.append(
|
||||
f"{self._ts_vtt(word_info['start'])} --> {self._ts_vtt(word_info['end'])}"
|
||||
)
|
||||
lines.append(word_info["word"])
|
||||
lines.append("")
|
||||
elif highlight_style == "karaoke":
|
||||
for cue in cues:
|
||||
words = cue.get("words", [])
|
||||
if not words:
|
||||
lines.append(f"{self._ts_vtt(cue['start'])} --> {self._ts_vtt(cue['end'])}")
|
||||
lines.append(cue["text"])
|
||||
lines.append("")
|
||||
continue
|
||||
for wi, word_info in enumerate(words):
|
||||
lines.append(
|
||||
f"{self._ts_vtt(word_info['start'])} --> {self._ts_vtt(word_info['end'])}"
|
||||
)
|
||||
parts = []
|
||||
for wj, w in enumerate(words):
|
||||
if wj == wi:
|
||||
parts.append(f"<b>{w['word']}</b>")
|
||||
else:
|
||||
parts.append(w["word"])
|
||||
lines.append(" ".join(parts))
|
||||
lines.append("")
|
||||
else:
|
||||
for cue in cues:
|
||||
lines.append(f"{self._ts_vtt(cue['start'])} --> {self._ts_vtt(cue['end'])}")
|
||||
lines.append(cue["text"])
|
||||
lines.append("")
|
||||
return "\n".join(lines)
|
||||
|
||||
@staticmethod
|
||||
def _ts_srt(seconds: float) -> str:
|
||||
"""Format seconds as SRT timestamp: HH:MM:SS,mmm"""
|
||||
h = int(seconds // 3600)
|
||||
m = int((seconds % 3600) // 60)
|
||||
s = int(seconds % 60)
|
||||
ms = int(round((seconds % 1) * 1000))
|
||||
return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"
|
||||
|
||||
@staticmethod
|
||||
def _ts_vtt(seconds: float) -> str:
|
||||
"""Format seconds as VTT timestamp: HH:MM:SS.mmm"""
|
||||
h = int(seconds // 3600)
|
||||
m = int((seconds % 3600) // 60)
|
||||
s = int(seconds % 60)
|
||||
ms = int(round((seconds % 1) * 1000))
|
||||
return f"{h:02d}:{m:02d}:{s:02d}.{ms:03d}"
|
||||
@@ -0,0 +1,263 @@
|
||||
"""Tool registry with status, stability, and support-envelope reporting.
|
||||
|
||||
The registry discovers all registered tools, reports their availability,
|
||||
and lets the orchestrator/agents query capabilities by tier, status, etc.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import inspect
|
||||
import pkgutil
|
||||
from types import ModuleType
|
||||
from typing import Any, Optional
|
||||
|
||||
from tools.base_tool import BaseTool, ToolStatus, ToolTier, ToolStability
|
||||
|
||||
|
||||
class ToolRegistry:
|
||||
"""Central registry of all OpenMontage tools."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._tools: dict[str, BaseTool] = {}
|
||||
self._discovered_packages: set[str] = set()
|
||||
|
||||
def register(self, tool: BaseTool) -> None:
|
||||
"""Register a tool instance."""
|
||||
if not tool.name:
|
||||
raise ValueError("Tool must have a non-empty name")
|
||||
self._tools[tool.name] = tool
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear registered tools and discovery state."""
|
||||
self._tools.clear()
|
||||
self._discovered_packages.clear()
|
||||
|
||||
def register_module(self, module: ModuleType) -> list[str]:
|
||||
"""Register all concrete BaseTool subclasses defined in a module."""
|
||||
registered: list[str] = []
|
||||
for _, cls in inspect.getmembers(module, inspect.isclass):
|
||||
if cls is BaseTool or not issubclass(cls, BaseTool):
|
||||
continue
|
||||
if cls.__module__ != module.__name__ or inspect.isabstract(cls):
|
||||
continue
|
||||
tool = cls()
|
||||
self.register(tool)
|
||||
registered.append(tool.name)
|
||||
return registered
|
||||
|
||||
@staticmethod
|
||||
def _load_dotenv() -> None:
|
||||
"""Load .env file into os.environ if present, so tools can find API keys."""
|
||||
from pathlib import Path
|
||||
import os
|
||||
env_path = Path(__file__).resolve().parent.parent / ".env"
|
||||
if not env_path.is_file():
|
||||
return
|
||||
with open(env_path, encoding="utf-8", errors="ignore") as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line or line.startswith("#") or "=" not in line:
|
||||
continue
|
||||
key, _, value = line.partition("=")
|
||||
key = key.strip()
|
||||
value = value.strip().strip("'\"")
|
||||
if key and key not in os.environ:
|
||||
os.environ[key] = value
|
||||
|
||||
def discover(self, package_name: str = "tools") -> list[str]:
|
||||
"""Import a package tree and register any concrete tools it defines."""
|
||||
self._load_dotenv()
|
||||
package = importlib.import_module(package_name)
|
||||
discovered: list[str] = []
|
||||
package_paths = getattr(package, "__path__", None)
|
||||
if package_paths is None:
|
||||
return self.register_module(package)
|
||||
|
||||
for module_info in pkgutil.walk_packages(package_paths, f"{package.__name__}."):
|
||||
if module_info.name.endswith(".base_tool") or module_info.name.endswith(".tool_registry"):
|
||||
continue
|
||||
module = importlib.import_module(module_info.name)
|
||||
discovered.extend(self.register_module(module))
|
||||
|
||||
self._discovered_packages.add(package_name)
|
||||
return discovered
|
||||
|
||||
def ensure_discovered(self, package_name: str = "tools") -> None:
|
||||
"""Load tool modules once before reporting capabilities."""
|
||||
if package_name not in self._discovered_packages:
|
||||
self.discover(package_name)
|
||||
|
||||
def get(self, name: str) -> Optional[BaseTool]:
|
||||
"""Get a tool by name."""
|
||||
return self._tools.get(name)
|
||||
|
||||
def list_all(self) -> list[str]:
|
||||
"""List all registered tool names."""
|
||||
return list(self._tools.keys())
|
||||
|
||||
def get_by_tier(self, tier: ToolTier) -> list[BaseTool]:
|
||||
"""Get all tools in a given tier."""
|
||||
return [t for t in self._tools.values() if t.tier == tier]
|
||||
|
||||
def get_by_capability(self, capability: str) -> list[BaseTool]:
|
||||
"""Get all tools registered for a top-level capability family."""
|
||||
return [t for t in self._tools.values() if t.capability == capability]
|
||||
|
||||
def get_by_provider(self, provider: str) -> list[BaseTool]:
|
||||
"""Get all tools backed by a specific provider."""
|
||||
return [t for t in self._tools.values() if t.provider == provider]
|
||||
|
||||
def get_by_status(self, status: ToolStatus) -> list[BaseTool]:
|
||||
"""Get all tools with a given status."""
|
||||
return [t for t in self._tools.values() if t.get_status() == status]
|
||||
|
||||
def get_available(self) -> list[BaseTool]:
|
||||
"""Get all tools that are currently available."""
|
||||
return self.get_by_status(ToolStatus.AVAILABLE)
|
||||
|
||||
def get_unavailable(self) -> list[BaseTool]:
|
||||
"""Get all tools that are currently unavailable."""
|
||||
return self.get_by_status(ToolStatus.UNAVAILABLE)
|
||||
|
||||
def get_by_stability(self, stability: ToolStability) -> list[BaseTool]:
|
||||
"""Get all tools at a given stability level."""
|
||||
return [t for t in self._tools.values() if t.stability == stability]
|
||||
|
||||
def find_by_capability(self, capability: str) -> list[BaseTool]:
|
||||
"""Find tools that declare a given capability."""
|
||||
return [
|
||||
t for t in self._tools.values()
|
||||
if capability in t.capabilities
|
||||
]
|
||||
|
||||
def find_fallback(self, tool_name: str) -> Optional[BaseTool]:
|
||||
"""Find the fallback tool for a given tool, if declared and available."""
|
||||
tool = self.get(tool_name)
|
||||
if tool is None:
|
||||
return None
|
||||
candidates = list(tool.fallback_tools or [])
|
||||
if tool.fallback and tool.fallback not in candidates:
|
||||
candidates.append(tool.fallback)
|
||||
for name in candidates:
|
||||
fb = self.get(name)
|
||||
if fb and fb.get_status() == ToolStatus.AVAILABLE:
|
||||
return fb
|
||||
return None
|
||||
|
||||
def support_envelope(self) -> dict[str, Any]:
|
||||
"""Generate a full support-envelope report for all tools.
|
||||
|
||||
Returns a dict mapping tool name to its contract info + live status.
|
||||
This is the primary report the orchestrator uses to understand
|
||||
what the system can and cannot do.
|
||||
"""
|
||||
self.ensure_discovered()
|
||||
report: dict[str, Any] = {}
|
||||
for name, tool in self._tools.items():
|
||||
info = tool.get_info()
|
||||
report[name] = info
|
||||
return report
|
||||
|
||||
def capability_catalog(self) -> dict[str, list[dict[str, Any]]]:
|
||||
"""Group the support envelope by top-level capability."""
|
||||
self.ensure_discovered()
|
||||
grouped: dict[str, list[dict[str, Any]]] = {}
|
||||
for tool in self._tools.values():
|
||||
grouped.setdefault(tool.capability, []).append(tool.get_info())
|
||||
for items in grouped.values():
|
||||
items.sort(key=lambda item: (item["provider"], item["name"]))
|
||||
return dict(sorted(grouped.items()))
|
||||
|
||||
def provider_catalog(self) -> dict[str, list[dict[str, Any]]]:
|
||||
"""Group the support envelope by provider."""
|
||||
self.ensure_discovered()
|
||||
grouped: dict[str, list[dict[str, Any]]] = {}
|
||||
for tool in self._tools.values():
|
||||
grouped.setdefault(tool.provider, []).append(tool.get_info())
|
||||
for items in grouped.values():
|
||||
items.sort(key=lambda item: (item["capability"], item["name"]))
|
||||
return dict(sorted(grouped.items()))
|
||||
|
||||
def tier_summary(self) -> dict[str, dict[str, int]]:
|
||||
"""Summarize tool counts by tier and status.
|
||||
|
||||
Returns:
|
||||
{"core": {"available": 5, "unavailable": 2, "degraded": 0}, ...}
|
||||
"""
|
||||
summary: dict[str, dict[str, int]] = {}
|
||||
for tier in ToolTier:
|
||||
tier_tools = self.get_by_tier(tier)
|
||||
counts = {"available": 0, "unavailable": 0, "degraded": 0}
|
||||
for t in tier_tools:
|
||||
status = t.get_status().value
|
||||
counts[status] = counts.get(status, 0) + 1
|
||||
if tier_tools:
|
||||
summary[tier.value] = counts
|
||||
return summary
|
||||
|
||||
def provider_menu(self) -> dict[str, dict[str, Any]]:
|
||||
"""Generate a capability-grouped provider menu for user-facing display.
|
||||
|
||||
Returns a dict like:
|
||||
{
|
||||
"video_generation": {
|
||||
"available": [{"name": ..., "provider": ..., "best_for": ...}],
|
||||
"unavailable": [{"name": ..., "provider": ..., "install_instructions": ...}],
|
||||
"total": 12,
|
||||
"configured": 2,
|
||||
},
|
||||
...
|
||||
}
|
||||
|
||||
This powers the agent's preflight provider menu — the agent reads this
|
||||
output and presents it to the user. Adding a new tool to tools/ is
|
||||
enough; this method auto-discovers it.
|
||||
"""
|
||||
self.ensure_discovered()
|
||||
menu: dict[str, dict[str, Any]] = {}
|
||||
|
||||
# Skip selectors — they aggregate, they aren't providers themselves
|
||||
tools = [t for t in self._tools.values() if t.provider != "selector"]
|
||||
|
||||
for tool in tools:
|
||||
cap = tool.capability
|
||||
if cap not in menu:
|
||||
menu[cap] = {"available": [], "unavailable": [], "total": 0, "configured": 0}
|
||||
|
||||
status = tool.get_status()
|
||||
entry = {
|
||||
"name": tool.name,
|
||||
"provider": tool.provider,
|
||||
"runtime": tool.runtime.value,
|
||||
"best_for": tool.best_for,
|
||||
"install_instructions": tool.install_instructions,
|
||||
"status": status.value,
|
||||
}
|
||||
|
||||
if status == ToolStatus.AVAILABLE:
|
||||
menu[cap]["available"].append(entry)
|
||||
menu[cap]["configured"] += 1
|
||||
else:
|
||||
menu[cap]["unavailable"].append(entry)
|
||||
menu[cap]["total"] += 1
|
||||
|
||||
return dict(sorted(menu.items()))
|
||||
|
||||
def gpu_required_tools(self) -> list[str]:
|
||||
"""List tools that require GPU (VRAM > 0)."""
|
||||
return [
|
||||
t.name for t in self._tools.values()
|
||||
if t.resource_profile.vram_mb > 0
|
||||
]
|
||||
|
||||
def network_required_tools(self) -> list[str]:
|
||||
"""List tools that require network access."""
|
||||
return [
|
||||
t.name for t in self._tools.values()
|
||||
if t.resource_profile.network_required
|
||||
]
|
||||
|
||||
|
||||
# Singleton registry instance
|
||||
registry = ToolRegistry()
|
||||
@@ -0,0 +1 @@
|
||||
"""Video tools — generation providers, composition, stitching, and trimming."""
|
||||
@@ -0,0 +1,576 @@
|
||||
"""Shared helpers for provider-specific video generation tools."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import ToolResult, ToolStatus
|
||||
|
||||
|
||||
HEYGEN_PROVIDERS = {
|
||||
"veo_3_1": {"name": "Google VEO 3.1", "quality": "highest", "speed": "slow"},
|
||||
"veo_3_1_fast": {"name": "Google VEO 3.1 Fast", "quality": "high", "speed": "medium"},
|
||||
"veo3": {"name": "Google VEO 3", "quality": "high", "speed": "slow"},
|
||||
"veo3_fast": {"name": "Google VEO 3 Fast", "quality": "high", "speed": "medium"},
|
||||
"veo2": {"name": "Google VEO 2", "quality": "medium", "speed": "medium"},
|
||||
"kling_pro": {"name": "Kling Pro", "quality": "high", "speed": "medium"},
|
||||
"kling_v2": {"name": "Kling v2", "quality": "medium", "speed": "fast"},
|
||||
"sora_v2": {"name": "Sora v2", "quality": "high", "speed": "slow"},
|
||||
"sora_v2_pro": {"name": "Sora v2 Pro", "quality": "highest", "speed": "slow"},
|
||||
"runway_gen4": {"name": "Runway Gen-4", "quality": "high", "speed": "medium"},
|
||||
"seedance_lite": {"name": "Seedance Lite", "quality": "medium", "speed": "fast"},
|
||||
"seedance_pro": {"name": "Seedance Pro", "quality": "high", "speed": "medium"},
|
||||
"ltx_distilled": {"name": "LTX Distilled", "quality": "low", "speed": "fastest"},
|
||||
}
|
||||
|
||||
WAN_VARIANTS = {
|
||||
"wan2.1-1.3b": {
|
||||
"name": "Wan 2.1 (1.3B)",
|
||||
"hf_id": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
||||
"hf_i2v_id": "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers",
|
||||
"pipeline_class": "WanPipeline",
|
||||
"vram_mb": 8000,
|
||||
"quality": "high",
|
||||
"speed": "medium",
|
||||
"t2v": True,
|
||||
"i2v": True,
|
||||
"license": "Apache-2.0",
|
||||
"default_width": 832,
|
||||
"default_height": 480,
|
||||
"default_num_frames": 81,
|
||||
"fps": 16,
|
||||
},
|
||||
"wan2.1-14b": {
|
||||
"name": "Wan 2.1 (14B)",
|
||||
"hf_id": "Wan-AI/Wan2.1-T2V-14B-Diffusers",
|
||||
"hf_i2v_id": "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers",
|
||||
"pipeline_class": "WanPipeline",
|
||||
"vram_mb": 24000,
|
||||
"quality": "highest",
|
||||
"speed": "slow",
|
||||
"t2v": True,
|
||||
"i2v": True,
|
||||
"license": "Apache-2.0",
|
||||
"default_width": 1280,
|
||||
"default_height": 720,
|
||||
"default_num_frames": 81,
|
||||
"fps": 16,
|
||||
},
|
||||
}
|
||||
|
||||
HUNYUAN_VARIANTS = {
|
||||
"hunyuan-1.5": {
|
||||
"name": "HunyuanVideo 1.5",
|
||||
"hf_id": "tencent/HunyuanVideo-1.5",
|
||||
"pipeline_class": "HunyuanVideoPipeline",
|
||||
"vram_mb": 14000,
|
||||
"quality": "high",
|
||||
"speed": "medium",
|
||||
"t2v": True,
|
||||
"i2v": True,
|
||||
"license": "Apache-2.0",
|
||||
"default_width": 848,
|
||||
"default_height": 480,
|
||||
"default_num_frames": 121,
|
||||
"fps": 24,
|
||||
},
|
||||
}
|
||||
|
||||
LTX_LOCAL_VARIANTS = {
|
||||
"ltx2-local": {
|
||||
"name": "LTX-2 (Local)",
|
||||
"hf_id": "Lightricks/LTX-2",
|
||||
"pipeline_class": "LTXPipeline",
|
||||
"vram_mb": 12000,
|
||||
"quality": "high",
|
||||
"speed": "medium",
|
||||
"t2v": True,
|
||||
"i2v": True,
|
||||
"license": "LTX-2-Community",
|
||||
"default_width": 768,
|
||||
"default_height": 512,
|
||||
"default_num_frames": 121,
|
||||
"fps": 24,
|
||||
},
|
||||
}
|
||||
|
||||
COGVIDEO_VARIANTS = {
|
||||
"cogvideo-5b": {
|
||||
"name": "CogVideoX 1.5 (5B)",
|
||||
"hf_id": "THUDM/CogVideoX-5b",
|
||||
"pipeline_class": "CogVideoXPipeline",
|
||||
"vram_mb": 12000,
|
||||
"quality": "medium",
|
||||
"speed": "medium",
|
||||
"t2v": True,
|
||||
"i2v": True,
|
||||
"license": "Apache-2.0",
|
||||
"default_width": 720,
|
||||
"default_height": 480,
|
||||
"default_num_frames": 49,
|
||||
"fps": 8,
|
||||
},
|
||||
"cogvideo-2b": {
|
||||
"name": "CogVideoX (2B)",
|
||||
"hf_id": "THUDM/CogVideoX-2b",
|
||||
"pipeline_class": "CogVideoXPipeline",
|
||||
"vram_mb": 6000,
|
||||
"quality": "medium",
|
||||
"speed": "fast",
|
||||
"t2v": True,
|
||||
"i2v": False,
|
||||
"license": "Apache-2.0",
|
||||
"default_width": 720,
|
||||
"default_height": 480,
|
||||
"default_num_frames": 49,
|
||||
"fps": 8,
|
||||
},
|
||||
}
|
||||
|
||||
LTX2_FRAME_COUNTS = {
|
||||
"1s": 25,
|
||||
"2s": 49,
|
||||
"3s": 73,
|
||||
"4s": 97,
|
||||
"5s": 121,
|
||||
"6.7s": 161,
|
||||
"8s": 193,
|
||||
}
|
||||
|
||||
|
||||
def local_generation_enabled() -> bool:
|
||||
return os.environ.get("VIDEO_GEN_LOCAL_ENABLED", "").lower() in {"true", "1", "yes"}
|
||||
|
||||
|
||||
def local_generation_status() -> ToolStatus:
|
||||
if not local_generation_enabled():
|
||||
return ToolStatus.UNAVAILABLE
|
||||
try:
|
||||
import diffusers # noqa: F401
|
||||
import torch # noqa: F401
|
||||
except ImportError:
|
||||
return ToolStatus.UNAVAILABLE
|
||||
return ToolStatus.AVAILABLE
|
||||
|
||||
|
||||
def local_install_instructions() -> str:
|
||||
return (
|
||||
"Enable local video generation and install the diffusers stack:\n"
|
||||
" set VIDEO_GEN_LOCAL_ENABLED=true\n"
|
||||
" pip install diffusers transformers accelerate torch pillow requests\n"
|
||||
"Use a GPU with the VRAM profile listed on the selected tool."
|
||||
)
|
||||
|
||||
|
||||
def estimate_quality_cost(quality: str) -> float:
|
||||
if quality == "highest":
|
||||
return 0.50
|
||||
if quality == "high":
|
||||
return 0.35
|
||||
if quality == "low":
|
||||
return 0.15
|
||||
return 0.20
|
||||
|
||||
|
||||
def estimate_speed_runtime(speed: str) -> float:
|
||||
return {"fastest": 30.0, "fast": 60.0, "medium": 120.0, "slow": 300.0}.get(speed, 120.0)
|
||||
|
||||
|
||||
def estimate_local_runtime(speed: str) -> float:
|
||||
return {"fast": 120.0, "medium": 240.0, "slow": 600.0}.get(speed, 240.0)
|
||||
|
||||
|
||||
def load_diffusers_pipeline(pipeline_class: str, model_id: str, enable_offload: bool):
|
||||
import diffusers
|
||||
import torch
|
||||
|
||||
pipeline_map = {
|
||||
"WanPipeline": "WanPipeline",
|
||||
"HunyuanVideoPipeline": "HunyuanVideoPipeline",
|
||||
"LTXPipeline": "LTXPipeline",
|
||||
"CogVideoXPipeline": "CogVideoXPipeline",
|
||||
}
|
||||
pipeline_name = pipeline_map.get(pipeline_class, pipeline_class)
|
||||
pipeline_class_obj = getattr(diffusers, pipeline_name)
|
||||
dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
|
||||
pipeline = pipeline_class_obj.from_pretrained(model_id, torch_dtype=dtype)
|
||||
|
||||
if enable_offload:
|
||||
pipeline.enable_model_cpu_offload()
|
||||
else:
|
||||
pipeline = pipeline.to("cuda")
|
||||
|
||||
if hasattr(pipeline, "vae") and pipeline.vae is not None:
|
||||
if hasattr(pipeline.vae, "enable_tiling"):
|
||||
pipeline.vae.enable_tiling()
|
||||
if hasattr(pipeline.vae, "enable_slicing"):
|
||||
pipeline.vae.enable_slicing()
|
||||
return pipeline
|
||||
|
||||
|
||||
def load_reference_image(inputs: dict[str, Any], width: int, height: int):
|
||||
from io import BytesIO
|
||||
|
||||
import requests
|
||||
from PIL import Image
|
||||
|
||||
ref_path = inputs.get("reference_image_path")
|
||||
ref_url = inputs.get("reference_image_url")
|
||||
|
||||
if ref_path:
|
||||
image = Image.open(ref_path).convert("RGB")
|
||||
elif ref_url:
|
||||
response = requests.get(ref_url, timeout=60)
|
||||
response.raise_for_status()
|
||||
image = Image.open(BytesIO(response.content)).convert("RGB")
|
||||
else:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="image_to_video requires reference_image_url or reference_image_path",
|
||||
)
|
||||
|
||||
return image.resize((width, height), Image.LANCZOS)
|
||||
|
||||
|
||||
def generate_local_video(
|
||||
*,
|
||||
tool_name: str,
|
||||
variants: dict[str, dict[str, Any]],
|
||||
default_variant: str,
|
||||
inputs: dict[str, Any],
|
||||
) -> ToolResult:
|
||||
import torch
|
||||
from diffusers.utils import export_to_video
|
||||
|
||||
variant = inputs.get("model_variant", default_variant)
|
||||
if variant not in variants:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"Unknown model_variant: {variant}. Available: {', '.join(sorted(variants))}",
|
||||
)
|
||||
|
||||
meta = variants[variant]
|
||||
prompt = inputs["prompt"]
|
||||
operation = inputs.get("operation", "text_to_video")
|
||||
seed = inputs.get("seed")
|
||||
enable_offload = inputs.get("enable_model_offload", True)
|
||||
|
||||
if operation == "image_to_video" and not meta.get("i2v"):
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"{meta['name']} does not support image_to_video.",
|
||||
)
|
||||
|
||||
width = inputs.get("width", meta["default_width"])
|
||||
height = inputs.get("height", meta["default_height"])
|
||||
num_frames = inputs.get("num_frames", meta["default_num_frames"])
|
||||
fps = meta["fps"]
|
||||
model_id = meta.get("hf_i2v_id") if operation == "image_to_video" and meta.get("hf_i2v_id") else meta["hf_id"]
|
||||
pipeline = load_diffusers_pipeline(meta["pipeline_class"], model_id, enable_offload)
|
||||
|
||||
generation_args: dict[str, Any] = {
|
||||
"prompt": prompt,
|
||||
"num_frames": num_frames,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"num_inference_steps": inputs.get("num_inference_steps", 30),
|
||||
}
|
||||
if seed is not None:
|
||||
generation_args["generator"] = torch.Generator(device="cpu").manual_seed(seed)
|
||||
if operation == "image_to_video":
|
||||
image = load_reference_image(inputs, width, height)
|
||||
if isinstance(image, ToolResult):
|
||||
return image
|
||||
generation_args["image"] = image
|
||||
if meta["pipeline_class"] == "CogVideoXPipeline":
|
||||
generation_args["negative_prompt"] = "worst quality, low quality, blurry, distorted, watermark"
|
||||
|
||||
output = pipeline(**generation_args)
|
||||
frames = output.frames[0] if hasattr(output, "frames") else output.images
|
||||
|
||||
output_path = Path(inputs.get("output_path", f"{tool_name}_{variant}.mp4"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
export_to_video(frames, str(output_path), fps=fps)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": tool_name,
|
||||
"model_variant": variant,
|
||||
"provider_name": meta["name"],
|
||||
"mode": "local",
|
||||
"prompt": prompt,
|
||||
"model_id": model_id,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"num_frames": num_frames,
|
||||
"fps": fps,
|
||||
"duration_seconds": round(num_frames / fps, 2),
|
||||
"operation": operation,
|
||||
"output": str(output_path),
|
||||
"format": "mp4",
|
||||
"license": meta["license"],
|
||||
**probe_output(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
seed=seed,
|
||||
model=model_id,
|
||||
)
|
||||
|
||||
|
||||
def poll_heygen(execution_id: str, api_key: str, timeout: int = 600) -> str:
|
||||
import requests
|
||||
|
||||
headers = {"X-Api-Key": api_key}
|
||||
url = f"https://api.heygen.com/v1/workflows/executions/{execution_id}"
|
||||
deadline = time.time() + timeout
|
||||
interval = 5.0
|
||||
|
||||
while time.time() < deadline:
|
||||
response = requests.get(url, headers=headers, timeout=30)
|
||||
response.raise_for_status()
|
||||
data = response.json().get("data", {})
|
||||
status = data.get("status", "")
|
||||
|
||||
if status == "completed":
|
||||
video_url = (
|
||||
data.get("output", {}).get("video", {}).get("video_url")
|
||||
or data.get("output", {}).get("video_url")
|
||||
)
|
||||
if video_url:
|
||||
return video_url
|
||||
raise RuntimeError(f"Completed but no video_url in output: {data}")
|
||||
|
||||
if status in {"failed", "error"}:
|
||||
raise RuntimeError(f"HeyGen generation failed: {data.get('error', 'Unknown')}")
|
||||
|
||||
time.sleep(min(interval, max(0.0, deadline - time.time())))
|
||||
interval = min(interval * 1.2, 30.0)
|
||||
|
||||
raise TimeoutError(f"HeyGen execution {execution_id} timed out after {timeout}s")
|
||||
|
||||
|
||||
def upload_image_heygen(image_path: str, api_key: str) -> str:
|
||||
import requests
|
||||
|
||||
path = Path(image_path)
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"Image not found: {image_path}")
|
||||
|
||||
with path.open("rb") as handle:
|
||||
response = requests.post(
|
||||
"https://api.heygen.com/v1/asset",
|
||||
headers={"X-Api-Key": api_key},
|
||||
files={"file": (path.name, handle, "image/png")},
|
||||
timeout=60,
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json().get("data", {}).get("url", "")
|
||||
|
||||
|
||||
def generate_heygen_video(inputs: dict[str, Any]) -> ToolResult:
|
||||
import requests
|
||||
|
||||
api_key = os.environ.get("HEYGEN_API_KEY")
|
||||
if not api_key:
|
||||
return ToolResult(success=False, error="HEYGEN_API_KEY not set.")
|
||||
|
||||
provider = inputs.get("provider_variant", "veo_3_1")
|
||||
if provider not in HEYGEN_PROVIDERS:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"Unknown provider_variant: {provider}. Available: {', '.join(sorted(HEYGEN_PROVIDERS))}",
|
||||
)
|
||||
|
||||
prompt = inputs["prompt"]
|
||||
aspect_ratio = inputs.get("aspect_ratio", "16:9")
|
||||
operation = inputs.get("operation", "text_to_video")
|
||||
workflow_input: dict[str, Any] = {
|
||||
"prompt": prompt,
|
||||
"provider": provider,
|
||||
"aspect_ratio": aspect_ratio,
|
||||
}
|
||||
if operation == "image_to_video":
|
||||
ref_url = inputs.get("reference_image_url")
|
||||
ref_path = inputs.get("reference_image_path")
|
||||
if ref_path and not ref_url:
|
||||
ref_url = upload_image_heygen(ref_path, api_key)
|
||||
if not ref_url:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="image_to_video requires reference_image_url or reference_image_path",
|
||||
)
|
||||
workflow_input["reference_image_url"] = ref_url
|
||||
|
||||
response = requests.post(
|
||||
"https://api.heygen.com/v1/workflows/executions",
|
||||
headers={"X-Api-Key": api_key, "Content-Type": "application/json"},
|
||||
json={"workflow_type": "GenerateVideoNode", "input": workflow_input},
|
||||
timeout=30,
|
||||
)
|
||||
response.raise_for_status()
|
||||
payload = response.json()
|
||||
execution_id = payload.get("data", {}).get("execution_id")
|
||||
if not execution_id:
|
||||
return ToolResult(success=False, error=f"No execution_id in response: {payload}")
|
||||
|
||||
video_url = poll_heygen(execution_id, api_key, timeout=600)
|
||||
output_path = Path(inputs.get("output_path", f"heygen_video_{execution_id}.mp4"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
download = requests.get(video_url, timeout=120)
|
||||
download.raise_for_status()
|
||||
output_path.write_bytes(download.content)
|
||||
|
||||
meta = HEYGEN_PROVIDERS[provider]
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "heygen",
|
||||
"provider_variant": provider,
|
||||
"provider_name": meta["name"],
|
||||
"mode": "api",
|
||||
"prompt": prompt,
|
||||
"aspect_ratio": aspect_ratio,
|
||||
"operation": operation,
|
||||
"execution_id": execution_id,
|
||||
"output": str(output_path),
|
||||
"format": "mp4",
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
model=provider,
|
||||
)
|
||||
|
||||
|
||||
def generate_ltx_modal_video(inputs: dict[str, Any]) -> ToolResult:
|
||||
import base64
|
||||
|
||||
import requests
|
||||
|
||||
endpoint_url = os.environ.get("MODAL_LTX2_ENDPOINT_URL")
|
||||
if not endpoint_url:
|
||||
return ToolResult(success=False, error="MODAL_LTX2_ENDPOINT_URL not set.")
|
||||
|
||||
prompt = inputs["prompt"]
|
||||
operation = inputs.get("operation", "text_to_video")
|
||||
aspect = inputs.get("aspect_ratio", "16:9")
|
||||
width = inputs.get("width")
|
||||
height = inputs.get("height")
|
||||
if width is None or height is None:
|
||||
if aspect == "16:9":
|
||||
width, height = 1024, 576
|
||||
elif aspect == "9:16":
|
||||
width, height = 576, 1024
|
||||
else:
|
||||
width, height = 512, 512
|
||||
|
||||
num_frames = inputs.get("num_frames", LTX2_FRAME_COUNTS.get(inputs.get("duration_hint", "5s"), 121))
|
||||
if (num_frames - 1) % 8 != 0:
|
||||
num_frames = ((num_frames - 1) // 8) * 8 + 1
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"prompt": prompt,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"num_frames": num_frames,
|
||||
"fps": 24,
|
||||
"steps": inputs.get("num_inference_steps", 30),
|
||||
"negative_prompt": "worst quality, low quality, blurry, distorted, watermark, text, logo",
|
||||
}
|
||||
if inputs.get("seed") is not None:
|
||||
payload["seed"] = inputs["seed"]
|
||||
|
||||
if operation == "image_to_video":
|
||||
ref_path = inputs.get("reference_image_path")
|
||||
ref_url = inputs.get("reference_image_url")
|
||||
if ref_path:
|
||||
payload["input_image"] = base64.b64encode(Path(ref_path).read_bytes()).decode()
|
||||
elif ref_url:
|
||||
payload["input_image_url"] = ref_url
|
||||
else:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="image_to_video requires reference_image_url or reference_image_path",
|
||||
)
|
||||
|
||||
response = requests.post(endpoint_url, json=payload, timeout=300)
|
||||
response.raise_for_status()
|
||||
output_path = Path(inputs.get("output_path", "ltx_video_modal.mp4"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
content_type = response.headers.get("content-type", "")
|
||||
if "video" in content_type or "octet-stream" in content_type:
|
||||
output_path.write_bytes(response.content)
|
||||
else:
|
||||
response_payload = response.json()
|
||||
video_url = response_payload.get("video_url") or response_payload.get("url")
|
||||
if not video_url:
|
||||
return ToolResult(success=False, error=f"No video data in response: {response_payload}")
|
||||
download = requests.get(video_url, timeout=120)
|
||||
download.raise_for_status()
|
||||
output_path.write_bytes(download.content)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "ltx-modal",
|
||||
"provider_name": "LTX-2 (Modal)",
|
||||
"mode": "modal",
|
||||
"prompt": prompt,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"num_frames": num_frames,
|
||||
"fps": 24,
|
||||
"duration_seconds": round(num_frames / 24, 2),
|
||||
"operation": operation,
|
||||
"output": str(output_path),
|
||||
"format": "mp4",
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
seed=inputs.get("seed"),
|
||||
model="ltx-2",
|
||||
)
|
||||
|
||||
|
||||
def probe_output(path: Path) -> dict[str, Any]:
|
||||
info: dict[str, Any] = {"file_size_bytes": path.stat().st_size}
|
||||
if not shutil.which("ffprobe"):
|
||||
return info
|
||||
|
||||
import json
|
||||
|
||||
try:
|
||||
proc = subprocess.run(
|
||||
[
|
||||
"ffprobe",
|
||||
"-v",
|
||||
"quiet",
|
||||
"-print_format",
|
||||
"json",
|
||||
"-show_format",
|
||||
"-show_streams",
|
||||
str(path),
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=10,
|
||||
check=False,
|
||||
)
|
||||
if proc.returncode == 0:
|
||||
probe = json.loads(proc.stdout)
|
||||
fmt = probe.get("format", {})
|
||||
info["duration_seconds"] = float(fmt.get("duration", 0))
|
||||
info["file_size_mb"] = round(path.stat().st_size / (1024 * 1024), 2)
|
||||
for stream in probe.get("streams", []):
|
||||
if stream.get("codec_type") == "video":
|
||||
info["video_width"] = int(stream.get("width", 0))
|
||||
info["video_height"] = int(stream.get("height", 0))
|
||||
info["video_codec"] = stream.get("codec_name", "")
|
||||
break
|
||||
except Exception:
|
||||
pass
|
||||
return info
|
||||
@@ -0,0 +1,97 @@
|
||||
"""CogVideo local video generation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
from tools.video._shared import COGVIDEO_VARIANTS, estimate_local_runtime, generate_local_video, local_generation_status, local_install_instructions
|
||||
|
||||
|
||||
class CogVideoVideo(BaseTool):
|
||||
name = "cogvideo_video"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "video_generation"
|
||||
provider = "cogvideo"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.STOCHASTIC
|
||||
runtime = ToolRuntime.LOCAL_GPU
|
||||
|
||||
install_instructions = local_install_instructions()
|
||||
fallback = "wan_video"
|
||||
fallback_tools = ["wan_video", "hunyuan_video", "ltx_video_local", "image_selector"]
|
||||
agent_skills = ["ltx2"]
|
||||
|
||||
capabilities = ["text_to_video", "image_to_video", "model_selection"]
|
||||
supports = {
|
||||
"reference_image": True,
|
||||
"offline": True,
|
||||
"native_audio": False,
|
||||
"local_gpu": True,
|
||||
}
|
||||
best_for = [
|
||||
"lower-VRAM local video experimentation",
|
||||
"teams that want an explicit CogVideo family path in the registry",
|
||||
]
|
||||
not_good_for = ["best-in-class local quality targets"]
|
||||
provider_matrix = {key: {"tool": "cogvideo_video", **value, "mode": "local_gpu"} for key, value in COGVIDEO_VARIANTS.items()}
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {"type": "string"},
|
||||
"operation": {"type": "string", "enum": ["text_to_video", "image_to_video"], "default": "text_to_video"},
|
||||
"model_variant": {"type": "string", "enum": sorted(COGVIDEO_VARIANTS), "default": "cogvideo-5b"},
|
||||
"reference_image_url": {"type": "string"},
|
||||
"reference_image_path": {"type": "string"},
|
||||
"width": {"type": "integer"},
|
||||
"height": {"type": "integer"},
|
||||
"num_frames": {"type": "integer"},
|
||||
"num_inference_steps": {"type": "integer"},
|
||||
"enable_model_offload": {"type": "boolean", "default": True},
|
||||
"seed": {"type": "integer"},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(cpu_cores=2, ram_mb=16000, vram_mb=6000, disk_mb=4000, network_required=False)
|
||||
retry_policy = RetryPolicy(max_retries=1)
|
||||
idempotency_key_fields = ["prompt", "model_variant", "operation", "seed"]
|
||||
side_effects = ["writes video file to output_path", "may download model weights"]
|
||||
user_visible_verification = ["Watch generated clip for motion coherence and artifacts"]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
return local_generation_status()
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, object]) -> float:
|
||||
return 0.0
|
||||
|
||||
def estimate_runtime(self, inputs: dict[str, object]) -> float:
|
||||
variant = COGVIDEO_VARIANTS.get(inputs.get("model_variant", "cogvideo-5b"), COGVIDEO_VARIANTS["cogvideo-5b"])
|
||||
return estimate_local_runtime(variant["speed"])
|
||||
|
||||
def execute(self, inputs: dict[str, object]) -> ToolResult:
|
||||
if self.get_status() != ToolStatus.AVAILABLE:
|
||||
return ToolResult(success=False, error="CogVideo local generation is unavailable. " + self.install_instructions)
|
||||
start = time.time()
|
||||
try:
|
||||
result = generate_local_video(tool_name=self.name, variants=COGVIDEO_VARIANTS, default_variant="cogvideo-5b", inputs=inputs)
|
||||
except Exception as exc:
|
||||
return ToolResult(success=False, error=f"CogVideo generation failed: {exc}")
|
||||
result.duration_seconds = round(time.time() - start, 2)
|
||||
return result
|
||||
|
||||
@@ -0,0 +1,117 @@
|
||||
"""HeyGen-backed cloud video generation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
from tools.video._shared import HEYGEN_PROVIDERS, estimate_quality_cost, estimate_speed_runtime, generate_heygen_video
|
||||
|
||||
|
||||
class HeyGenVideo(BaseTool):
|
||||
name = "heygen_video"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "video_generation"
|
||||
provider = "heygen"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.STOCHASTIC
|
||||
runtime = ToolRuntime.API
|
||||
|
||||
install_instructions = (
|
||||
"Set the HEYGEN_API_KEY environment variable:\n"
|
||||
" set HEYGEN_API_KEY=your_key_here\n"
|
||||
"Get a key at https://app.heygen.com/settings/api"
|
||||
)
|
||||
fallback = "wan_video"
|
||||
fallback_tools = ["wan_video", "hunyuan_video", "ltx_video_local", "cogvideo_video", "ltx_video_modal", "image_selector"]
|
||||
agent_skills = ["ai-video-gen", "create-video"]
|
||||
|
||||
capabilities = ["text_to_video", "image_to_video", "provider_selection"]
|
||||
supports = {
|
||||
"reference_image": True,
|
||||
"offline": False,
|
||||
"native_audio": False,
|
||||
"cloud_generation": True,
|
||||
}
|
||||
best_for = [
|
||||
"premium cloud video generation without local GPU setup",
|
||||
"fast access to VEO, Sora, Kling, Runway, and Seedance providers",
|
||||
]
|
||||
not_good_for = [
|
||||
"offline or privacy-constrained rendering",
|
||||
"free local-first production",
|
||||
]
|
||||
provider_matrix = {
|
||||
key: {"tool": "heygen_video", **value, "mode": "api"} for key, value in HEYGEN_PROVIDERS.items()
|
||||
}
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {"type": "string"},
|
||||
"operation": {
|
||||
"type": "string",
|
||||
"enum": ["text_to_video", "image_to_video"],
|
||||
"default": "text_to_video",
|
||||
},
|
||||
"provider_variant": {
|
||||
"type": "string",
|
||||
"enum": sorted(HEYGEN_PROVIDERS),
|
||||
"default": "veo_3_1",
|
||||
},
|
||||
"reference_image_url": {"type": "string"},
|
||||
"reference_image_path": {"type": "string"},
|
||||
"aspect_ratio": {
|
||||
"type": "string",
|
||||
"enum": ["16:9", "9:16", "1:1"],
|
||||
"default": "16:9",
|
||||
},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(cpu_cores=1, ram_mb=512, vram_mb=0, disk_mb=500, network_required=True)
|
||||
retry_policy = RetryPolicy(max_retries=2, backoff_seconds=10.0, retryable_errors=["rate_limit", "timeout", "server_error"])
|
||||
idempotency_key_fields = ["prompt", "provider_variant", "aspect_ratio"]
|
||||
side_effects = ["writes video file to output_path", "calls HeyGen API"]
|
||||
user_visible_verification = ["Watch generated clip for motion quality and prompt adherence"]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
return ToolStatus.AVAILABLE if os.environ.get("HEYGEN_API_KEY") else ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
meta = HEYGEN_PROVIDERS.get(inputs.get("provider_variant", "veo_3_1"), HEYGEN_PROVIDERS["veo_3_1"])
|
||||
return estimate_quality_cost(meta["quality"])
|
||||
|
||||
def estimate_runtime(self, inputs: dict[str, Any]) -> float:
|
||||
meta = HEYGEN_PROVIDERS.get(inputs.get("provider_variant", "veo_3_1"), HEYGEN_PROVIDERS["veo_3_1"])
|
||||
return estimate_speed_runtime(meta["speed"])
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
if self.get_status() != ToolStatus.AVAILABLE:
|
||||
return ToolResult(success=False, error="HeyGen video generation is unavailable. " + self.install_instructions)
|
||||
start = time.time()
|
||||
try:
|
||||
result = generate_heygen_video(inputs)
|
||||
except Exception as exc:
|
||||
return ToolResult(success=False, error=f"HeyGen video generation failed: {exc}")
|
||||
result.duration_seconds = round(time.time() - start, 2)
|
||||
result.cost_usd = self.estimate_cost(inputs)
|
||||
return result
|
||||
|
||||
@@ -0,0 +1,95 @@
|
||||
"""Hunyuan local video generation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
from tools.video._shared import HUNYUAN_VARIANTS, estimate_local_runtime, generate_local_video, local_generation_status, local_install_instructions
|
||||
|
||||
|
||||
class HunyuanVideo(BaseTool):
|
||||
name = "hunyuan_video"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "video_generation"
|
||||
provider = "hunyuan"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.STOCHASTIC
|
||||
runtime = ToolRuntime.LOCAL_GPU
|
||||
|
||||
install_instructions = local_install_instructions()
|
||||
fallback = "wan_video"
|
||||
fallback_tools = ["wan_video", "ltx_video_local", "cogvideo_video", "image_selector"]
|
||||
agent_skills = ["ltx2"]
|
||||
|
||||
capabilities = ["text_to_video", "image_to_video"]
|
||||
supports = {
|
||||
"reference_image": True,
|
||||
"offline": True,
|
||||
"native_audio": False,
|
||||
"local_gpu": True,
|
||||
}
|
||||
best_for = [
|
||||
"local generation when Hunyuan motion behavior fits the brief",
|
||||
"teams that want one known Hunyuan baseline instead of multiple variants",
|
||||
]
|
||||
not_good_for = ["CPU-only machines"]
|
||||
provider_matrix = {key: {"tool": "hunyuan_video", **value, "mode": "local_gpu"} for key, value in HUNYUAN_VARIANTS.items()}
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {"type": "string"},
|
||||
"operation": {"type": "string", "enum": ["text_to_video", "image_to_video"], "default": "text_to_video"},
|
||||
"model_variant": {"type": "string", "enum": ["hunyuan-1.5"], "default": "hunyuan-1.5"},
|
||||
"reference_image_url": {"type": "string"},
|
||||
"reference_image_path": {"type": "string"},
|
||||
"width": {"type": "integer"},
|
||||
"height": {"type": "integer"},
|
||||
"num_frames": {"type": "integer"},
|
||||
"num_inference_steps": {"type": "integer"},
|
||||
"enable_model_offload": {"type": "boolean", "default": True},
|
||||
"seed": {"type": "integer"},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(cpu_cores=2, ram_mb=16000, vram_mb=14000, disk_mb=4000, network_required=False)
|
||||
retry_policy = RetryPolicy(max_retries=1)
|
||||
idempotency_key_fields = ["prompt", "model_variant", "operation", "seed"]
|
||||
side_effects = ["writes video file to output_path", "may download model weights"]
|
||||
user_visible_verification = ["Watch generated clip for motion coherence and artifacts"]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
return local_generation_status()
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
return 0.0
|
||||
|
||||
def estimate_runtime(self, inputs: dict[str, Any]) -> float:
|
||||
return estimate_local_runtime(HUNYUAN_VARIANTS["hunyuan-1.5"]["speed"])
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
if self.get_status() != ToolStatus.AVAILABLE:
|
||||
return ToolResult(success=False, error="Hunyuan local video generation is unavailable. " + self.install_instructions)
|
||||
start = time.time()
|
||||
try:
|
||||
result = generate_local_video(tool_name=self.name, variants=HUNYUAN_VARIANTS, default_variant="hunyuan-1.5", inputs=inputs)
|
||||
except Exception as exc:
|
||||
return ToolResult(success=False, error=f"Hunyuan video generation failed: {exc}")
|
||||
result.duration_seconds = round(time.time() - start, 2)
|
||||
return result
|
||||
@@ -0,0 +1,178 @@
|
||||
"""Kling video generation via fal.ai API.
|
||||
|
||||
Best for cinematic B-roll with high visual fidelity and fluid motion.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class KlingVideo(BaseTool):
|
||||
name = "kling_video"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "video_generation"
|
||||
provider = "kling"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.STOCHASTIC
|
||||
runtime = ToolRuntime.API
|
||||
|
||||
dependencies = []
|
||||
install_instructions = (
|
||||
"Set FAL_KEY to your fal.ai API key.\n"
|
||||
" Get one at https://fal.ai/dashboard/keys"
|
||||
)
|
||||
agent_skills = ["ai-video-gen"]
|
||||
|
||||
capabilities = ["text_to_video", "image_to_video"]
|
||||
supports = {
|
||||
"text_to_video": True,
|
||||
"image_to_video": True,
|
||||
"native_audio": True,
|
||||
"cinematic_quality": True,
|
||||
}
|
||||
best_for = [
|
||||
"cinematic B-roll with highest visual fidelity",
|
||||
"fluid motion and camera direction",
|
||||
"professional video clips",
|
||||
]
|
||||
not_good_for = ["budget-constrained projects", "offline generation", "quick iteration"]
|
||||
fallback_tools = ["minimax_video", "veo_video", "wan_video"]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {"type": "string"},
|
||||
"operation": {
|
||||
"type": "string",
|
||||
"enum": ["text_to_video", "image_to_video"],
|
||||
"default": "text_to_video",
|
||||
},
|
||||
"model_variant": {
|
||||
"type": "string",
|
||||
"enum": ["v3/standard", "v2.1/master", "v2.1/pro", "v2.1/standard"],
|
||||
"default": "v3/standard",
|
||||
},
|
||||
"duration": {
|
||||
"type": "string",
|
||||
"enum": ["5", "10"],
|
||||
"default": "5",
|
||||
"description": "Duration in seconds",
|
||||
},
|
||||
"aspect_ratio": {
|
||||
"type": "string",
|
||||
"enum": ["16:9", "9:16", "1:1"],
|
||||
"default": "16:9",
|
||||
},
|
||||
"image_url": {"type": "string", "description": "Reference image URL for image_to_video"},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=1, ram_mb=512, vram_mb=0, disk_mb=500, network_required=True
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=2, retryable_errors=["rate_limit", "timeout"])
|
||||
idempotency_key_fields = ["prompt", "model_variant", "operation", "duration"]
|
||||
side_effects = ["writes video file to output_path", "calls fal.ai API"]
|
||||
user_visible_verification = ["Watch generated clip for motion coherence and visual quality"]
|
||||
|
||||
def _get_api_key(self) -> str | None:
|
||||
return os.environ.get("FAL_KEY") or os.environ.get("FAL_AI_API_KEY")
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if self._get_api_key():
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
variant = inputs.get("model_variant", "v3/standard")
|
||||
duration = int(inputs.get("duration", "5"))
|
||||
if "master" in variant:
|
||||
return 0.30 * (duration / 5)
|
||||
if "pro" in variant:
|
||||
return 0.20 * (duration / 5)
|
||||
return 0.10 * (duration / 5) # standard
|
||||
|
||||
def estimate_runtime(self, inputs: dict[str, Any]) -> float:
|
||||
return 60.0 # ~1 minute typical
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
api_key = self._get_api_key()
|
||||
if not api_key:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="FAL_KEY not set. " + self.install_instructions,
|
||||
)
|
||||
|
||||
import requests
|
||||
|
||||
start = time.time()
|
||||
operation = inputs.get("operation", "text_to_video")
|
||||
variant = inputs.get("model_variant", "v3/standard")
|
||||
model_path = f"kling-video/{variant}/{operation}"
|
||||
|
||||
payload: dict[str, Any] = {"prompt": inputs["prompt"]}
|
||||
if inputs.get("duration"):
|
||||
payload["duration"] = inputs["duration"]
|
||||
if inputs.get("aspect_ratio"):
|
||||
payload["aspect_ratio"] = inputs["aspect_ratio"]
|
||||
if operation == "image_to_video" and inputs.get("image_url"):
|
||||
payload["image_url"] = inputs["image_url"]
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
f"https://fal.run/fal-ai/{model_path}",
|
||||
headers={
|
||||
"Authorization": f"Key {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
json=payload,
|
||||
timeout=300,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
video_url = data["video"]["url"]
|
||||
video_response = requests.get(video_url, timeout=120)
|
||||
video_response.raise_for_status()
|
||||
|
||||
output_path = Path(inputs.get("output_path", "kling_output.mp4"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(video_response.content)
|
||||
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Kling video generation failed: {e}")
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "kling",
|
||||
"model": f"fal-ai/{model_path}",
|
||||
"prompt": inputs["prompt"],
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
cost_usd=self.estimate_cost(inputs),
|
||||
duration_seconds=round(time.time() - start, 2),
|
||||
model=f"fal-ai/{model_path}",
|
||||
)
|
||||
@@ -0,0 +1,96 @@
|
||||
"""LTX local video generation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
from tools.video._shared import LTX_LOCAL_VARIANTS, estimate_local_runtime, generate_local_video, local_generation_status, local_install_instructions
|
||||
|
||||
|
||||
class LTXVideoLocal(BaseTool):
|
||||
name = "ltx_video_local"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "video_generation"
|
||||
provider = "ltx"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.STOCHASTIC
|
||||
runtime = ToolRuntime.LOCAL_GPU
|
||||
|
||||
install_instructions = local_install_instructions()
|
||||
fallback = "wan_video"
|
||||
fallback_tools = ["wan_video", "hunyuan_video", "cogvideo_video", "ltx_video_modal", "image_selector"]
|
||||
agent_skills = ["ltx2"]
|
||||
|
||||
capabilities = ["text_to_video", "image_to_video"]
|
||||
supports = {
|
||||
"reference_image": True,
|
||||
"offline": True,
|
||||
"native_audio": False,
|
||||
"local_gpu": True,
|
||||
}
|
||||
best_for = [
|
||||
"local LTX workflows already tuned around LTX prompting",
|
||||
"teams that want one dedicated LTX local path in the registry",
|
||||
]
|
||||
not_good_for = ["CPU-only machines"]
|
||||
provider_matrix = {key: {"tool": "ltx_video_local", **value, "mode": "local_gpu"} for key, value in LTX_LOCAL_VARIANTS.items()}
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {"type": "string"},
|
||||
"operation": {"type": "string", "enum": ["text_to_video", "image_to_video"], "default": "text_to_video"},
|
||||
"model_variant": {"type": "string", "enum": ["ltx2-local"], "default": "ltx2-local"},
|
||||
"reference_image_url": {"type": "string"},
|
||||
"reference_image_path": {"type": "string"},
|
||||
"width": {"type": "integer"},
|
||||
"height": {"type": "integer"},
|
||||
"num_frames": {"type": "integer"},
|
||||
"num_inference_steps": {"type": "integer"},
|
||||
"enable_model_offload": {"type": "boolean", "default": True},
|
||||
"seed": {"type": "integer"},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(cpu_cores=2, ram_mb=16000, vram_mb=12000, disk_mb=4000, network_required=False)
|
||||
retry_policy = RetryPolicy(max_retries=1)
|
||||
idempotency_key_fields = ["prompt", "model_variant", "operation", "seed"]
|
||||
side_effects = ["writes video file to output_path", "may download model weights"]
|
||||
user_visible_verification = ["Watch generated clip for motion coherence and artifacts"]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
return local_generation_status()
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, object]) -> float:
|
||||
return 0.0
|
||||
|
||||
def estimate_runtime(self, inputs: dict[str, object]) -> float:
|
||||
return estimate_local_runtime(LTX_LOCAL_VARIANTS["ltx2-local"]["speed"])
|
||||
|
||||
def execute(self, inputs: dict[str, object]) -> ToolResult:
|
||||
if self.get_status() != ToolStatus.AVAILABLE:
|
||||
return ToolResult(success=False, error="Local LTX video generation is unavailable. " + self.install_instructions)
|
||||
start = time.time()
|
||||
try:
|
||||
result = generate_local_video(tool_name=self.name, variants=LTX_LOCAL_VARIANTS, default_variant="ltx2-local", inputs=inputs)
|
||||
except Exception as exc:
|
||||
return ToolResult(success=False, error=f"Local LTX video generation failed: {exc}")
|
||||
result.duration_seconds = round(time.time() - start, 2)
|
||||
return result
|
||||
|
||||
@@ -0,0 +1,106 @@
|
||||
"""Modal-hosted LTX video generation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
from tools.video._shared import generate_ltx_modal_video
|
||||
|
||||
|
||||
class LTXVideoModal(BaseTool):
|
||||
name = "ltx_video_modal"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "video_generation"
|
||||
provider = "ltx-modal"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.STOCHASTIC
|
||||
runtime = ToolRuntime.API
|
||||
|
||||
install_instructions = (
|
||||
"Set the MODAL_LTX2_ENDPOINT_URL environment variable to your deployed LTX endpoint:\n"
|
||||
" set MODAL_LTX2_ENDPOINT_URL=https://<your-modal-endpoint>"
|
||||
)
|
||||
fallback = "ltx_video_local"
|
||||
fallback_tools = ["ltx_video_local", "wan_video", "hunyuan_video", "cogvideo_video", "image_selector"]
|
||||
agent_skills = ["ltx2"]
|
||||
|
||||
capabilities = ["text_to_video", "image_to_video"]
|
||||
supports = {
|
||||
"reference_image": True,
|
||||
"offline": False,
|
||||
"native_audio": False,
|
||||
"self_hosted_cloud": True,
|
||||
}
|
||||
best_for = ["self-hosted cloud GPU rendering for LTX without local workstation dependence"]
|
||||
not_good_for = ["zero-setup local workflows"]
|
||||
provider_matrix = {
|
||||
"ltx2-modal": {
|
||||
"tool": "ltx_video_modal",
|
||||
"name": "LTX-2 (Modal)",
|
||||
"mode": "api",
|
||||
"quality": "high",
|
||||
"speed": "medium",
|
||||
}
|
||||
}
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {"type": "string"},
|
||||
"operation": {"type": "string", "enum": ["text_to_video", "image_to_video"], "default": "text_to_video"},
|
||||
"reference_image_url": {"type": "string"},
|
||||
"reference_image_path": {"type": "string"},
|
||||
"aspect_ratio": {"type": "string", "enum": ["16:9", "9:16", "1:1"], "default": "16:9"},
|
||||
"duration_hint": {"type": "string"},
|
||||
"width": {"type": "integer"},
|
||||
"height": {"type": "integer"},
|
||||
"num_frames": {"type": "integer"},
|
||||
"num_inference_steps": {"type": "integer"},
|
||||
"seed": {"type": "integer"},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(cpu_cores=1, ram_mb=512, vram_mb=0, disk_mb=500, network_required=True)
|
||||
retry_policy = RetryPolicy(max_retries=2, backoff_seconds=10.0, retryable_errors=["timeout", "server_error"])
|
||||
idempotency_key_fields = ["prompt", "aspect_ratio", "num_frames", "seed"]
|
||||
side_effects = ["writes video file to output_path", "calls modal endpoint"]
|
||||
user_visible_verification = ["Watch generated clip for motion quality and prompt adherence"]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
return ToolStatus.AVAILABLE if os.environ.get("MODAL_LTX2_ENDPOINT_URL") else ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, object]) -> float:
|
||||
return 0.25
|
||||
|
||||
def estimate_runtime(self, inputs: dict[str, object]) -> float:
|
||||
return 180.0
|
||||
|
||||
def execute(self, inputs: dict[str, object]) -> ToolResult:
|
||||
if self.get_status() != ToolStatus.AVAILABLE:
|
||||
return ToolResult(success=False, error="Modal LTX video generation is unavailable. " + self.install_instructions)
|
||||
start = time.time()
|
||||
try:
|
||||
result = generate_ltx_modal_video(inputs)
|
||||
except Exception as exc:
|
||||
return ToolResult(success=False, error=f"Modal LTX video generation failed: {exc}")
|
||||
result.duration_seconds = round(time.time() - start, 2)
|
||||
result.cost_usd = self.estimate_cost(inputs)
|
||||
return result
|
||||
|
||||
@@ -0,0 +1,176 @@
|
||||
"""MiniMax (Hailuo AI) video generation via fal.ai API.
|
||||
|
||||
Rewards prompt craft — follows camera directions well and produces high-texture footage.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class MiniMaxVideo(BaseTool):
|
||||
name = "minimax_video"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "video_generation"
|
||||
provider = "minimax"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.STOCHASTIC
|
||||
runtime = ToolRuntime.API
|
||||
|
||||
dependencies = []
|
||||
install_instructions = (
|
||||
"Set FAL_KEY to your fal.ai API key.\n"
|
||||
" Get one at https://fal.ai/dashboard/keys"
|
||||
)
|
||||
agent_skills = ["ai-video-gen"]
|
||||
|
||||
capabilities = ["text_to_video", "image_to_video"]
|
||||
supports = {
|
||||
"text_to_video": True,
|
||||
"image_to_video": True,
|
||||
"camera_direction": True,
|
||||
}
|
||||
best_for = [
|
||||
"prompt-following with camera directions (framing, motion, composition)",
|
||||
"high-texture footage with minimal hallucination",
|
||||
"cost-effective video generation",
|
||||
]
|
||||
not_good_for = ["offline generation", "very long clips"]
|
||||
fallback_tools = ["kling_video", "veo_video", "wan_video"]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {"type": "string"},
|
||||
"operation": {
|
||||
"type": "string",
|
||||
"enum": ["text_to_video", "image_to_video"],
|
||||
"default": "text_to_video",
|
||||
},
|
||||
"model_variant": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"video-01", "hailuo-02/pro", "hailuo-02/standard",
|
||||
"hailuo-2.3-fast/pro", "hailuo-2.3-fast/standard",
|
||||
],
|
||||
"default": "hailuo-02/pro",
|
||||
},
|
||||
"image_url": {"type": "string", "description": "Reference image URL for image_to_video"},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=1, ram_mb=512, vram_mb=0, disk_mb=500, network_required=True
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=2, retryable_errors=["rate_limit", "timeout"])
|
||||
idempotency_key_fields = ["prompt", "model_variant", "operation"]
|
||||
side_effects = ["writes video file to output_path", "calls fal.ai API"]
|
||||
user_visible_verification = ["Watch generated clip for motion coherence and prompt adherence"]
|
||||
|
||||
def _get_api_key(self) -> str | None:
|
||||
return os.environ.get("FAL_KEY") or os.environ.get("FAL_AI_API_KEY")
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if self._get_api_key():
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
variant = inputs.get("model_variant", "hailuo-02/pro")
|
||||
if "pro" in variant:
|
||||
return 0.15
|
||||
if "fast" in variant:
|
||||
return 0.08
|
||||
return 0.10 # standard
|
||||
|
||||
def estimate_runtime(self, inputs: dict[str, Any]) -> float:
|
||||
variant = inputs.get("model_variant", "hailuo-02/pro")
|
||||
if "fast" in variant:
|
||||
return 30.0
|
||||
return 60.0
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
api_key = self._get_api_key()
|
||||
if not api_key:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="FAL_KEY not set. " + self.install_instructions,
|
||||
)
|
||||
|
||||
import requests
|
||||
|
||||
start = time.time()
|
||||
operation = inputs.get("operation", "text_to_video")
|
||||
variant = inputs.get("model_variant", "hailuo-02/pro")
|
||||
|
||||
# Build fal.ai model path
|
||||
if operation == "text_to_video":
|
||||
model_path = f"minimax/{variant}/text-to-video"
|
||||
if variant == "video-01":
|
||||
model_path = "minimax/video-01"
|
||||
else:
|
||||
model_path = f"minimax/{variant}/image-to-video"
|
||||
if variant == "video-01":
|
||||
model_path = "minimax/video-01/image-to-video"
|
||||
|
||||
payload: dict[str, Any] = {"prompt": inputs["prompt"]}
|
||||
if operation == "image_to_video" and inputs.get("image_url"):
|
||||
payload["image_url"] = inputs["image_url"]
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
f"https://fal.run/fal-ai/{model_path}",
|
||||
headers={
|
||||
"Authorization": f"Key {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
json=payload,
|
||||
timeout=300,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
video_url = data["video"]["url"]
|
||||
video_response = requests.get(video_url, timeout=120)
|
||||
video_response.raise_for_status()
|
||||
|
||||
output_path = Path(inputs.get("output_path", "minimax_output.mp4"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(video_response.content)
|
||||
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"MiniMax video generation failed: {e}")
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "minimax",
|
||||
"model": f"fal-ai/{model_path}",
|
||||
"prompt": inputs["prompt"],
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
cost_usd=self.estimate_cost(inputs),
|
||||
duration_seconds=round(time.time() - start, 2),
|
||||
model=f"fal-ai/{model_path}",
|
||||
)
|
||||
@@ -0,0 +1,213 @@
|
||||
"""Stock video acquisition from Pexels API (free)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class PexelsVideo(BaseTool):
|
||||
name = "pexels_video"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.SOURCE
|
||||
capability = "video_generation"
|
||||
provider = "pexels"
|
||||
stability = ToolStability.BETA
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
runtime = ToolRuntime.API
|
||||
|
||||
dependencies = []
|
||||
install_instructions = (
|
||||
"Set PEXELS_API_KEY to your Pexels API key.\n"
|
||||
" Get one free at https://www.pexels.com/api/"
|
||||
)
|
||||
agent_skills = []
|
||||
|
||||
capabilities = ["search_video", "download_video", "stock_video"]
|
||||
supports = {
|
||||
"orientation_filter": True,
|
||||
"size_filter": True,
|
||||
"free_commercial_use": True,
|
||||
}
|
||||
best_for = [
|
||||
"real-world B-roll footage (cities, nature, people, offices)",
|
||||
"establishing shots and transitions",
|
||||
"free stock video — no cost, no attribution required",
|
||||
]
|
||||
not_good_for = [
|
||||
"custom/specific scenes",
|
||||
"animated or stylized content",
|
||||
"offline use",
|
||||
]
|
||||
fallback_tools = ["pixabay_video"]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["query"],
|
||||
"properties": {
|
||||
"query": {"type": "string", "description": "Search term"},
|
||||
"orientation": {
|
||||
"type": "string",
|
||||
"enum": ["landscape", "portrait", "square"],
|
||||
},
|
||||
"size": {
|
||||
"type": "string",
|
||||
"enum": ["large", "medium", "small"],
|
||||
"description": "large=4K, medium=Full HD, small=HD",
|
||||
},
|
||||
"min_duration": {
|
||||
"type": "integer",
|
||||
"description": "Minimum duration in seconds",
|
||||
},
|
||||
"max_duration": {
|
||||
"type": "integer",
|
||||
"description": "Maximum duration in seconds",
|
||||
},
|
||||
"per_page": {"type": "integer", "default": 5, "minimum": 1, "maximum": 80},
|
||||
"page": {"type": "integer", "default": 1},
|
||||
"preferred_quality": {
|
||||
"type": "string",
|
||||
"enum": ["hd", "sd"],
|
||||
"default": "hd",
|
||||
},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=1, ram_mb=256, vram_mb=0, disk_mb=200, network_required=True
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=2, retryable_errors=["rate_limit", "timeout"])
|
||||
idempotency_key_fields = ["query", "orientation", "size", "page"]
|
||||
side_effects = ["writes video file to output_path", "calls Pexels API"]
|
||||
user_visible_verification = ["Watch downloaded clip to verify it matches the intended scene"]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if os.environ.get("PEXELS_API_KEY"):
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
return 0.0
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
api_key = os.environ.get("PEXELS_API_KEY")
|
||||
if not api_key:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="PEXELS_API_KEY not set. " + self.install_instructions,
|
||||
)
|
||||
|
||||
import requests
|
||||
|
||||
start = time.time()
|
||||
query = inputs["query"]
|
||||
|
||||
params: dict[str, Any] = {
|
||||
"query": query,
|
||||
"per_page": inputs.get("per_page", 5),
|
||||
"page": inputs.get("page", 1),
|
||||
}
|
||||
if inputs.get("orientation"):
|
||||
params["orientation"] = inputs["orientation"]
|
||||
if inputs.get("size"):
|
||||
params["size"] = inputs["size"]
|
||||
|
||||
try:
|
||||
search_response = requests.get(
|
||||
"https://api.pexels.com/videos/search",
|
||||
headers={"Authorization": api_key},
|
||||
params=params,
|
||||
timeout=30,
|
||||
)
|
||||
search_response.raise_for_status()
|
||||
data = search_response.json()
|
||||
|
||||
videos = data.get("videos", [])
|
||||
|
||||
# Filter by duration if specified
|
||||
min_dur = inputs.get("min_duration")
|
||||
max_dur = inputs.get("max_duration")
|
||||
if min_dur or max_dur:
|
||||
filtered = []
|
||||
for v in videos:
|
||||
dur = v.get("duration", 0)
|
||||
if min_dur and dur < min_dur:
|
||||
continue
|
||||
if max_dur and dur > max_dur:
|
||||
continue
|
||||
filtered.append(v)
|
||||
videos = filtered
|
||||
|
||||
if not videos:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"No videos found for query: {query}",
|
||||
data={"total_results": data.get("total_results", 0)},
|
||||
)
|
||||
|
||||
video = videos[0]
|
||||
preferred_quality = inputs.get("preferred_quality", "hd")
|
||||
|
||||
# Pick the best matching video file
|
||||
video_files = video.get("video_files", [])
|
||||
selected_file = None
|
||||
for vf in sorted(video_files, key=lambda x: x.get("width", 0), reverse=True):
|
||||
if vf.get("quality") == preferred_quality:
|
||||
selected_file = vf
|
||||
break
|
||||
if not selected_file and video_files:
|
||||
selected_file = video_files[0]
|
||||
|
||||
if not selected_file:
|
||||
return ToolResult(success=False, error="No downloadable video file found.")
|
||||
|
||||
video_url = selected_file["link"]
|
||||
video_response = requests.get(video_url, timeout=120)
|
||||
video_response.raise_for_status()
|
||||
|
||||
output_path = Path(inputs.get("output_path", f"pexels_video_{video['id']}.mp4"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(video_response.content)
|
||||
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Pexels video search failed: {e}")
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "pexels",
|
||||
"video_id": video["id"],
|
||||
"user": video.get("user", {}).get("name", "Unknown"),
|
||||
"duration_seconds": video.get("duration"),
|
||||
"width": selected_file.get("width"),
|
||||
"height": selected_file.get("height"),
|
||||
"fps": selected_file.get("fps"),
|
||||
"quality": selected_file.get("quality"),
|
||||
"query": query,
|
||||
"output": str(output_path),
|
||||
"total_results": data.get("total_results", 0),
|
||||
"results_returned": len(videos),
|
||||
"license": "Pexels License (free, no attribution required)",
|
||||
"pexels_url": video.get("url", ""),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
cost_usd=0.0,
|
||||
duration_seconds=round(time.time() - start, 2),
|
||||
)
|
||||
@@ -0,0 +1,221 @@
|
||||
"""Stock video acquisition from Pixabay API (free)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class PixabayVideo(BaseTool):
|
||||
name = "pixabay_video"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.SOURCE
|
||||
capability = "video_generation"
|
||||
provider = "pixabay"
|
||||
stability = ToolStability.BETA
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
runtime = ToolRuntime.API
|
||||
|
||||
dependencies = []
|
||||
install_instructions = (
|
||||
"Set PIXABAY_API_KEY to your Pixabay API key.\n"
|
||||
" Get one free at https://pixabay.com/api/docs/"
|
||||
)
|
||||
agent_skills = []
|
||||
|
||||
capabilities = ["search_video", "download_video", "stock_video"]
|
||||
supports = {
|
||||
"video_type_filter": True,
|
||||
"category_filter": True,
|
||||
"editors_choice": True,
|
||||
"free_commercial_use": True,
|
||||
}
|
||||
best_for = [
|
||||
"large royalty-free video library",
|
||||
"category-based filtering",
|
||||
"free stock video — no cost, no attribution required",
|
||||
]
|
||||
not_good_for = [
|
||||
"4K footage (max 1080p on standard API)",
|
||||
"custom scenes",
|
||||
"offline use",
|
||||
]
|
||||
fallback_tools = ["pexels_video"]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["query"],
|
||||
"properties": {
|
||||
"query": {"type": "string", "description": "Search term (max 100 chars)"},
|
||||
"video_type": {
|
||||
"type": "string",
|
||||
"enum": ["all", "film", "animation"],
|
||||
"default": "all",
|
||||
},
|
||||
"category": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"backgrounds", "fashion", "nature", "science", "education",
|
||||
"feelings", "health", "people", "religion", "places",
|
||||
"animals", "industry", "computer", "food", "sports",
|
||||
"transportation", "travel", "buildings", "business", "music",
|
||||
],
|
||||
},
|
||||
"min_duration": {
|
||||
"type": "integer",
|
||||
"description": "Minimum duration in seconds",
|
||||
},
|
||||
"max_duration": {
|
||||
"type": "integer",
|
||||
"description": "Maximum duration in seconds",
|
||||
},
|
||||
"editors_choice": {"type": "boolean", "default": False},
|
||||
"safesearch": {"type": "boolean", "default": True},
|
||||
"per_page": {"type": "integer", "default": 5, "minimum": 3, "maximum": 200},
|
||||
"page": {"type": "integer", "default": 1},
|
||||
"preferred_quality": {
|
||||
"type": "string",
|
||||
"enum": ["large", "medium", "small", "tiny"],
|
||||
"default": "large",
|
||||
},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=1, ram_mb=256, vram_mb=0, disk_mb=200, network_required=True
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=2, retryable_errors=["rate_limit", "timeout"])
|
||||
idempotency_key_fields = ["query", "video_type", "category", "page"]
|
||||
side_effects = ["writes video file to output_path", "calls Pixabay API"]
|
||||
user_visible_verification = ["Watch downloaded clip to verify it matches the intended scene"]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if os.environ.get("PIXABAY_API_KEY"):
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
return 0.0
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
api_key = os.environ.get("PIXABAY_API_KEY")
|
||||
if not api_key:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="PIXABAY_API_KEY not set. " + self.install_instructions,
|
||||
)
|
||||
|
||||
import requests
|
||||
|
||||
start = time.time()
|
||||
query = inputs["query"]
|
||||
|
||||
params: dict[str, Any] = {
|
||||
"key": api_key,
|
||||
"q": query,
|
||||
"per_page": inputs.get("per_page", 5),
|
||||
"page": inputs.get("page", 1),
|
||||
"safesearch": str(inputs.get("safesearch", True)).lower(),
|
||||
}
|
||||
if inputs.get("video_type") and inputs["video_type"] != "all":
|
||||
params["video_type"] = inputs["video_type"]
|
||||
if inputs.get("category"):
|
||||
params["category"] = inputs["category"]
|
||||
if inputs.get("editors_choice"):
|
||||
params["editors_choice"] = "true"
|
||||
|
||||
try:
|
||||
search_response = requests.get(
|
||||
"https://pixabay.com/api/videos/",
|
||||
params=params,
|
||||
timeout=30,
|
||||
)
|
||||
search_response.raise_for_status()
|
||||
data = search_response.json()
|
||||
|
||||
hits = data.get("hits", [])
|
||||
|
||||
# Filter by duration if specified
|
||||
min_dur = inputs.get("min_duration")
|
||||
max_dur = inputs.get("max_duration")
|
||||
if min_dur or max_dur:
|
||||
filtered = []
|
||||
for h in hits:
|
||||
dur = h.get("duration", 0)
|
||||
if min_dur and dur < min_dur:
|
||||
continue
|
||||
if max_dur and dur > max_dur:
|
||||
continue
|
||||
filtered.append(h)
|
||||
hits = filtered
|
||||
|
||||
if not hits:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"No videos found for query: {query}",
|
||||
data={"total_results": data.get("total", 0)},
|
||||
)
|
||||
|
||||
hit = hits[0]
|
||||
preferred = inputs.get("preferred_quality", "large")
|
||||
video_info = hit.get("videos", {}).get(preferred)
|
||||
if not video_info:
|
||||
# Fallback to best available
|
||||
for quality in ["large", "medium", "small", "tiny"]:
|
||||
video_info = hit.get("videos", {}).get(quality)
|
||||
if video_info:
|
||||
break
|
||||
|
||||
if not video_info:
|
||||
return ToolResult(success=False, error="No downloadable video file found.")
|
||||
|
||||
# Download immediately — Pixabay URLs expire
|
||||
video_url = video_info["url"]
|
||||
video_response = requests.get(video_url, timeout=120)
|
||||
video_response.raise_for_status()
|
||||
|
||||
output_path = Path(inputs.get("output_path", f"pixabay_video_{hit['id']}.mp4"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(video_response.content)
|
||||
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Pixabay video search failed: {e}")
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "pixabay",
|
||||
"video_id": hit["id"],
|
||||
"user": hit.get("user", "Unknown"),
|
||||
"tags": hit.get("tags", ""),
|
||||
"duration_seconds": hit.get("duration"),
|
||||
"width": video_info.get("width"),
|
||||
"height": video_info.get("height"),
|
||||
"query": query,
|
||||
"output": str(output_path),
|
||||
"total_results": data.get("total", 0),
|
||||
"results_returned": len(hits),
|
||||
"license": "Pixabay Content License (free, no attribution required)",
|
||||
"page_url": hit.get("pageURL", ""),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
cost_usd=0.0,
|
||||
duration_seconds=round(time.time() - start, 2),
|
||||
)
|
||||
@@ -0,0 +1,204 @@
|
||||
"""Runway Gen-4 video generation via Runway API.
|
||||
|
||||
Highest Elo-rated video generation model — professional quality and control.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class RunwayVideo(BaseTool):
|
||||
name = "runway_video"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "video_generation"
|
||||
provider = "runway"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.STOCHASTIC
|
||||
runtime = ToolRuntime.API
|
||||
|
||||
dependencies = []
|
||||
install_instructions = (
|
||||
"Set RUNWAY_API_KEY to your Runway API key.\n"
|
||||
" Get one at https://app.runwayml.com/settings/api-keys"
|
||||
)
|
||||
agent_skills = ["ai-video-gen"]
|
||||
|
||||
capabilities = ["text_to_video", "image_to_video"]
|
||||
supports = {
|
||||
"text_to_video": True,
|
||||
"image_to_video": True,
|
||||
"professional_control": True,
|
||||
}
|
||||
best_for = [
|
||||
"highest overall video quality (#1 Elo rating)",
|
||||
"professional video production",
|
||||
"precise control over generation",
|
||||
]
|
||||
not_good_for = ["budget projects", "offline generation", "very long clips"]
|
||||
fallback_tools = ["kling_video", "veo_video", "minimax_video", "wan_video"]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {"type": "string"},
|
||||
"operation": {
|
||||
"type": "string",
|
||||
"enum": ["text_to_video", "image_to_video"],
|
||||
"default": "text_to_video",
|
||||
},
|
||||
"model": {
|
||||
"type": "string",
|
||||
"enum": ["gen4_turbo", "gen4"],
|
||||
"default": "gen4_turbo",
|
||||
},
|
||||
"duration": {
|
||||
"type": "integer",
|
||||
"enum": [5, 10],
|
||||
"default": 5,
|
||||
"description": "Duration in seconds",
|
||||
},
|
||||
"ratio": {
|
||||
"type": "string",
|
||||
"enum": ["16:9", "9:16", "1:1"],
|
||||
"default": "16:9",
|
||||
},
|
||||
"image_url": {"type": "string", "description": "Reference image URL for image_to_video"},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=1, ram_mb=512, vram_mb=0, disk_mb=500, network_required=True
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=2, retryable_errors=["rate_limit", "timeout"])
|
||||
idempotency_key_fields = ["prompt", "model", "operation", "duration"]
|
||||
side_effects = ["writes video file to output_path", "calls Runway API"]
|
||||
user_visible_verification = ["Watch generated clip for visual quality and motion coherence"]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if os.environ.get("RUNWAY_API_KEY"):
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
duration = inputs.get("duration", 5)
|
||||
# Runway charges per second of generated video
|
||||
return 0.05 * duration # ~$0.25 for 5s, ~$0.50 for 10s
|
||||
|
||||
def estimate_runtime(self, inputs: dict[str, Any]) -> float:
|
||||
model = inputs.get("model", "gen4_turbo")
|
||||
if "turbo" in model:
|
||||
return 30.0
|
||||
return 60.0
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
api_key = os.environ.get("RUNWAY_API_KEY")
|
||||
if not api_key:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="RUNWAY_API_KEY not set. " + self.install_instructions,
|
||||
)
|
||||
|
||||
import requests
|
||||
|
||||
start = time.time()
|
||||
model = inputs.get("model", "gen4_turbo")
|
||||
operation = inputs.get("operation", "text_to_video")
|
||||
|
||||
# Runway API v1 — submit task
|
||||
task_payload: dict[str, Any] = {
|
||||
"model": model,
|
||||
"promptText": inputs["prompt"],
|
||||
"duration": inputs.get("duration", 5),
|
||||
"ratio": inputs.get("ratio", "16:9"),
|
||||
}
|
||||
if operation == "image_to_video" and inputs.get("image_url"):
|
||||
task_payload["promptImage"] = inputs["image_url"]
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
"X-Runway-Version": "2024-11-06",
|
||||
}
|
||||
|
||||
try:
|
||||
# Submit generation task
|
||||
submit_response = requests.post(
|
||||
"https://api.dev.runwayml.com/v1/image_to_video" if operation == "image_to_video"
|
||||
else "https://api.dev.runwayml.com/v1/text_to_video",
|
||||
headers=headers,
|
||||
json=task_payload,
|
||||
timeout=30,
|
||||
)
|
||||
submit_response.raise_for_status()
|
||||
task_id = submit_response.json()["id"]
|
||||
|
||||
# Poll for completion
|
||||
video_url = None
|
||||
for _ in range(60): # max 5 minutes
|
||||
time.sleep(5)
|
||||
poll_response = requests.get(
|
||||
f"https://api.dev.runwayml.com/v1/tasks/{task_id}",
|
||||
headers=headers,
|
||||
timeout=15,
|
||||
)
|
||||
poll_response.raise_for_status()
|
||||
task_data = poll_response.json()
|
||||
|
||||
if task_data["status"] == "SUCCEEDED":
|
||||
video_url = task_data["output"][0]
|
||||
break
|
||||
if task_data["status"] == "FAILED":
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"Runway generation failed: {task_data.get('failure', 'unknown error')}",
|
||||
)
|
||||
|
||||
if not video_url:
|
||||
return ToolResult(success=False, error="Runway generation timed out.")
|
||||
|
||||
# Download video
|
||||
video_response = requests.get(video_url, timeout=120)
|
||||
video_response.raise_for_status()
|
||||
|
||||
output_path = Path(inputs.get("output_path", "runway_output.mp4"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(video_response.content)
|
||||
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Runway video generation failed: {e}")
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "runway",
|
||||
"model": model,
|
||||
"prompt": inputs["prompt"],
|
||||
"output": str(output_path),
|
||||
"task_id": task_id,
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
cost_usd=self.estimate_cost(inputs),
|
||||
duration_seconds=round(time.time() - start, 2),
|
||||
model=model,
|
||||
)
|
||||
@@ -0,0 +1,198 @@
|
||||
"""Google Veo 3 video generation via fal.ai API.
|
||||
|
||||
State-of-the-art video generation with native audio/dialogue synthesis.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class VeoVideo(BaseTool):
|
||||
name = "veo_video"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "video_generation"
|
||||
provider = "veo"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.STOCHASTIC
|
||||
runtime = ToolRuntime.API
|
||||
|
||||
dependencies = []
|
||||
install_instructions = (
|
||||
"Set FAL_KEY to your fal.ai API key.\n"
|
||||
" Get one at https://fal.ai/dashboard/keys"
|
||||
)
|
||||
agent_skills = ["ai-video-gen"]
|
||||
|
||||
capabilities = ["text_to_video", "image_to_video"]
|
||||
supports = {
|
||||
"text_to_video": True,
|
||||
"image_to_video": True,
|
||||
"native_audio": True,
|
||||
"dialogue_generation": True,
|
||||
"ambient_sound": True,
|
||||
}
|
||||
best_for = [
|
||||
"videos with synchronized dialogue and audio",
|
||||
"cutting-edge quality from Google DeepMind",
|
||||
"ambient sound and music generation built in",
|
||||
]
|
||||
not_good_for = ["budget projects", "offline generation", "quick iteration"]
|
||||
fallback_tools = ["kling_video", "minimax_video", "wan_video"]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {"type": "string"},
|
||||
"operation": {
|
||||
"type": "string",
|
||||
"enum": ["text_to_video", "image_to_video"],
|
||||
"default": "text_to_video",
|
||||
},
|
||||
"model_variant": {
|
||||
"type": "string",
|
||||
"enum": ["veo3", "veo3/fast", "veo3.1", "veo3.1/fast"],
|
||||
"default": "veo3",
|
||||
},
|
||||
"duration": {
|
||||
"type": "string",
|
||||
"enum": ["5", "8"],
|
||||
"default": "8",
|
||||
"description": "Duration in seconds",
|
||||
},
|
||||
"aspect_ratio": {
|
||||
"type": "string",
|
||||
"enum": ["16:9", "9:16"],
|
||||
"default": "16:9",
|
||||
},
|
||||
"generate_audio": {
|
||||
"type": "boolean",
|
||||
"default": True,
|
||||
"description": "Whether to generate synchronized audio",
|
||||
},
|
||||
"image_url": {"type": "string", "description": "Reference image URL for image_to_video"},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=1, ram_mb=512, vram_mb=0, disk_mb=500, network_required=True
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=2, retryable_errors=["rate_limit", "timeout"])
|
||||
idempotency_key_fields = ["prompt", "model_variant", "operation", "duration"]
|
||||
side_effects = ["writes video file to output_path", "calls fal.ai API"]
|
||||
user_visible_verification = [
|
||||
"Watch generated clip for visual quality and motion",
|
||||
"Listen for audio synchronization and quality",
|
||||
]
|
||||
|
||||
def _get_api_key(self) -> str | None:
|
||||
return os.environ.get("FAL_KEY") or os.environ.get("FAL_AI_API_KEY")
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if self._get_api_key():
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
variant = inputs.get("model_variant", "veo3")
|
||||
duration = int(inputs.get("duration", "8"))
|
||||
if "fast" in variant:
|
||||
per_second = 0.12
|
||||
else:
|
||||
per_second = 0.30
|
||||
return per_second * duration
|
||||
|
||||
def estimate_runtime(self, inputs: dict[str, Any]) -> float:
|
||||
variant = inputs.get("model_variant", "veo3")
|
||||
if "fast" in variant:
|
||||
return 45.0
|
||||
return 120.0
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
api_key = self._get_api_key()
|
||||
if not api_key:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="FAL_KEY not set. " + self.install_instructions,
|
||||
)
|
||||
|
||||
import requests
|
||||
|
||||
start = time.time()
|
||||
operation = inputs.get("operation", "text_to_video")
|
||||
variant = inputs.get("model_variant", "veo3")
|
||||
|
||||
# Build fal.ai model path
|
||||
if operation == "image_to_video":
|
||||
model_path = f"{variant}/image-to-video"
|
||||
else:
|
||||
model_path = variant # text-to-video is the default endpoint
|
||||
|
||||
payload: dict[str, Any] = {"prompt": inputs["prompt"]}
|
||||
if inputs.get("duration"):
|
||||
payload["duration"] = inputs["duration"]
|
||||
if inputs.get("aspect_ratio"):
|
||||
payload["aspect_ratio"] = inputs["aspect_ratio"]
|
||||
if inputs.get("generate_audio") is not None:
|
||||
payload["generate_audio"] = inputs["generate_audio"]
|
||||
if operation == "image_to_video" and inputs.get("image_url"):
|
||||
payload["image_url"] = inputs["image_url"]
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
f"https://fal.run/fal-ai/{model_path}",
|
||||
headers={
|
||||
"Authorization": f"Key {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
json=payload,
|
||||
timeout=300,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
video_url = data["video"]["url"]
|
||||
video_response = requests.get(video_url, timeout=120)
|
||||
video_response.raise_for_status()
|
||||
|
||||
output_path = Path(inputs.get("output_path", "veo_output.mp4"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(video_response.content)
|
||||
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Veo video generation failed: {e}")
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"provider": "veo",
|
||||
"model": f"fal-ai/{model_path}",
|
||||
"prompt": inputs["prompt"],
|
||||
"output": str(output_path),
|
||||
"has_audio": inputs.get("generate_audio", True),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
cost_usd=self.estimate_cost(inputs),
|
||||
duration_seconds=round(time.time() - start, 2),
|
||||
model=f"fal-ai/{model_path}",
|
||||
)
|
||||
@@ -0,0 +1,787 @@
|
||||
"""Video composition tool — FFmpeg + Remotion.
|
||||
|
||||
Final composition path that takes edit decisions, assets, and audio
|
||||
and renders the complete output video. Supports subtitle burn-in,
|
||||
overlay compositing, and platform-specific encoding profiles.
|
||||
|
||||
For compositions with still images, animated scenes, or component types
|
||||
(text cards, stat cards, etc.), the render operation auto-routes to
|
||||
Remotion for frame-accurate spring animations and React-based rendering.
|
||||
For pure video cuts (talking-head, etc.), FFmpeg handles trimming and concat.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ResumeSupport,
|
||||
ToolResult,
|
||||
ToolStability,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class VideoCompose(BaseTool):
|
||||
name = "video_compose"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.CORE
|
||||
capability = "video_post"
|
||||
provider = "ffmpeg"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
|
||||
dependencies = ["cmd:ffmpeg"]
|
||||
install_instructions = "Install FFmpeg: https://ffmpeg.org/download.html"
|
||||
agent_skills = ["remotion-best-practices", "remotion", "ffmpeg"]
|
||||
|
||||
capabilities = [
|
||||
"compose_cuts",
|
||||
"burn_subtitles",
|
||||
"overlay_assets",
|
||||
"encode_profile",
|
||||
"remotion_render",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["operation"],
|
||||
"properties": {
|
||||
"operation": {
|
||||
"type": "string",
|
||||
"enum": ["compose", "render", "remotion_render", "burn_subtitles", "overlay", "encode"],
|
||||
"description": (
|
||||
"compose: low-level concat cuts + audio + subtitles. "
|
||||
"render: high-level — resolves asset IDs, auto-routes to Remotion "
|
||||
"for images/animations or FFmpeg for video-only. Preferred for compose-director. "
|
||||
"remotion_render: render via Remotion (Node.js). "
|
||||
"burn_subtitles: burn subtitle file into existing video. "
|
||||
"overlay: composite overlays onto base video. "
|
||||
"encode: re-encode to a target profile/codec."
|
||||
),
|
||||
},
|
||||
"input_path": {"type": "string"},
|
||||
"output_path": {"type": "string"},
|
||||
"edit_decisions": {
|
||||
"type": "object",
|
||||
"description": "Full edit_decisions artifact (required for compose/render)",
|
||||
},
|
||||
"asset_manifest": {
|
||||
"type": "object",
|
||||
"description": (
|
||||
"Full asset_manifest artifact (required for render). "
|
||||
"Used to resolve asset IDs in cuts[].source to file paths."
|
||||
),
|
||||
},
|
||||
"subtitle_path": {"type": "string"},
|
||||
"subtitle_style": {
|
||||
"type": "object",
|
||||
"description": "ASS subtitle styling. Also extracted from edit_decisions.subtitles if not provided.",
|
||||
"properties": {
|
||||
"font": {"type": "string", "default": "Arial"},
|
||||
"font_size": {"type": "integer", "default": 24},
|
||||
"primary_color": {"type": "string", "default": "&HFFFFFF"},
|
||||
"outline_color": {"type": "string", "default": "&H000000"},
|
||||
"outline_width": {"type": "number", "default": 2},
|
||||
"margin_v": {"type": "integer", "default": 40},
|
||||
"alignment": {"type": "integer", "default": 2},
|
||||
},
|
||||
},
|
||||
"overlays": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"asset_path": {"type": "string"},
|
||||
"x": {"type": "number"},
|
||||
"y": {"type": "number"},
|
||||
"width": {"type": "number"},
|
||||
"height": {"type": "number"},
|
||||
"start_seconds": {"type": "number"},
|
||||
"end_seconds": {"type": "number"},
|
||||
"opacity": {"type": "number", "minimum": 0, "maximum": 1},
|
||||
},
|
||||
},
|
||||
},
|
||||
"audio_path": {"type": "string", "description": "Mixed audio to mux into output"},
|
||||
"profile": {
|
||||
"type": "string",
|
||||
"description": (
|
||||
"Media profile name from media_profiles.py "
|
||||
"(e.g. youtube_landscape, tiktok, instagram_reels). "
|
||||
"Applied in render and encode operations."
|
||||
),
|
||||
},
|
||||
"options": {
|
||||
"type": "object",
|
||||
"description": "Render options (used by the render operation)",
|
||||
"properties": {
|
||||
"subtitle_burn": {"type": "boolean", "default": True},
|
||||
"two_pass_encode": {"type": "boolean", "default": False},
|
||||
},
|
||||
},
|
||||
"codec": {"type": "string", "default": "libx264"},
|
||||
"crf": {"type": "integer", "default": 23},
|
||||
"preset": {"type": "string", "default": "medium"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=4, ram_mb=2048, vram_mb=0, disk_mb=5000, network_required=False
|
||||
)
|
||||
|
||||
# Remotion scene types that trigger React-based rendering
|
||||
_REMOTION_COMPONENTS = [
|
||||
"text_card", "stat_card", "callout", "comparison",
|
||||
"progress", "chart", "bar_chart", "line_chart", "pie_chart", "kpi_grid",
|
||||
]
|
||||
|
||||
best_for = [
|
||||
"Final render for explainer and animation pipelines",
|
||||
"Image-to-video with spring animations (Remotion)",
|
||||
"Animated text cards, stat cards, charts (Remotion)",
|
||||
"Complex transitions between scenes (Remotion)",
|
||||
"Pure video concat and trim (FFmpeg)",
|
||||
]
|
||||
retry_policy = RetryPolicy(max_retries=1, retryable_errors=["Conversion failed"])
|
||||
resume_support = ResumeSupport.FROM_START
|
||||
idempotency_key_fields = ["operation", "input_path", "edit_decisions"]
|
||||
side_effects = ["writes video file to output_path"]
|
||||
user_visible_verification = [
|
||||
"Play the composed output and verify cuts, subtitles, and overlays",
|
||||
]
|
||||
|
||||
def _remotion_available(self) -> bool:
|
||||
"""Check if Remotion rendering is available (requires npx + composer project)."""
|
||||
import shutil as _shutil
|
||||
|
||||
if not _shutil.which("npx"):
|
||||
return False
|
||||
composer_dir = Path(__file__).resolve().parent.parent.parent / "remotion-composer"
|
||||
return composer_dir.exists() and (composer_dir / "package.json").exists()
|
||||
|
||||
def get_info(self) -> dict[str, Any]:
|
||||
"""Extend base get_info to surface Remotion sub-capability.
|
||||
|
||||
This lets preflight report 'video_compose: AVAILABLE (FFmpeg + Remotion)'
|
||||
so the agent knows Remotion is usable for motion graphics, animated text
|
||||
cards, stat cards, charts, and image-to-video rendering — rather than
|
||||
falling back to Ken Burns pan-and-zoom over still images.
|
||||
"""
|
||||
info = super().get_info()
|
||||
remotion_ok = self._remotion_available()
|
||||
info["render_engines"] = {
|
||||
"ffmpeg": True,
|
||||
"remotion": remotion_ok,
|
||||
}
|
||||
if remotion_ok:
|
||||
info["remotion_components"] = self._REMOTION_COMPONENTS
|
||||
info["remotion_note"] = (
|
||||
"Remotion is available for React-based rendering. Use it for "
|
||||
"image-to-video with spring animations, animated text/stat cards, "
|
||||
"charts, callouts, comparisons, and complex transitions. "
|
||||
"Prefer Remotion over Ken Burns pan-and-zoom for explainer "
|
||||
"and motion-graphics pipelines."
|
||||
)
|
||||
else:
|
||||
info["remotion_note"] = (
|
||||
"Remotion is NOT available (needs Node.js/npx + remotion-composer). "
|
||||
"Falling back to FFmpeg Ken Burns for image-based compositions."
|
||||
)
|
||||
return info
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
operation = inputs["operation"]
|
||||
start = time.time()
|
||||
|
||||
try:
|
||||
if operation == "compose":
|
||||
result = self._compose(inputs)
|
||||
elif operation == "render":
|
||||
result = self._render(inputs)
|
||||
elif operation == "remotion_render":
|
||||
result = self._remotion_render(inputs)
|
||||
elif operation == "burn_subtitles":
|
||||
result = self._burn_subtitles(inputs)
|
||||
elif operation == "overlay":
|
||||
result = self._overlay(inputs)
|
||||
elif operation == "encode":
|
||||
result = self._encode(inputs)
|
||||
else:
|
||||
return ToolResult(success=False, error=f"Unknown operation: {operation}")
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=str(e))
|
||||
|
||||
result.duration_seconds = round(time.time() - start, 2)
|
||||
return result
|
||||
|
||||
_IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".tif", ".webp"}
|
||||
|
||||
@staticmethod
|
||||
def _is_image(path: Path) -> bool:
|
||||
"""Check if a file is a still image (routes to Remotion, not FFmpeg)."""
|
||||
return path.suffix.lower() in VideoCompose._IMAGE_EXTENSIONS
|
||||
|
||||
def _compose(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
"""FFmpeg composition: concat video cuts, add audio, burn subtitles.
|
||||
|
||||
Handles video sources only. Still images and animated scene types
|
||||
are routed to Remotion via the render operation — call compose
|
||||
directly only for pure video pipelines (e.g. talking-head).
|
||||
"""
|
||||
edit_decisions = inputs.get("edit_decisions")
|
||||
if not edit_decisions:
|
||||
return ToolResult(success=False, error="edit_decisions required for compose")
|
||||
|
||||
output_path = Path(inputs.get("output_path", "composed_output.mp4"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
audio_path = inputs.get("audio_path")
|
||||
subtitle_path = inputs.get("subtitle_path")
|
||||
codec = inputs.get("codec", "libx264")
|
||||
crf = inputs.get("crf", 23)
|
||||
preset = inputs.get("preset", "medium")
|
||||
profile_name = inputs.get("profile")
|
||||
|
||||
# Resolve target resolution from profile or default
|
||||
resolution = "1920x1080"
|
||||
if profile_name:
|
||||
try:
|
||||
from lib.media_profiles import get_profile
|
||||
p = get_profile(profile_name)
|
||||
resolution = f"{p.width}x{p.height}"
|
||||
except (ImportError, ValueError):
|
||||
pass
|
||||
|
||||
cuts = edit_decisions.get("cuts", [])
|
||||
if not cuts:
|
||||
return ToolResult(success=False, error="No cuts in edit_decisions")
|
||||
|
||||
# Extract subtitle style from edit_decisions if not provided directly
|
||||
if not inputs.get("subtitle_style"):
|
||||
ed_subs = edit_decisions.get("subtitles", {})
|
||||
if ed_subs:
|
||||
inputs = dict(inputs)
|
||||
inputs["subtitle_style"] = {
|
||||
k: v for k, v in ed_subs.items()
|
||||
if k in ("font", "font_size", "color", "outline_color", "background")
|
||||
}
|
||||
if ed_subs.get("source") and not subtitle_path:
|
||||
subtitle_path = ed_subs["source"]
|
||||
|
||||
temp_dir = output_path.parent / ".compose_tmp"
|
||||
temp_dir.mkdir(parents=True, exist_ok=True)
|
||||
temp_segments: list[Path] = []
|
||||
|
||||
try:
|
||||
for i, cut in enumerate(cuts):
|
||||
source = Path(cut["source"])
|
||||
if not source.exists():
|
||||
return ToolResult(success=False, error=f"Cut source not found: {source}")
|
||||
|
||||
seg_path = temp_dir / f"seg_{i:04d}.mp4"
|
||||
in_s = cut["in_seconds"]
|
||||
out_s = cut["out_seconds"]
|
||||
duration = out_s - in_s
|
||||
speed = cut.get("speed", 1.0)
|
||||
|
||||
if self._is_image(source):
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=(
|
||||
f"Still image '{source.name}' in cuts. "
|
||||
"Use operation='render' (auto-routes to Remotion) "
|
||||
"or operation='remotion_render' for compositions "
|
||||
"with images, animations, or component scenes."
|
||||
),
|
||||
)
|
||||
else:
|
||||
# Video source: trim to segment
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", str(source),
|
||||
"-ss", str(in_s),
|
||||
"-to", str(out_s),
|
||||
]
|
||||
|
||||
if speed != 1.0:
|
||||
vf = f"setpts={1.0/speed}*PTS"
|
||||
af = self._build_atempo(speed)
|
||||
cmd.extend(["-filter:v", vf, "-filter:a", af])
|
||||
cmd.extend(["-c:v", codec, "-crf", str(crf), "-c:a", "aac"])
|
||||
else:
|
||||
cmd.extend(["-c", "copy"])
|
||||
|
||||
cmd.append(str(seg_path))
|
||||
self.run_command(cmd)
|
||||
|
||||
temp_segments.append(seg_path)
|
||||
|
||||
# Step 2: Concat segments
|
||||
concat_path = temp_dir / "concat_list.txt"
|
||||
with open(concat_path, "w", encoding="utf-8") as f:
|
||||
for seg in temp_segments:
|
||||
safe = str(seg.resolve()).replace("\\", "/")
|
||||
f.write(f"file '{safe}'\n")
|
||||
|
||||
concat_out = temp_dir / "concat.mp4"
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-f", "concat", "-safe", "0",
|
||||
"-i", str(concat_path),
|
||||
"-c", "copy",
|
||||
str(concat_out),
|
||||
]
|
||||
self.run_command(cmd)
|
||||
|
||||
# Step 3: Apply subtitles and/or replace audio
|
||||
final_input = concat_out
|
||||
vfilters = []
|
||||
|
||||
if subtitle_path and Path(subtitle_path).exists():
|
||||
style = inputs.get("subtitle_style", {})
|
||||
ass_style = self._build_subtitle_style(style)
|
||||
sub_escaped = str(Path(subtitle_path).resolve()).replace("\\", "/").replace(":", "\\:")
|
||||
vfilters.append(f"subtitles='{sub_escaped}':force_style='{ass_style}'")
|
||||
|
||||
cmd = ["ffmpeg", "-y", "-i", str(final_input)]
|
||||
|
||||
if audio_path and Path(audio_path).exists():
|
||||
cmd.extend(["-i", audio_path])
|
||||
|
||||
if vfilters:
|
||||
cmd.extend(["-vf", ",".join(vfilters)])
|
||||
cmd.extend(["-c:v", codec, "-crf", str(crf), "-preset", preset])
|
||||
else:
|
||||
cmd.extend(["-c:v", "copy"])
|
||||
|
||||
if audio_path and Path(audio_path).exists():
|
||||
cmd.extend(["-map", "0:v:0", "-map", "1:a:0", "-c:a", "aac", "-shortest"])
|
||||
else:
|
||||
cmd.extend(["-c:a", "copy"])
|
||||
|
||||
# Apply profile resolution/fps at final output
|
||||
if profile_name:
|
||||
try:
|
||||
from lib.media_profiles import get_profile
|
||||
p = get_profile(profile_name)
|
||||
cmd.extend(["-s", f"{p.width}x{p.height}", "-r", str(p.fps)])
|
||||
except (ImportError, ValueError):
|
||||
pass
|
||||
|
||||
cmd.append(str(output_path))
|
||||
self.run_command(cmd)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"operation": "compose",
|
||||
"cut_count": len(cuts),
|
||||
"has_subtitles": subtitle_path is not None,
|
||||
"has_mixed_audio": audio_path is not None,
|
||||
"profile": profile_name,
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
)
|
||||
finally:
|
||||
# Cleanup temp files
|
||||
for f in temp_segments:
|
||||
if f.exists():
|
||||
f.unlink()
|
||||
for f in [concat_path, concat_out]:
|
||||
if f.exists():
|
||||
f.unlink()
|
||||
if temp_dir.exists():
|
||||
try:
|
||||
temp_dir.rmdir()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
_REMOTION_SCENE_TYPES = {
|
||||
"text_card", "stat_card", "callout", "comparison", "progress", "chart",
|
||||
}
|
||||
|
||||
def _needs_remotion(self, cuts: list[dict]) -> bool:
|
||||
"""Determine if the composition requires Remotion.
|
||||
|
||||
Returns True when any cut contains still images, animated scene types,
|
||||
component types (text_card, stat_card, etc.), or transitions — all of
|
||||
which benefit from Remotion's React-based rendering over FFmpeg.
|
||||
"""
|
||||
for cut in cuts:
|
||||
source = cut.get("source", "")
|
||||
if source and Path(source).suffix.lower() in self._IMAGE_EXTENSIONS:
|
||||
return True
|
||||
if cut.get("type") in self._REMOTION_SCENE_TYPES:
|
||||
return True
|
||||
if cut.get("animation") or cut.get("transition_in") or cut.get("transition_out"):
|
||||
return True
|
||||
transform = cut.get("transform", {})
|
||||
if transform and transform.get("animation"):
|
||||
return True
|
||||
return False
|
||||
|
||||
def _render(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
"""High-level render: assemble edit decisions + asset manifest into final video.
|
||||
|
||||
This is the primary entry point for the compose-director skill.
|
||||
It resolves asset IDs, then auto-routes to Remotion (for images,
|
||||
animations, component scenes) or FFmpeg (for pure video cuts).
|
||||
|
||||
The agent should pass edit_decisions, asset_manifest, and optionally
|
||||
profile, subtitle_path, audio_path, and options.
|
||||
"""
|
||||
edit_decisions = inputs.get("edit_decisions")
|
||||
asset_manifest = inputs.get("asset_manifest")
|
||||
if not edit_decisions:
|
||||
return ToolResult(success=False, error="edit_decisions required for render")
|
||||
if not asset_manifest:
|
||||
return ToolResult(success=False, error="asset_manifest required for render")
|
||||
|
||||
output_path = Path(inputs.get("output_path", "renders/output.mp4"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Build asset lookup: id -> asset info
|
||||
asset_lookup = {a["id"]: a for a in asset_manifest.get("assets", [])}
|
||||
|
||||
cuts = edit_decisions.get("cuts", [])
|
||||
if not cuts:
|
||||
return ToolResult(success=False, error="No cuts in edit_decisions")
|
||||
|
||||
# Resolve asset IDs in cuts to file paths
|
||||
resolved_cuts = []
|
||||
for cut in cuts:
|
||||
source_id = cut.get("source", "")
|
||||
resolved_cut = dict(cut)
|
||||
if source_id in asset_lookup:
|
||||
resolved_cut["source"] = asset_lookup[source_id]["path"]
|
||||
resolved_cuts.append(resolved_cut)
|
||||
|
||||
# Also accept profile as "output_profile" (skill convention) or "profile"
|
||||
profile = inputs.get("profile") or inputs.get("output_profile")
|
||||
|
||||
# --- Route: Remotion for rich content, FFmpeg for pure video ---
|
||||
if self._needs_remotion(resolved_cuts):
|
||||
remotion_inputs: dict[str, Any] = {
|
||||
"edit_decisions": dict(edit_decisions, cuts=resolved_cuts),
|
||||
"output_path": str(output_path),
|
||||
}
|
||||
if profile:
|
||||
remotion_inputs["profile"] = profile
|
||||
return self._remotion_render(remotion_inputs)
|
||||
|
||||
# --- FFmpeg path: pure video cuts (talking-head, etc.) ---
|
||||
# Handle options
|
||||
options = inputs.get("options", {})
|
||||
subtitle_burn = options.get("subtitle_burn", True)
|
||||
|
||||
# Resolve subtitle_path from edit_decisions if not provided
|
||||
subtitle_path = inputs.get("subtitle_path")
|
||||
if subtitle_burn and not subtitle_path:
|
||||
ed_subs = edit_decisions.get("subtitles", {})
|
||||
if ed_subs.get("enabled") and ed_subs.get("source"):
|
||||
subtitle_path = ed_subs["source"]
|
||||
|
||||
# Build compose inputs
|
||||
compose_inputs = dict(inputs)
|
||||
compose_inputs["edit_decisions"] = dict(edit_decisions, cuts=resolved_cuts)
|
||||
compose_inputs["output_path"] = str(output_path)
|
||||
if subtitle_path:
|
||||
compose_inputs["subtitle_path"] = subtitle_path
|
||||
if profile:
|
||||
compose_inputs["profile"] = profile
|
||||
|
||||
return self._compose(compose_inputs)
|
||||
|
||||
def _remotion_render(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
"""Render via Remotion (requires Node.js + npx).
|
||||
|
||||
Handles compositions with still images, animated scenes, component
|
||||
types, and transitions using React-based frame-accurate rendering.
|
||||
Accepts edit_decisions (with resolved file paths) or raw composition_data.
|
||||
"""
|
||||
import shutil
|
||||
|
||||
if not shutil.which("npx"):
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="npx not found. Install Node.js to use Remotion rendering.",
|
||||
)
|
||||
|
||||
composition_data = inputs.get("edit_decisions") or inputs.get("composition_data")
|
||||
if not composition_data:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="edit_decisions or composition_data required for remotion_render",
|
||||
)
|
||||
|
||||
output_path = Path(inputs.get("output_path", "renders/remotion_output.mp4"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Deep-copy props so we don't mutate the original
|
||||
props = json.loads(json.dumps(composition_data))
|
||||
|
||||
# Convert absolute file paths to file:// URIs for Remotion's
|
||||
# Img and OffthreadVideo components
|
||||
for cut in props.get("cuts", []):
|
||||
source = cut.get("source", "")
|
||||
if source and not source.startswith(("http://", "https://", "file://")):
|
||||
resolved = Path(source).resolve()
|
||||
if resolved.exists():
|
||||
posix = resolved.as_posix()
|
||||
cut["source"] = f"file:///{posix}" if not posix.startswith("/") else f"file://{posix}"
|
||||
|
||||
# Write props to temp file for Remotion CLI
|
||||
props_path = output_path.parent / ".remotion_props.json"
|
||||
with open(props_path, "w", encoding="utf-8") as f:
|
||||
json.dump(props, f)
|
||||
|
||||
# remotion-composer lives at project root
|
||||
composer_dir = Path(__file__).resolve().parent.parent.parent / "remotion-composer"
|
||||
if not composer_dir.exists():
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"Remotion composer project not found at {composer_dir}",
|
||||
)
|
||||
|
||||
cmd = [
|
||||
"npx", "remotion", "render",
|
||||
str(composer_dir / "src" / "index.tsx"),
|
||||
"Explainer",
|
||||
str(output_path),
|
||||
"--props", str(props_path),
|
||||
]
|
||||
|
||||
# Apply media profile dimensions
|
||||
profile_name = inputs.get("profile")
|
||||
if profile_name:
|
||||
try:
|
||||
from lib.media_profiles import get_profile
|
||||
p = get_profile(profile_name)
|
||||
cmd.extend(["--width", str(p.width), "--height", str(p.height)])
|
||||
except (ImportError, ValueError):
|
||||
pass
|
||||
|
||||
try:
|
||||
self.run_command(cmd, timeout=600)
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=f"Remotion render failed: {e}")
|
||||
finally:
|
||||
if props_path.exists():
|
||||
props_path.unlink()
|
||||
|
||||
if not output_path.exists():
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"Remotion render completed but output file missing: {output_path}",
|
||||
)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"operation": "remotion_render",
|
||||
"output": str(output_path),
|
||||
"profile": profile_name,
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
)
|
||||
|
||||
def _burn_subtitles(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
"""Burn subtitle file into video."""
|
||||
input_path = Path(inputs["input_path"])
|
||||
subtitle_path = Path(inputs["subtitle_path"])
|
||||
output_path = Path(inputs.get("output_path", str(input_path.with_stem(f"{input_path.stem}_subtitled"))))
|
||||
|
||||
if not input_path.exists():
|
||||
return ToolResult(success=False, error=f"Input not found: {input_path}")
|
||||
if not subtitle_path.exists():
|
||||
return ToolResult(success=False, error=f"Subtitle file not found: {subtitle_path}")
|
||||
|
||||
style = inputs.get("subtitle_style", {})
|
||||
ass_style = self._build_subtitle_style(style)
|
||||
sub_escaped = str(subtitle_path.resolve()).replace("\\", "/").replace(":", "\\:")
|
||||
codec = inputs.get("codec", "libx264")
|
||||
crf = inputs.get("crf", 23)
|
||||
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", str(input_path),
|
||||
"-vf", f"subtitles='{sub_escaped}':force_style='{ass_style}'",
|
||||
"-c:v", codec, "-crf", str(crf),
|
||||
"-c:a", "copy",
|
||||
str(output_path),
|
||||
]
|
||||
|
||||
self.run_command(cmd)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"operation": "burn_subtitles",
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
)
|
||||
|
||||
def _overlay(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
"""Composite overlay images/videos on top of base video."""
|
||||
input_path = Path(inputs["input_path"])
|
||||
overlays = inputs.get("overlays", [])
|
||||
output_path = Path(inputs.get("output_path", str(input_path.with_stem(f"{input_path.stem}_overlay"))))
|
||||
codec = inputs.get("codec", "libx264")
|
||||
crf = inputs.get("crf", 23)
|
||||
|
||||
if not input_path.exists():
|
||||
return ToolResult(success=False, error=f"Input not found: {input_path}")
|
||||
if not overlays:
|
||||
return ToolResult(success=False, error="No overlays provided")
|
||||
|
||||
# Build complex filter for each overlay
|
||||
input_args = ["-i", str(input_path)]
|
||||
filter_parts = []
|
||||
prev_label = "0:v"
|
||||
|
||||
for i, ov in enumerate(overlays):
|
||||
asset_path = Path(ov["asset_path"])
|
||||
if not asset_path.exists():
|
||||
return ToolResult(success=False, error=f"Overlay asset not found: {asset_path}")
|
||||
|
||||
input_args.extend(["-i", str(asset_path)])
|
||||
|
||||
x = int(ov.get("x", 0))
|
||||
y = int(ov.get("y", 0))
|
||||
start = ov.get("start_seconds", 0)
|
||||
end = ov.get("end_seconds")
|
||||
opacity = ov.get("opacity", 1.0)
|
||||
|
||||
overlay_input = f"{i + 1}:v"
|
||||
|
||||
# Scale overlay if dimensions specified
|
||||
if "width" in ov and "height" in ov:
|
||||
w = int(ov["width"])
|
||||
h = int(ov["height"])
|
||||
filter_parts.append(f"[{overlay_input}]scale={w}:{h}[ov_scaled_{i}]")
|
||||
overlay_input = f"ov_scaled_{i}"
|
||||
|
||||
# Build enable expression for timed overlays
|
||||
enable = f"between(t,{start},{end})" if end else f"gte(t,{start})"
|
||||
out_label = f"v{i}"
|
||||
|
||||
filter_parts.append(
|
||||
f"[{prev_label}][{overlay_input}]overlay={x}:{y}:enable='{enable}'[{out_label}]"
|
||||
)
|
||||
prev_label = out_label
|
||||
|
||||
filter_complex = ";".join(filter_parts)
|
||||
|
||||
cmd = ["ffmpeg", "-y"]
|
||||
cmd.extend(input_args)
|
||||
cmd.extend(["-filter_complex", filter_complex])
|
||||
cmd.extend(["-map", f"[{prev_label}]", "-map", "0:a?"])
|
||||
cmd.extend(["-c:v", codec, "-crf", str(crf), "-c:a", "copy"])
|
||||
cmd.append(str(output_path))
|
||||
|
||||
self.run_command(cmd)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"operation": "overlay",
|
||||
"overlay_count": len(overlays),
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
)
|
||||
|
||||
def _encode(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
"""Re-encode video with a specific profile/codec settings."""
|
||||
input_path = Path(inputs["input_path"])
|
||||
output_path = Path(inputs.get("output_path", str(input_path.with_stem(f"{input_path.stem}_encoded"))))
|
||||
codec = inputs.get("codec", "libx264")
|
||||
crf = inputs.get("crf", 23)
|
||||
preset = inputs.get("preset", "medium")
|
||||
profile_name = inputs.get("profile")
|
||||
|
||||
if not input_path.exists():
|
||||
return ToolResult(success=False, error=f"Input not found: {input_path}")
|
||||
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", str(input_path),
|
||||
"-c:v", codec, "-crf", str(crf), "-preset", preset,
|
||||
"-c:a", "aac", "-b:a", "192k",
|
||||
]
|
||||
|
||||
# Apply media profile if specified
|
||||
if profile_name:
|
||||
try:
|
||||
from lib.media_profiles import get_profile, ffmpeg_output_args
|
||||
profile = get_profile(profile_name)
|
||||
cmd.extend(["-s", f"{profile.width}x{profile.height}"])
|
||||
cmd.extend(["-r", str(profile.fps)])
|
||||
except (ImportError, ValueError):
|
||||
pass # proceed without profile
|
||||
|
||||
cmd.append(str(output_path))
|
||||
self.run_command(cmd)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"operation": "encode",
|
||||
"codec": codec,
|
||||
"crf": crf,
|
||||
"profile": profile_name,
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _build_subtitle_style(style: dict) -> str:
|
||||
"""Build ASS force_style string from style dict.
|
||||
|
||||
Produces modern social-media-style captions by default:
|
||||
bold, outlined, positioned in the lower portion of the frame.
|
||||
"""
|
||||
parts = []
|
||||
parts.append(f"FontName={style.get('font', 'Arial')}")
|
||||
parts.append(f"FontSize={style.get('font_size', 16)}")
|
||||
parts.append(f"Bold={1 if style.get('bold', True) else 0}")
|
||||
if style.get("primary_color"):
|
||||
parts.append(f"PrimaryColour={style['primary_color']}")
|
||||
if style.get("outline_color"):
|
||||
parts.append(f"OutlineColour={style['outline_color']}")
|
||||
if style.get("back_color"):
|
||||
parts.append(f"BackColour={style['back_color']}")
|
||||
# BorderStyle: 1=outline+shadow (default), 4=opaque box
|
||||
border_style = style.get("border_style", 1)
|
||||
parts.append(f"BorderStyle={border_style}")
|
||||
parts.append(f"Outline={style.get('outline_width', 3)}")
|
||||
parts.append(f"Shadow={style.get('shadow', 1)}")
|
||||
parts.append(f"MarginV={style.get('margin_v', 40)}")
|
||||
parts.append(f"Alignment={style.get('alignment', 2)}")
|
||||
return ",".join(parts)
|
||||
|
||||
@staticmethod
|
||||
def _build_atempo(factor: float) -> str:
|
||||
"""Build atempo filter chain for audio speed adjustment."""
|
||||
filters = []
|
||||
remaining = factor
|
||||
while remaining > 100.0:
|
||||
filters.append("atempo=100.0")
|
||||
remaining /= 100.0
|
||||
while remaining < 0.5:
|
||||
filters.append("atempo=0.5")
|
||||
remaining /= 0.5
|
||||
filters.append(f"atempo={remaining:.4f}")
|
||||
return ",".join(filters)
|
||||
@@ -0,0 +1,136 @@
|
||||
"""Capability-level video selector that routes between generation and stock providers.
|
||||
|
||||
Provider discovery is automatic — any BaseTool with capability="video_generation"
|
||||
is picked up from the registry. Adding a new video provider requires only creating
|
||||
the tool file in tools/video/; no changes to this selector are needed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
from tools.base_tool import BaseTool, ToolResult, ToolRuntime, ToolStability, ToolStatus, ToolTier
|
||||
|
||||
|
||||
class VideoSelector(BaseTool):
|
||||
name = "video_selector"
|
||||
version = "0.3.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "video_generation"
|
||||
provider = "selector"
|
||||
stability = ToolStability.BETA
|
||||
runtime = ToolRuntime.HYBRID
|
||||
agent_skills = ["ai-video-gen", "create-video", "ltx2"]
|
||||
|
||||
capabilities = [
|
||||
"text_to_video", "image_to_video", "stock_video",
|
||||
"provider_selection", "search_video", "download_video",
|
||||
]
|
||||
supports = {
|
||||
"user_preference_routing": True,
|
||||
"offline_fallback": True,
|
||||
"reference_image": True,
|
||||
"stock_fallback": True,
|
||||
}
|
||||
best_for = [
|
||||
"preflight routing",
|
||||
"user-facing recommendation flows",
|
||||
"switching between cloud, local, and stock video tools",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {"type": "string"},
|
||||
"preferred_provider": {
|
||||
"type": "string",
|
||||
"description": "Provider name or 'auto'. Valid values are discovered at runtime from the registry.",
|
||||
"default": "auto",
|
||||
},
|
||||
"allowed_providers": {"type": "array", "items": {"type": "string"}},
|
||||
"operation": {"type": "string", "enum": ["text_to_video", "image_to_video"], "default": "text_to_video"},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
def _providers(self) -> list[BaseTool]:
|
||||
"""Auto-discover video generation providers from the registry."""
|
||||
from tools.tool_registry import registry
|
||||
registry.ensure_discovered()
|
||||
return [t for t in registry.get_by_capability("video_generation")
|
||||
if t.name != self.name]
|
||||
|
||||
@property
|
||||
def fallback_tools(self) -> list[str]:
|
||||
"""Dynamically built from discovered providers + image_selector as last resort."""
|
||||
return [t.name for t in self._providers()] + ["image_selector"]
|
||||
|
||||
@property
|
||||
def provider_matrix(self) -> dict[str, dict[str, str]]:
|
||||
"""Built at runtime from each provider's best_for field."""
|
||||
matrix = {}
|
||||
for tool in self._providers():
|
||||
strength = ", ".join(tool.best_for) if tool.best_for else tool.name
|
||||
matrix[tool.provider] = {"tool": tool.name, "strength": strength}
|
||||
return matrix
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
if any(tool.get_status() == ToolStatus.AVAILABLE for tool in self._providers()):
|
||||
return ToolStatus.AVAILABLE
|
||||
return ToolStatus.UNAVAILABLE
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, object]) -> float:
|
||||
tool = self._select_tool(inputs)
|
||||
return tool.estimate_cost(inputs) if tool else 0.0
|
||||
|
||||
def estimate_runtime(self, inputs: dict[str, object]) -> float:
|
||||
tool = self._select_tool(inputs)
|
||||
return tool.estimate_runtime(inputs) if tool else 0.0
|
||||
|
||||
def execute(self, inputs: dict[str, object]) -> ToolResult:
|
||||
tool = self._select_tool(inputs)
|
||||
if tool is None:
|
||||
return ToolResult(success=False, error="No video generation provider available.")
|
||||
|
||||
# Adapt input keys: stock tools use 'query' while generators use 'prompt'
|
||||
adapted = dict(inputs)
|
||||
if hasattr(tool, 'input_schema'):
|
||||
required = tool.input_schema.get("properties", {})
|
||||
if "query" in required and "query" not in adapted:
|
||||
adapted["query"] = adapted.get("prompt", "")
|
||||
|
||||
result = tool.execute(adapted)
|
||||
if result.success:
|
||||
result.data.setdefault("selected_tool", tool.name)
|
||||
return result
|
||||
|
||||
def _select_tool(self, inputs: dict[str, object]) -> BaseTool | None:
|
||||
preferred = inputs.get("preferred_provider", "auto")
|
||||
allowed = set(inputs.get("allowed_providers") or [])
|
||||
candidates = self._providers()
|
||||
if allowed:
|
||||
candidates = [tool for tool in candidates if tool.provider in allowed]
|
||||
|
||||
env_hint = os.environ.get("VIDEO_GEN_LOCAL_MODEL", "").lower()
|
||||
env_map = {
|
||||
"wan2.1-1.3b": "wan",
|
||||
"wan2.1-14b": "wan",
|
||||
"hunyuan-1.5": "hunyuan",
|
||||
"ltx2-local": "ltx",
|
||||
"cogvideo-5b": "cogvideo",
|
||||
"cogvideo-2b": "cogvideo",
|
||||
}
|
||||
if preferred == "auto" and env_hint in env_map:
|
||||
preferred = env_map[env_hint]
|
||||
|
||||
if preferred != "auto":
|
||||
ordered = [tool for tool in candidates if tool.provider == preferred]
|
||||
ordered.extend([tool for tool in candidates if tool.provider != preferred])
|
||||
else:
|
||||
ordered = candidates
|
||||
|
||||
for tool in ordered:
|
||||
if tool.get_status() == ToolStatus.AVAILABLE:
|
||||
return tool
|
||||
return None
|
||||
@@ -0,0 +1,962 @@
|
||||
"""Video stitch tool wrapping FFmpeg.
|
||||
|
||||
Multi-clip assembly with validation, transitions, and spatial layouts.
|
||||
Supports sequential concatenation (TikTok-style stitch), crossfade/fade
|
||||
transitions, and spatial compositions (side-by-side, vertical stack,
|
||||
picture-in-picture) for duet-style content.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ResumeSupport,
|
||||
ToolResult,
|
||||
ToolStability,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class VideoStitch(BaseTool):
|
||||
name = "video_stitch"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.CORE
|
||||
capability = "video_post"
|
||||
provider = "ffmpeg"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
|
||||
dependencies = ["cmd:ffmpeg", "cmd:ffprobe"]
|
||||
install_instructions = (
|
||||
"Install FFmpeg: https://ffmpeg.org/download.html\n"
|
||||
"Windows: winget install FFmpeg\n"
|
||||
"macOS: brew install ffmpeg\n"
|
||||
"Linux: sudo apt install ffmpeg"
|
||||
)
|
||||
agent_skills = ["ffmpeg", "video_toolkit"]
|
||||
|
||||
capabilities = [
|
||||
"validate_clips",
|
||||
"stitch",
|
||||
"crossfade",
|
||||
"fade_through_black",
|
||||
"preview_stitch",
|
||||
"spatial_side_by_side",
|
||||
"spatial_vertical_stack",
|
||||
"spatial_picture_in_picture",
|
||||
]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["operation"],
|
||||
"properties": {
|
||||
"operation": {
|
||||
"type": "string",
|
||||
"enum": ["validate", "stitch", "preview_stitch", "spatial"],
|
||||
},
|
||||
"clips": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"},
|
||||
"description": "List of input video file paths",
|
||||
},
|
||||
"output_path": {"type": "string"},
|
||||
"transition": {
|
||||
"type": "string",
|
||||
"enum": ["cut", "crossfade", "fade"],
|
||||
"default": "cut",
|
||||
"description": "Transition type: cut (default), crossfade, or fade (fade-through-black)",
|
||||
},
|
||||
"transition_duration": {
|
||||
"type": "number",
|
||||
"minimum": 0.1,
|
||||
"maximum": 5.0,
|
||||
"default": 0.5,
|
||||
"description": "Transition duration in seconds",
|
||||
},
|
||||
"auto_normalize": {
|
||||
"type": "boolean",
|
||||
"default": False,
|
||||
"description": "Re-encode clips to a common format before concat if they differ",
|
||||
},
|
||||
"target_resolution": {
|
||||
"type": "string",
|
||||
"description": "Target resolution for normalization (e.g. '1920x1080')",
|
||||
},
|
||||
"target_fps": {
|
||||
"type": "integer",
|
||||
"description": "Target FPS for normalization",
|
||||
},
|
||||
"codec": {"type": "string", "default": "libx264"},
|
||||
"crf": {"type": "integer", "default": 23},
|
||||
"preset": {"type": "string", "default": "medium"},
|
||||
"profile": {
|
||||
"type": "string",
|
||||
"description": "Media profile name from media_profiles.py",
|
||||
},
|
||||
"layout": {
|
||||
"type": "string",
|
||||
"enum": ["side_by_side", "vertical_stack", "picture_in_picture"],
|
||||
"description": "Spatial layout for the spatial operation",
|
||||
},
|
||||
"pip_position": {
|
||||
"type": "string",
|
||||
"enum": ["top_left", "top_right", "bottom_left", "bottom_right"],
|
||||
"default": "bottom_right",
|
||||
"description": "Position of the PiP overlay",
|
||||
},
|
||||
"pip_scale": {
|
||||
"type": "number",
|
||||
"minimum": 0.1,
|
||||
"maximum": 0.5,
|
||||
"default": 0.3,
|
||||
"description": "Scale of PiP overlay relative to base video",
|
||||
},
|
||||
"pip_margin": {
|
||||
"type": "integer",
|
||||
"default": 10,
|
||||
"description": "Margin in pixels for PiP overlay from edges",
|
||||
},
|
||||
"dry_run": {
|
||||
"type": "boolean",
|
||||
"default": False,
|
||||
"description": "If true, return what would be done without executing",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=4, ram_mb=2048, vram_mb=0, disk_mb=5000, network_required=False
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=1, retryable_errors=["Conversion failed"])
|
||||
resume_support = ResumeSupport.FROM_START
|
||||
idempotency_key_fields = ["operation", "clips", "transition", "layout"]
|
||||
side_effects = ["writes video file to output_path"]
|
||||
user_visible_verification = [
|
||||
"Play the stitched output and verify clip ordering, transitions, and A/V sync",
|
||||
]
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
operation = inputs["operation"]
|
||||
start = time.time()
|
||||
|
||||
if inputs.get("dry_run"):
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data=self.dry_run(inputs),
|
||||
)
|
||||
|
||||
try:
|
||||
if operation == "validate":
|
||||
result = self._validate(inputs)
|
||||
elif operation == "stitch":
|
||||
result = self._stitch(inputs)
|
||||
elif operation == "preview_stitch":
|
||||
result = self._preview_stitch(inputs)
|
||||
elif operation == "spatial":
|
||||
result = self._spatial(inputs)
|
||||
else:
|
||||
return ToolResult(success=False, error=f"Unknown operation: {operation}")
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=str(e))
|
||||
|
||||
result.duration_seconds = round(time.time() - start, 2)
|
||||
return result
|
||||
|
||||
def dry_run(self, inputs: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Preflight check: validate clips and report what would happen."""
|
||||
clips = inputs.get("clips", [])
|
||||
operation = inputs.get("operation", "stitch")
|
||||
info = {
|
||||
"tool": self.name,
|
||||
"operation": operation,
|
||||
"clip_count": len(clips),
|
||||
"transition": inputs.get("transition", "cut"),
|
||||
"auto_normalize": inputs.get("auto_normalize", False),
|
||||
"estimated_cost_usd": self.estimate_cost(inputs),
|
||||
"estimated_runtime_seconds": self.estimate_runtime(inputs),
|
||||
"status": self.get_status().value,
|
||||
"would_execute": True,
|
||||
}
|
||||
if clips:
|
||||
probe_results = []
|
||||
for clip in clips:
|
||||
if Path(clip).exists():
|
||||
probe = self._probe_clip(clip)
|
||||
if probe:
|
||||
probe_results.append(probe)
|
||||
info["clip_info"] = probe_results
|
||||
return info
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Audio-stream detection and silent-audio helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _clip_has_audio(self, clip_path: str) -> bool:
|
||||
"""Return True if *clip_path* contains at least one audio stream."""
|
||||
cmd = [
|
||||
"ffprobe", "-v", "quiet",
|
||||
"-select_streams", "a",
|
||||
"-show_entries", "stream=codec_type",
|
||||
"-of", "json",
|
||||
str(clip_path),
|
||||
]
|
||||
try:
|
||||
proc = self.run_command(cmd)
|
||||
data = json.loads(proc.stdout)
|
||||
return len(data.get("streams", [])) > 0
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def _ensure_audio_for_clips(
|
||||
self,
|
||||
clips: list[str],
|
||||
temp_dir: Path,
|
||||
temp_files: list[Path],
|
||||
) -> list[str]:
|
||||
"""Return a list of clip paths where every clip is guaranteed to have
|
||||
an audio stream. Clips that already contain audio are returned as-is.
|
||||
For clips without audio, a silent stereo AAC track is muxed in and the
|
||||
path to the new file is returned instead. All generated temp files are
|
||||
appended to *temp_files* so the caller can clean them up.
|
||||
"""
|
||||
result: list[str] = []
|
||||
for i, clip in enumerate(clips):
|
||||
if self._clip_has_audio(clip):
|
||||
result.append(clip)
|
||||
else:
|
||||
augmented = temp_dir / f"audio_aug_{i:04d}.mp4"
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", str(clip),
|
||||
"-f", "lavfi", "-i", "anullsrc=r=44100:cl=stereo",
|
||||
"-c:v", "copy",
|
||||
"-c:a", "aac",
|
||||
"-shortest",
|
||||
str(augmented),
|
||||
]
|
||||
self.run_command(cmd)
|
||||
temp_files.append(augmented)
|
||||
result.append(str(augmented))
|
||||
return result
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Probe helper
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _probe_clip(self, clip_path: str) -> Optional[dict[str, Any]]:
|
||||
"""Probe a single clip with ffprobe and return metadata dict."""
|
||||
cmd = [
|
||||
"ffprobe", "-v", "quiet",
|
||||
"-print_format", "json",
|
||||
"-show_streams",
|
||||
"-show_format",
|
||||
str(clip_path),
|
||||
]
|
||||
try:
|
||||
proc = self.run_command(cmd)
|
||||
data = json.loads(proc.stdout)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
info: dict[str, Any] = {"path": str(clip_path)}
|
||||
|
||||
# Extract video stream info
|
||||
for stream in data.get("streams", []):
|
||||
if stream.get("codec_type") == "video":
|
||||
info["width"] = stream.get("width")
|
||||
info["height"] = stream.get("height")
|
||||
info["video_codec"] = stream.get("codec_name")
|
||||
info["pixel_format"] = stream.get("pix_fmt")
|
||||
# Parse fps from r_frame_rate (e.g. "30/1")
|
||||
rfr = stream.get("r_frame_rate", "0/1")
|
||||
try:
|
||||
num, den = rfr.split("/")
|
||||
info["fps"] = round(int(num) / int(den), 2)
|
||||
except (ValueError, ZeroDivisionError):
|
||||
info["fps"] = None
|
||||
break
|
||||
|
||||
# Extract audio stream info
|
||||
for stream in data.get("streams", []):
|
||||
if stream.get("codec_type") == "audio":
|
||||
info["audio_codec"] = stream.get("codec_name")
|
||||
info["sample_rate"] = stream.get("sample_rate")
|
||||
info["audio_channels"] = stream.get("channels")
|
||||
break
|
||||
|
||||
# Duration from format
|
||||
fmt = data.get("format", {})
|
||||
try:
|
||||
info["duration"] = float(fmt.get("duration", 0))
|
||||
except (TypeError, ValueError):
|
||||
info["duration"] = 0.0
|
||||
try:
|
||||
info["file_size_bytes"] = int(fmt.get("size", 0))
|
||||
except (TypeError, ValueError):
|
||||
info["file_size_bytes"] = 0
|
||||
|
||||
return info
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# validate
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _validate(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
"""Check clip compatibility: resolution, fps, codec, audio format.
|
||||
|
||||
Returns a detailed report of mismatches.
|
||||
"""
|
||||
clips = inputs.get("clips", [])
|
||||
if not clips:
|
||||
return ToolResult(success=False, error="No clips provided")
|
||||
|
||||
# Probe all clips
|
||||
probes: list[dict[str, Any]] = []
|
||||
missing: list[str] = []
|
||||
probe_errors: list[str] = []
|
||||
|
||||
for clip in clips:
|
||||
if not Path(clip).exists():
|
||||
missing.append(clip)
|
||||
continue
|
||||
info = self._probe_clip(clip)
|
||||
if info is None:
|
||||
probe_errors.append(clip)
|
||||
else:
|
||||
probes.append(info)
|
||||
|
||||
if missing:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"Clips not found: {', '.join(missing)}",
|
||||
)
|
||||
if probe_errors:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=f"Failed to probe clips: {', '.join(probe_errors)}",
|
||||
)
|
||||
|
||||
# Compare properties across clips
|
||||
mismatches: list[dict[str, Any]] = []
|
||||
reference = probes[0]
|
||||
check_fields = [
|
||||
("width", "resolution width"),
|
||||
("height", "resolution height"),
|
||||
("fps", "frame rate"),
|
||||
("video_codec", "video codec"),
|
||||
("pixel_format", "pixel format"),
|
||||
("audio_codec", "audio codec"),
|
||||
("sample_rate", "audio sample rate"),
|
||||
("audio_channels", "audio channels"),
|
||||
]
|
||||
|
||||
for i, probe in enumerate(probes[1:], start=1):
|
||||
clip_mismatches: list[str] = []
|
||||
for field_key, label in check_fields:
|
||||
ref_val = reference.get(field_key)
|
||||
cur_val = probe.get(field_key)
|
||||
if ref_val is not None and cur_val is not None and ref_val != cur_val:
|
||||
clip_mismatches.append(
|
||||
f"{label}: clip[0]={ref_val} vs clip[{i}]={cur_val}"
|
||||
)
|
||||
if clip_mismatches:
|
||||
mismatches.append({
|
||||
"clip_index": i,
|
||||
"clip_path": probe["path"],
|
||||
"differences": clip_mismatches,
|
||||
})
|
||||
|
||||
compatible = len(mismatches) == 0
|
||||
total_duration = sum(p.get("duration", 0) for p in probes)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"operation": "validate",
|
||||
"clip_count": len(clips),
|
||||
"compatible": compatible,
|
||||
"total_duration": round(total_duration, 2),
|
||||
"reference_clip": {
|
||||
"path": reference["path"],
|
||||
"resolution": f"{reference.get('width')}x{reference.get('height')}",
|
||||
"fps": reference.get("fps"),
|
||||
"video_codec": reference.get("video_codec"),
|
||||
"audio_codec": reference.get("audio_codec"),
|
||||
},
|
||||
"mismatches": mismatches,
|
||||
"clips": probes,
|
||||
},
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Normalization helper
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _resolve_normalization_target(
|
||||
self, inputs: dict[str, Any], probes: list[dict[str, Any]]
|
||||
) -> tuple[int, int, int, str, str]:
|
||||
"""Determine the target resolution, fps, and codecs for normalization.
|
||||
|
||||
Returns (width, height, fps, video_codec, audio_codec).
|
||||
"""
|
||||
# If a media profile is specified, use it
|
||||
profile_name = inputs.get("profile")
|
||||
if profile_name:
|
||||
try:
|
||||
from lib.media_profiles import get_profile
|
||||
profile = get_profile(profile_name)
|
||||
return (profile.width, profile.height, profile.fps, profile.codec, profile.audio_codec)
|
||||
except (ImportError, ValueError):
|
||||
pass
|
||||
|
||||
# Explicit target overrides
|
||||
target_w, target_h = None, None
|
||||
if inputs.get("target_resolution"):
|
||||
parts = inputs["target_resolution"].split("x")
|
||||
if len(parts) == 2:
|
||||
target_w, target_h = int(parts[0]), int(parts[1])
|
||||
|
||||
target_fps = inputs.get("target_fps")
|
||||
|
||||
# Fall back to first clip as reference
|
||||
ref = probes[0] if probes else {}
|
||||
width = target_w or ref.get("width", 1920)
|
||||
height = target_h or ref.get("height", 1080)
|
||||
fps = target_fps or ref.get("fps", 30)
|
||||
video_codec = inputs.get("codec", "libx264")
|
||||
audio_codec = "aac"
|
||||
|
||||
return (width, height, int(fps), video_codec, audio_codec)
|
||||
|
||||
def _normalize_clip(
|
||||
self,
|
||||
clip_path: str,
|
||||
output_path: Path,
|
||||
width: int,
|
||||
height: int,
|
||||
fps: int,
|
||||
video_codec: str,
|
||||
audio_codec: str,
|
||||
crf: int,
|
||||
preset: str,
|
||||
) -> None:
|
||||
"""Re-encode a clip to the target format."""
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", str(clip_path),
|
||||
"-vf", f"scale={width}:{height}:force_original_aspect_ratio=decrease,pad={width}:{height}:(ow-iw)/2:(oh-ih)/2",
|
||||
"-r", str(fps),
|
||||
"-c:v", video_codec, "-crf", str(crf), "-preset", preset,
|
||||
"-c:a", audio_codec, "-ar", "44100", "-ac", "2",
|
||||
"-pix_fmt", "yuv420p",
|
||||
str(output_path),
|
||||
]
|
||||
self.run_command(cmd)
|
||||
|
||||
def _needs_normalization(self, probes: list[dict[str, Any]]) -> bool:
|
||||
"""Check whether clips need normalization to be concat-compatible."""
|
||||
if len(probes) < 2:
|
||||
return False
|
||||
ref = probes[0]
|
||||
for probe in probes[1:]:
|
||||
for key in ("width", "height", "fps", "video_codec", "audio_codec", "sample_rate"):
|
||||
if ref.get(key) != probe.get(key) and ref.get(key) is not None:
|
||||
return True
|
||||
return False
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# stitch
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _stitch(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
"""Concatenate clips sequentially with FFmpeg concat demuxer.
|
||||
|
||||
Supports transitions: cut (default), crossfade, fade-through-black.
|
||||
"""
|
||||
clips = inputs.get("clips", [])
|
||||
if not clips:
|
||||
return ToolResult(success=False, error="No clips provided")
|
||||
if len(clips) < 2:
|
||||
return ToolResult(success=False, error="At least 2 clips required for stitch")
|
||||
|
||||
output_path = Path(inputs.get("output_path", "stitched_output.mp4"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
transition = inputs.get("transition", "cut")
|
||||
transition_dur = inputs.get("transition_duration", 0.5)
|
||||
auto_normalize = inputs.get("auto_normalize", False)
|
||||
codec = inputs.get("codec", "libx264")
|
||||
crf = inputs.get("crf", 23)
|
||||
preset = inputs.get("preset", "medium")
|
||||
|
||||
# Verify all clips exist
|
||||
for clip in clips:
|
||||
if not Path(clip).exists():
|
||||
return ToolResult(success=False, error=f"Clip not found: {clip}")
|
||||
|
||||
# Probe clips for compatibility check
|
||||
probes: list[dict[str, Any]] = []
|
||||
for clip in clips:
|
||||
info = self._probe_clip(clip)
|
||||
if info is None:
|
||||
return ToolResult(success=False, error=f"Failed to probe clip: {clip}")
|
||||
probes.append(info)
|
||||
|
||||
needs_norm = self._needs_normalization(probes)
|
||||
|
||||
# If clips are incompatible and auto_normalize is off, fail with advice
|
||||
if needs_norm and not auto_normalize and transition == "cut":
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error=(
|
||||
"Clips have mismatched properties (resolution/fps/codec). "
|
||||
"Set auto_normalize=true to re-encode to a common format, "
|
||||
"or use a transition type other than 'cut'."
|
||||
),
|
||||
)
|
||||
|
||||
temp_dir = output_path.parent / ".stitch_tmp"
|
||||
temp_dir.mkdir(parents=True, exist_ok=True)
|
||||
temp_files: list[Path] = []
|
||||
|
||||
try:
|
||||
# Normalize clips if needed
|
||||
working_clips: list[str] = []
|
||||
if needs_norm or auto_normalize or transition != "cut":
|
||||
width, height, fps, vid_codec, aud_codec = self._resolve_normalization_target(inputs, probes)
|
||||
for i, clip in enumerate(clips):
|
||||
norm_path = temp_dir / f"norm_{i:04d}.mp4"
|
||||
self._normalize_clip(clip, norm_path, width, height, fps, vid_codec, aud_codec, crf, preset)
|
||||
working_clips.append(str(norm_path))
|
||||
temp_files.append(norm_path)
|
||||
else:
|
||||
working_clips = list(clips)
|
||||
|
||||
# For crossfade/fade transitions, ensure every clip has an audio
|
||||
# stream so that the acrossfade filter does not fail. Image-derived
|
||||
# video clips typically lack audio; we add a silent track for those.
|
||||
if transition in ("crossfade", "fade"):
|
||||
working_clips = self._ensure_audio_for_clips(
|
||||
working_clips, temp_dir, temp_files,
|
||||
)
|
||||
|
||||
if transition == "cut":
|
||||
result_data = self._stitch_cut(working_clips, output_path, temp_dir, temp_files)
|
||||
elif transition == "crossfade":
|
||||
result_data = self._stitch_crossfade(working_clips, output_path, transition_dur, probes)
|
||||
elif transition == "fade":
|
||||
result_data = self._stitch_fade_through_black(working_clips, output_path, transition_dur, probes)
|
||||
else:
|
||||
return ToolResult(success=False, error=f"Unknown transition type: {transition}")
|
||||
|
||||
# Get output file info
|
||||
file_size = output_path.stat().st_size if output_path.exists() else 0
|
||||
out_probe = self._probe_clip(str(output_path))
|
||||
out_duration = out_probe.get("duration", 0) if out_probe else 0
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"operation": "stitch",
|
||||
"clip_count": len(clips),
|
||||
"transition": transition,
|
||||
"transition_duration": transition_dur if transition != "cut" else 0,
|
||||
"auto_normalized": needs_norm or auto_normalize,
|
||||
"output": str(output_path),
|
||||
"duration": round(out_duration, 2),
|
||||
"file_size_bytes": file_size,
|
||||
**result_data,
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
)
|
||||
finally:
|
||||
self._cleanup_temp(temp_dir, temp_files)
|
||||
|
||||
def _stitch_cut(
|
||||
self,
|
||||
clips: list[str],
|
||||
output_path: Path,
|
||||
temp_dir: Path,
|
||||
temp_files: list[Path],
|
||||
) -> dict[str, Any]:
|
||||
"""Simple concat via FFmpeg concat demuxer (no transition)."""
|
||||
concat_list = temp_dir / "concat_list.txt"
|
||||
temp_files.append(concat_list)
|
||||
with open(concat_list, "w", encoding="utf-8") as f:
|
||||
for clip in clips:
|
||||
safe_path = str(Path(clip).resolve()).replace("\\", "/")
|
||||
f.write(f"file '{safe_path}'\n")
|
||||
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-f", "concat", "-safe", "0",
|
||||
"-i", str(concat_list),
|
||||
"-c", "copy",
|
||||
str(output_path),
|
||||
]
|
||||
self.run_command(cmd)
|
||||
return {"method": "concat_demuxer"}
|
||||
|
||||
def _stitch_crossfade(
|
||||
self,
|
||||
clips: list[str],
|
||||
output_path: Path,
|
||||
duration: float,
|
||||
probes: list[dict[str, Any]],
|
||||
) -> dict[str, Any]:
|
||||
"""Crossfade between adjacent clips using xfade filter."""
|
||||
if len(clips) == 2:
|
||||
# Simple two-clip crossfade
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", clips[0],
|
||||
"-i", clips[1],
|
||||
"-filter_complex",
|
||||
f"[0:v][1:v]xfade=transition=fade:duration={duration}:offset={self._get_xfade_offset(probes, 0, duration)}[v];"
|
||||
f"[0:a][1:a]acrossfade=d={duration}[a]",
|
||||
"-map", "[v]", "-map", "[a]",
|
||||
str(output_path),
|
||||
]
|
||||
self.run_command(cmd)
|
||||
else:
|
||||
# Chain crossfades for N clips
|
||||
self._chain_xfade(clips, output_path, duration, probes, transition="fade")
|
||||
return {"method": "xfade_crossfade"}
|
||||
|
||||
def _stitch_fade_through_black(
|
||||
self,
|
||||
clips: list[str],
|
||||
output_path: Path,
|
||||
duration: float,
|
||||
probes: list[dict[str, Any]],
|
||||
) -> dict[str, Any]:
|
||||
"""Fade-through-black between adjacent clips using xfade fadeblack."""
|
||||
if len(clips) == 2:
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", clips[0],
|
||||
"-i", clips[1],
|
||||
"-filter_complex",
|
||||
f"[0:v][1:v]xfade=transition=fadeblack:duration={duration}:offset={self._get_xfade_offset(probes, 0, duration)}[v];"
|
||||
f"[0:a][1:a]acrossfade=d={duration}[a]",
|
||||
"-map", "[v]", "-map", "[a]",
|
||||
str(output_path),
|
||||
]
|
||||
self.run_command(cmd)
|
||||
else:
|
||||
self._chain_xfade(clips, output_path, duration, probes, transition="fadeblack")
|
||||
return {"method": "xfade_fadeblack"}
|
||||
|
||||
def _get_xfade_offset(
|
||||
self, probes: list[dict[str, Any]], clip_index: int, duration: float
|
||||
) -> float:
|
||||
"""Calculate xfade offset for a given clip pair.
|
||||
|
||||
The offset is the timestamp in the output where the transition starts,
|
||||
which equals the duration of the first clip minus the transition duration.
|
||||
"""
|
||||
clip_dur = probes[clip_index].get("duration", 0) if clip_index < len(probes) else 0
|
||||
offset = max(0, clip_dur - duration)
|
||||
return round(offset, 3)
|
||||
|
||||
def _chain_xfade(
|
||||
self,
|
||||
clips: list[str],
|
||||
output_path: Path,
|
||||
duration: float,
|
||||
probes: list[dict[str, Any]],
|
||||
transition: str,
|
||||
) -> None:
|
||||
"""Chain xfade filters for N > 2 clips.
|
||||
|
||||
Builds a complex filtergraph that progressively applies xfade
|
||||
between each adjacent pair of clips.
|
||||
"""
|
||||
n = len(clips)
|
||||
input_args: list[str] = []
|
||||
for clip in clips:
|
||||
input_args.extend(["-i", clip])
|
||||
|
||||
# Calculate cumulative offsets
|
||||
# Each xfade offset = cumulative duration of all previous segments
|
||||
# minus cumulative transition overlaps minus current transition duration
|
||||
video_filters: list[str] = []
|
||||
audio_filters: list[str] = []
|
||||
cumulative_offset = 0.0
|
||||
|
||||
for i in range(n - 1):
|
||||
clip_dur = probes[i].get("duration", 0) if i < len(probes) else 0
|
||||
offset = round(cumulative_offset + clip_dur - duration, 3)
|
||||
offset = max(0, offset)
|
||||
|
||||
if i == 0:
|
||||
v_in1 = "[0:v]"
|
||||
a_in1 = "[0:a]"
|
||||
else:
|
||||
v_in1 = f"[vfade{i-1}]"
|
||||
a_in1 = f"[afade{i-1}]"
|
||||
|
||||
v_in2 = f"[{i+1}:v]"
|
||||
a_in2 = f"[{i+1}:a]"
|
||||
|
||||
if i < n - 2:
|
||||
v_out = f"[vfade{i}]"
|
||||
a_out = f"[afade{i}]"
|
||||
else:
|
||||
v_out = "[vout]"
|
||||
a_out = "[aout]"
|
||||
|
||||
video_filters.append(
|
||||
f"{v_in1}{v_in2}xfade=transition={transition}:duration={duration}:offset={offset}{v_out}"
|
||||
)
|
||||
audio_filters.append(
|
||||
f"{a_in1}{a_in2}acrossfade=d={duration}{a_out}"
|
||||
)
|
||||
|
||||
# Cumulative offset advances by clip duration minus overlap
|
||||
cumulative_offset = offset
|
||||
|
||||
filter_complex = ";".join(video_filters + audio_filters)
|
||||
|
||||
cmd = ["ffmpeg", "-y"]
|
||||
cmd.extend(input_args)
|
||||
cmd.extend(["-filter_complex", filter_complex])
|
||||
cmd.extend(["-map", "[vout]", "-map", "[aout]"])
|
||||
cmd.append(str(output_path))
|
||||
self.run_command(cmd)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# preview_stitch
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _preview_stitch(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
"""Generate a low-resolution preview of the stitched result."""
|
||||
clips = inputs.get("clips", [])
|
||||
if not clips:
|
||||
return ToolResult(success=False, error="No clips provided")
|
||||
if len(clips) < 2:
|
||||
return ToolResult(success=False, error="At least 2 clips required for preview")
|
||||
|
||||
output_path = Path(inputs.get("output_path", "stitch_preview.mp4"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Verify all clips exist
|
||||
for clip in clips:
|
||||
if not Path(clip).exists():
|
||||
return ToolResult(success=False, error=f"Clip not found: {clip}")
|
||||
|
||||
# Build preview by normalizing to low-res and stitching
|
||||
preview_inputs = dict(inputs)
|
||||
preview_inputs["auto_normalize"] = True
|
||||
preview_inputs["target_resolution"] = "640x360"
|
||||
preview_inputs["target_fps"] = 24
|
||||
preview_inputs["crf"] = 30
|
||||
preview_inputs["preset"] = "ultrafast"
|
||||
preview_inputs["output_path"] = str(output_path)
|
||||
|
||||
# Delegate to _stitch with preview settings
|
||||
result = self._stitch(preview_inputs)
|
||||
|
||||
if result.success:
|
||||
result.data["operation"] = "preview_stitch"
|
||||
result.data["preview"] = True
|
||||
result.data["preview_resolution"] = "640x360"
|
||||
|
||||
return result
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# spatial
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _spatial(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
"""Side-by-side, vertical stack, or picture-in-picture layouts.
|
||||
|
||||
Designed for TikTok Stitch/Duet style compositions (D3.5.8).
|
||||
"""
|
||||
clips = inputs.get("clips", [])
|
||||
if not clips or len(clips) < 2:
|
||||
return ToolResult(
|
||||
success=False,
|
||||
error="At least 2 clips required for spatial layout",
|
||||
)
|
||||
|
||||
layout = inputs.get("layout")
|
||||
if not layout:
|
||||
return ToolResult(success=False, error="layout is required for spatial operation")
|
||||
|
||||
output_path = Path(inputs.get("output_path", "spatial_output.mp4"))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
codec = inputs.get("codec", "libx264")
|
||||
crf = inputs.get("crf", 23)
|
||||
|
||||
# Verify all clips exist
|
||||
for clip in clips:
|
||||
if not Path(clip).exists():
|
||||
return ToolResult(success=False, error=f"Clip not found: {clip}")
|
||||
|
||||
temp_dir = output_path.parent / ".spatial_tmp"
|
||||
temp_dir.mkdir(parents=True, exist_ok=True)
|
||||
temp_files: list[Path] = []
|
||||
|
||||
try:
|
||||
# side_by_side and vertical_stack use amix which requires audio
|
||||
# on both inputs. Ensure silent tracks for audio-less clips.
|
||||
working_clips = list(clips)
|
||||
if layout in ("side_by_side", "vertical_stack"):
|
||||
working_clips = self._ensure_audio_for_clips(
|
||||
working_clips, temp_dir, temp_files,
|
||||
)
|
||||
|
||||
if layout == "side_by_side":
|
||||
self._spatial_side_by_side(working_clips, output_path, codec, crf)
|
||||
elif layout == "vertical_stack":
|
||||
self._spatial_vertical_stack(working_clips, output_path, codec, crf)
|
||||
elif layout == "picture_in_picture":
|
||||
self._spatial_pip(working_clips, output_path, inputs, codec, crf)
|
||||
else:
|
||||
return ToolResult(success=False, error=f"Unknown layout: {layout}")
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=str(e))
|
||||
finally:
|
||||
self._cleanup_temp(temp_dir, temp_files)
|
||||
|
||||
file_size = output_path.stat().st_size if output_path.exists() else 0
|
||||
out_probe = self._probe_clip(str(output_path))
|
||||
out_duration = out_probe.get("duration", 0) if out_probe else 0
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"operation": "spatial",
|
||||
"layout": layout,
|
||||
"clip_count": len(clips),
|
||||
"output": str(output_path),
|
||||
"duration": round(out_duration, 2),
|
||||
"file_size_bytes": file_size,
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
)
|
||||
|
||||
def _spatial_side_by_side(
|
||||
self, clips: list[str], output_path: Path, codec: str, crf: int
|
||||
) -> None:
|
||||
"""Place clips side by side (horizontal split).
|
||||
|
||||
Both clips are scaled to the same height and placed left-right.
|
||||
Uses the first two clips; additional clips are ignored.
|
||||
"""
|
||||
input_args = ["-i", clips[0], "-i", clips[1]]
|
||||
filter_complex = (
|
||||
"[0:v]scale=-2:480[left];"
|
||||
"[1:v]scale=-2:480[right];"
|
||||
"[left][right]hstack=inputs=2[v];"
|
||||
"[0:a][1:a]amix=inputs=2:duration=shortest[a]"
|
||||
)
|
||||
cmd = ["ffmpeg", "-y"]
|
||||
cmd.extend(input_args)
|
||||
cmd.extend([
|
||||
"-filter_complex", filter_complex,
|
||||
"-map", "[v]", "-map", "[a]",
|
||||
"-c:v", codec, "-crf", str(crf),
|
||||
"-c:a", "aac",
|
||||
"-shortest",
|
||||
str(output_path),
|
||||
])
|
||||
self.run_command(cmd)
|
||||
|
||||
def _spatial_vertical_stack(
|
||||
self, clips: list[str], output_path: Path, codec: str, crf: int
|
||||
) -> None:
|
||||
"""Place clips in a vertical stack (top-bottom).
|
||||
|
||||
Both clips are scaled to the same width and stacked vertically.
|
||||
Ideal for portrait/mobile viewing.
|
||||
"""
|
||||
input_args = ["-i", clips[0], "-i", clips[1]]
|
||||
filter_complex = (
|
||||
"[0:v]scale=540:-2[top];"
|
||||
"[1:v]scale=540:-2[bottom];"
|
||||
"[top][bottom]vstack=inputs=2[v];"
|
||||
"[0:a][1:a]amix=inputs=2:duration=shortest[a]"
|
||||
)
|
||||
cmd = ["ffmpeg", "-y"]
|
||||
cmd.extend(input_args)
|
||||
cmd.extend([
|
||||
"-filter_complex", filter_complex,
|
||||
"-map", "[v]", "-map", "[a]",
|
||||
"-c:v", codec, "-crf", str(crf),
|
||||
"-c:a", "aac",
|
||||
"-shortest",
|
||||
str(output_path),
|
||||
])
|
||||
self.run_command(cmd)
|
||||
|
||||
def _spatial_pip(
|
||||
self,
|
||||
clips: list[str],
|
||||
output_path: Path,
|
||||
inputs: dict[str, Any],
|
||||
codec: str,
|
||||
crf: int,
|
||||
) -> None:
|
||||
"""Picture-in-picture: overlay second clip on first.
|
||||
|
||||
clips[0] is the base (full-screen), clips[1] is the PiP overlay.
|
||||
"""
|
||||
pip_position = inputs.get("pip_position", "bottom_right")
|
||||
pip_scale = inputs.get("pip_scale", 0.3)
|
||||
pip_margin = inputs.get("pip_margin", 10)
|
||||
|
||||
# Build position expression based on corner
|
||||
position_map = {
|
||||
"top_left": f"{pip_margin}:{pip_margin}",
|
||||
"top_right": f"main_w-overlay_w-{pip_margin}:{pip_margin}",
|
||||
"bottom_left": f"{pip_margin}:main_h-overlay_h-{pip_margin}",
|
||||
"bottom_right": f"main_w-overlay_w-{pip_margin}:main_h-overlay_h-{pip_margin}",
|
||||
}
|
||||
position = position_map.get(pip_position, position_map["bottom_right"])
|
||||
|
||||
input_args = ["-i", clips[0], "-i", clips[1]]
|
||||
filter_complex = (
|
||||
f"[1:v]scale=iw*{pip_scale}:ih*{pip_scale}[pip];"
|
||||
f"[0:v][pip]overlay={position}:shortest=1[v]"
|
||||
)
|
||||
cmd = ["ffmpeg", "-y"]
|
||||
cmd.extend(input_args)
|
||||
cmd.extend([
|
||||
"-filter_complex", filter_complex,
|
||||
"-map", "[v]", "-map", "0:a?",
|
||||
"-c:v", codec, "-crf", str(crf),
|
||||
"-c:a", "aac",
|
||||
"-shortest",
|
||||
str(output_path),
|
||||
])
|
||||
self.run_command(cmd)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Cleanup
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _cleanup_temp(temp_dir: Path, temp_files: list[Path]) -> None:
|
||||
"""Remove temporary files and directory."""
|
||||
for f in temp_files:
|
||||
if f.exists():
|
||||
try:
|
||||
f.unlink()
|
||||
except OSError:
|
||||
pass
|
||||
if temp_dir.exists():
|
||||
try:
|
||||
temp_dir.rmdir()
|
||||
except OSError:
|
||||
pass
|
||||
@@ -0,0 +1,270 @@
|
||||
"""Video trimmer tool wrapping FFmpeg.
|
||||
|
||||
Provides cut, trim, speed adjustment, and concatenation of video segments.
|
||||
All operations are deterministic and produce lossless or near-lossless output
|
||||
by default.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ResumeSupport,
|
||||
ToolResult,
|
||||
ToolStability,
|
||||
ToolTier,
|
||||
)
|
||||
|
||||
|
||||
class VideoTrimmer(BaseTool):
|
||||
name = "video_trimmer"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.CORE
|
||||
capability = "video_post"
|
||||
provider = "ffmpeg"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.DETERMINISTIC
|
||||
|
||||
dependencies = ["cmd:ffmpeg"]
|
||||
install_instructions = (
|
||||
"Install FFmpeg: https://ffmpeg.org/download.html\n"
|
||||
"Windows: winget install FFmpeg\n"
|
||||
"macOS: brew install ffmpeg\n"
|
||||
"Linux: sudo apt install ffmpeg"
|
||||
)
|
||||
agent_skills = ["ffmpeg", "video_toolkit"]
|
||||
|
||||
capabilities = ["cut", "trim", "speed_adjust", "concat"]
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["operation"],
|
||||
"properties": {
|
||||
"operation": {
|
||||
"type": "string",
|
||||
"enum": ["cut", "speed", "concat"],
|
||||
},
|
||||
"input_path": {"type": "string"},
|
||||
"output_path": {"type": "string"},
|
||||
"start_seconds": {"type": "number", "minimum": 0},
|
||||
"end_seconds": {"type": "number", "minimum": 0},
|
||||
"speed_factor": {"type": "number", "minimum": 0.1, "maximum": 100.0},
|
||||
"segments": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"input_path": {"type": "string"},
|
||||
"start_seconds": {"type": "number"},
|
||||
"end_seconds": {"type": "number"},
|
||||
},
|
||||
},
|
||||
},
|
||||
"codec": {"type": "string", "default": "copy"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(
|
||||
cpu_cores=2, ram_mb=1024, vram_mb=0, disk_mb=2000, network_required=False
|
||||
)
|
||||
retry_policy = RetryPolicy(max_retries=1, retryable_errors=["FFmpeg error"])
|
||||
resume_support = ResumeSupport.FROM_START
|
||||
idempotency_key_fields = ["operation", "input_path", "start_seconds", "end_seconds", "speed_factor"]
|
||||
side_effects = ["writes video file to output_path"]
|
||||
user_visible_verification = ["Play trimmed output and verify cut points"]
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
operation = inputs["operation"]
|
||||
start = time.time()
|
||||
|
||||
try:
|
||||
if operation == "cut":
|
||||
result = self._cut(inputs)
|
||||
elif operation == "speed":
|
||||
result = self._speed(inputs)
|
||||
elif operation == "concat":
|
||||
result = self._concat(inputs)
|
||||
else:
|
||||
return ToolResult(success=False, error=f"Unknown operation: {operation}")
|
||||
except Exception as e:
|
||||
return ToolResult(success=False, error=str(e))
|
||||
|
||||
result.duration_seconds = round(time.time() - start, 2)
|
||||
return result
|
||||
|
||||
def _cut(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
input_path = Path(inputs["input_path"])
|
||||
if not input_path.exists():
|
||||
return ToolResult(success=False, error=f"Input not found: {input_path}")
|
||||
|
||||
start_s = inputs.get("start_seconds", 0)
|
||||
end_s = inputs.get("end_seconds")
|
||||
codec = inputs.get("codec", "copy")
|
||||
output_path = Path(
|
||||
inputs.get("output_path", str(input_path.with_stem(f"{input_path.stem}_cut")))
|
||||
)
|
||||
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", str(input_path),
|
||||
"-ss", str(start_s),
|
||||
]
|
||||
if end_s is not None:
|
||||
cmd.extend(["-to", str(end_s)])
|
||||
if codec == "copy":
|
||||
cmd.extend(["-c", "copy"])
|
||||
else:
|
||||
cmd.extend(["-c:v", codec, "-c:a", "aac"])
|
||||
cmd.append(str(output_path))
|
||||
|
||||
self.run_command(cmd)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"operation": "cut",
|
||||
"input": str(input_path),
|
||||
"output": str(output_path),
|
||||
"start_seconds": start_s,
|
||||
"end_seconds": end_s,
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
)
|
||||
|
||||
def _speed(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
input_path = Path(inputs["input_path"])
|
||||
if not input_path.exists():
|
||||
return ToolResult(success=False, error=f"Input not found: {input_path}")
|
||||
|
||||
factor = inputs.get("speed_factor", 1.0)
|
||||
output_path = Path(
|
||||
inputs.get("output_path", str(input_path.with_stem(f"{input_path.stem}_speed")))
|
||||
)
|
||||
|
||||
# Video: setpts adjusts presentation timestamps (inverse of speed)
|
||||
# Audio: atempo adjusts audio speed (must chain for >2x)
|
||||
video_filter = f"setpts={1.0/factor}*PTS"
|
||||
audio_filters = self._build_atempo_chain(factor)
|
||||
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-i", str(input_path),
|
||||
"-filter:v", video_filter,
|
||||
"-filter:a", audio_filters,
|
||||
"-c:v", "libx264", "-preset", "fast",
|
||||
"-c:a", "aac",
|
||||
str(output_path),
|
||||
]
|
||||
|
||||
self.run_command(cmd)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"operation": "speed",
|
||||
"input": str(input_path),
|
||||
"output": str(output_path),
|
||||
"speed_factor": factor,
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
)
|
||||
|
||||
def _concat(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
segments = inputs.get("segments", [])
|
||||
if not segments:
|
||||
return ToolResult(success=False, error="No segments provided for concat")
|
||||
|
||||
output_path = Path(inputs.get("output_path", "concat_output.mp4"))
|
||||
|
||||
# First, cut each segment to a temp file if start/end are specified
|
||||
temp_files: list[Path] = []
|
||||
temp_dir = output_path.parent / ".concat_tmp"
|
||||
temp_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
try:
|
||||
for i, seg in enumerate(segments):
|
||||
seg_input = Path(seg["input_path"])
|
||||
if not seg_input.exists():
|
||||
return ToolResult(success=False, error=f"Segment input not found: {seg_input}")
|
||||
|
||||
seg_start = seg.get("start_seconds")
|
||||
seg_end = seg.get("end_seconds")
|
||||
|
||||
if seg_start is not None or seg_end is not None:
|
||||
temp_path = temp_dir / f"seg_{i:04d}{seg_input.suffix}"
|
||||
cmd = ["ffmpeg", "-y", "-i", str(seg_input)]
|
||||
if seg_start is not None:
|
||||
cmd.extend(["-ss", str(seg_start)])
|
||||
if seg_end is not None:
|
||||
cmd.extend(["-to", str(seg_end)])
|
||||
cmd.extend(["-c", "copy", str(temp_path)])
|
||||
self.run_command(cmd)
|
||||
temp_files.append(temp_path)
|
||||
else:
|
||||
temp_files.append(seg_input)
|
||||
|
||||
# Write concat file list
|
||||
list_path = temp_dir / "concat_list.txt"
|
||||
with open(list_path, "w", encoding="utf-8") as f:
|
||||
for tf in temp_files:
|
||||
# FFmpeg concat demuxer needs forward slashes and escaped quotes
|
||||
safe_path = str(tf.resolve()).replace("\\", "/")
|
||||
f.write(f"file '{safe_path}'\n")
|
||||
|
||||
cmd = [
|
||||
"ffmpeg", "-y",
|
||||
"-f", "concat", "-safe", "0",
|
||||
"-i", str(list_path),
|
||||
"-c", "copy",
|
||||
str(output_path),
|
||||
]
|
||||
self.run_command(cmd)
|
||||
|
||||
return ToolResult(
|
||||
success=True,
|
||||
data={
|
||||
"operation": "concat",
|
||||
"segment_count": len(segments),
|
||||
"output": str(output_path),
|
||||
},
|
||||
artifacts=[str(output_path)],
|
||||
)
|
||||
finally:
|
||||
# Clean up temp segment files (but not the originals)
|
||||
for tf in temp_files:
|
||||
if tf.parent == temp_dir and tf.exists():
|
||||
tf.unlink()
|
||||
if list_path.exists():
|
||||
list_path.unlink()
|
||||
if temp_dir.exists():
|
||||
try:
|
||||
temp_dir.rmdir()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _build_atempo_chain(factor: float) -> str:
|
||||
"""Build an atempo filter chain. atempo only accepts [0.5, 100.0]."""
|
||||
if factor <= 0:
|
||||
factor = 1.0
|
||||
# Chain multiple atempo filters for extreme values
|
||||
filters = []
|
||||
remaining = factor
|
||||
while remaining > 100.0:
|
||||
filters.append("atempo=100.0")
|
||||
remaining /= 100.0
|
||||
while remaining < 0.5:
|
||||
filters.append("atempo=0.5")
|
||||
remaining /= 0.5
|
||||
filters.append(f"atempo={remaining:.4f}")
|
||||
return ",".join(filters)
|
||||
@@ -0,0 +1,98 @@
|
||||
"""Wan local video generation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from tools.base_tool import (
|
||||
BaseTool,
|
||||
Determinism,
|
||||
ExecutionMode,
|
||||
ResourceProfile,
|
||||
RetryPolicy,
|
||||
ToolResult,
|
||||
ToolRuntime,
|
||||
ToolStability,
|
||||
ToolStatus,
|
||||
ToolTier,
|
||||
)
|
||||
from tools.video._shared import WAN_VARIANTS, estimate_local_runtime, generate_local_video, local_generation_status, local_install_instructions
|
||||
|
||||
|
||||
class WanVideo(BaseTool):
|
||||
name = "wan_video"
|
||||
version = "0.1.0"
|
||||
tier = ToolTier.GENERATE
|
||||
capability = "video_generation"
|
||||
provider = "wan"
|
||||
stability = ToolStability.EXPERIMENTAL
|
||||
execution_mode = ExecutionMode.SYNC
|
||||
determinism = Determinism.STOCHASTIC
|
||||
runtime = ToolRuntime.LOCAL_GPU
|
||||
|
||||
install_instructions = local_install_instructions()
|
||||
fallback = "hunyuan_video"
|
||||
fallback_tools = ["hunyuan_video", "ltx_video_local", "cogvideo_video", "image_selector"]
|
||||
agent_skills = ["ltx2"]
|
||||
|
||||
capabilities = ["text_to_video", "image_to_video", "model_selection"]
|
||||
supports = {
|
||||
"reference_image": True,
|
||||
"offline": True,
|
||||
"native_audio": False,
|
||||
"local_gpu": True,
|
||||
}
|
||||
best_for = [
|
||||
"best quality-to-VRAM ratio for local generation",
|
||||
"local pipelines that still want image-to-video support",
|
||||
]
|
||||
not_good_for = ["CPU-only machines", "instant iteration on low-end hardware"]
|
||||
provider_matrix = {key: {"tool": "wan_video", **value, "mode": "local_gpu"} for key, value in WAN_VARIANTS.items()}
|
||||
|
||||
input_schema = {
|
||||
"type": "object",
|
||||
"required": ["prompt"],
|
||||
"properties": {
|
||||
"prompt": {"type": "string"},
|
||||
"operation": {"type": "string", "enum": ["text_to_video", "image_to_video"], "default": "text_to_video"},
|
||||
"model_variant": {"type": "string", "enum": sorted(WAN_VARIANTS), "default": "wan2.1-1.3b"},
|
||||
"reference_image_url": {"type": "string"},
|
||||
"reference_image_path": {"type": "string"},
|
||||
"width": {"type": "integer"},
|
||||
"height": {"type": "integer"},
|
||||
"num_frames": {"type": "integer"},
|
||||
"num_inference_steps": {"type": "integer"},
|
||||
"enable_model_offload": {"type": "boolean", "default": True},
|
||||
"seed": {"type": "integer"},
|
||||
"output_path": {"type": "string"},
|
||||
},
|
||||
}
|
||||
|
||||
resource_profile = ResourceProfile(cpu_cores=2, ram_mb=16000, vram_mb=8000, disk_mb=4000, network_required=False)
|
||||
retry_policy = RetryPolicy(max_retries=1)
|
||||
idempotency_key_fields = ["prompt", "model_variant", "operation", "seed"]
|
||||
side_effects = ["writes video file to output_path", "may download model weights"]
|
||||
user_visible_verification = ["Watch generated clip for motion coherence and artifacts"]
|
||||
|
||||
def get_status(self) -> ToolStatus:
|
||||
return local_generation_status()
|
||||
|
||||
def estimate_cost(self, inputs: dict[str, Any]) -> float:
|
||||
return 0.0
|
||||
|
||||
def estimate_runtime(self, inputs: dict[str, Any]) -> float:
|
||||
variant = WAN_VARIANTS.get(inputs.get("model_variant", "wan2.1-1.3b"), WAN_VARIANTS["wan2.1-1.3b"])
|
||||
return estimate_local_runtime(variant["speed"])
|
||||
|
||||
def execute(self, inputs: dict[str, Any]) -> ToolResult:
|
||||
if self.get_status() != ToolStatus.AVAILABLE:
|
||||
return ToolResult(success=False, error="Wan local video generation is unavailable. " + self.install_instructions)
|
||||
start = time.time()
|
||||
try:
|
||||
result = generate_local_video(tool_name=self.name, variants=WAN_VARIANTS, default_variant="wan2.1-1.3b", inputs=inputs)
|
||||
except Exception as exc:
|
||||
return ToolResult(success=False, error=f"Wan video generation failed: {exc}")
|
||||
result.duration_seconds = round(time.time() - start, 2)
|
||||
return result
|
||||
|
||||
Reference in New Issue
Block a user