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:
calesthio
@@ -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
|
||||
Reference in New Issue
Block a user