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,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