@_vrsen/openswarm 0.1.0

package/video_generation_agent/tools/TrimVideo.py

@@ -0,0 +1,233 @@
+"""Tool for trimming videos from start and/or end."""
+
+import os
+import subprocess
+import asyncio
+import cv2
+from typing import Optional
+from pathlib import Path
+
+from pydantic import Field, field_validator, model_validator
+from PIL import Image
+
+from agency_swarm import BaseTool, ToolOutputText
+
+from .utils.video_utils import (
+    get_videos_dir,
+    ensure_not_blank,
+    generate_spritesheet,
+    extract_last_frame,
+    create_image_output,
+    resolve_ffmpeg_executable,
+)
+
+
+class TrimVideo(BaseTool):
+    """
+    Trim a video by removing seconds from the start and/or end.
+    
+    Useful for removing unwanted intro/outro segments or adjusting video length.
+    
+    Videos are saved to: mnt/{product_name}/generated_videos/
+    """
+    
+    product_name: str = Field(
+        ...,
+        description="Name of the product this video is for (e.g., 'Acme_Widget_Pro', 'Green_Tea_Extract'). Used to organize files into product-specific folders.",
+    )
+    input_video: str = Field(
+        ...,
+        description=(
+            "The video to trim. Can be: "
+            "1) Video name without extension (searches generated_videos folder), "
+            "2) Full local path to video file."
+        ),
+    )
+    output_name: str = Field(
+        ...,
+        description="The name for the trimmed video file (without extension)",
+    )
+    trim_start: Optional[float] = Field(
+        default=None,
+        description="Seconds to trim from the start of the video (optional, defaults to 0.0 if not provided)",
+    )
+    trim_end: Optional[float] = Field(
+        default=None,
+        description="Seconds to trim from the end of the video (optional, defaults to 0.0 if not provided)",
+    )
+    
+    @field_validator("input_video")
+    @classmethod
+    def _input_not_blank(cls, value: str) -> str:
+        return ensure_not_blank(value, "input_video")
+    
+    @field_validator("output_name")
+    @classmethod
+    def _output_not_blank(cls, value: str) -> str:
+        return ensure_not_blank(value, "output_name")
+    
+    @model_validator(mode='after')
+    def _validate_and_set_defaults(self):
+        """Set defaults and validate that at least one trim parameter is provided (> 0)"""
+        # Set defaults
+        if self.trim_start is None:
+            self.trim_start = 0.0
+        if self.trim_end is None:
+            self.trim_end = 0.0
+        
+        # Validate non-negative
+        if self.trim_start < 0:
+            raise ValueError("trim_start must be non-negative")
+        if self.trim_end < 0:
+            raise ValueError("trim_end must be non-negative")
+        
+        # Ensure at least one is provided
+        if self.trim_start == 0.0 and self.trim_end == 0.0:
+            raise ValueError(
+                "At least one of trim_start or trim_end must be greater than 0. "
+                "Provide at least one trim value."
+            )
+        
+        return self
+    
+    async def run(self) -> list:
+        """Trim the video and save with metadata."""
+        input_path = self._resolve_video_path()
+        videos_dir = get_videos_dir(self.product_name)
+        output_path = os.path.join(videos_dir, f"{self.output_name}.mp4")
+
+        loop = asyncio.get_event_loop()
+        await loop.run_in_executor(None, self._trim_video_blocking, input_path, output_path)
+
+        output = []
+
+        spritesheet_path = os.path.join(videos_dir, f"{self.output_name}_spritesheet.jpg")
+        spritesheet = await loop.run_in_executor(None, generate_spritesheet, output_path, spritesheet_path)
+        if spritesheet:
+            output.extend(create_image_output(spritesheet_path, f"{self.output_name}_spritesheet.jpg"))
+
+        thumbnail_path = os.path.join(videos_dir, f"{self.output_name}_thumbnail.jpg")
+        thumbnail = await loop.run_in_executor(None, self._extract_first_frame, output_path, thumbnail_path)
+        if thumbnail:
+            output.extend(create_image_output(thumbnail_path, f"{self.output_name}_thumbnail.jpg"))
+
+        last_frame_path = os.path.join(videos_dir, f"{self.output_name}_last_frame.jpg")
+        last_frame = await loop.run_in_executor(None, extract_last_frame, output_path, last_frame_path)
+        if last_frame:
+            output.extend(create_image_output(last_frame_path, f"{self.output_name}_last_frame.jpg"))
+
+        output.append(ToolOutputText(
+            type="text",
+            text=f"Video trimmed successfully!\nSaved to: `{self.output_name}.mp4`\nPath: {output_path}\nTrimmed: {self.trim_start}s from start, {self.trim_end}s from end"
+        ))
+        
+        return output
+    
+    def _resolve_video_path(self) -> str:
+        """Resolve input video to full path."""
+        # Try as full path first
+        path = Path(self.input_video).expanduser().resolve()
+        
+        if path.exists():
+            return str(path)
+        
+        # Try as video name without extension in generated_videos
+        videos_dir = get_videos_dir(self.product_name)
+        
+        for ext in [".mp4", ".mov", ".avi", ".webm"]:
+            potential_path = os.path.join(videos_dir, f"{self.input_video}{ext}")
+            if os.path.exists(potential_path):
+                return potential_path
+        
+        raise FileNotFoundError(
+            f"Video '{self.input_video}' not found in {videos_dir}. "
+            f"Tried extensions: .mp4, .mov, .avi, .webm"
+        )
+    
+    def _trim_video_blocking(self, input_path: str, output_path: str) -> None:
+        """Trim video using ffmpeg (blocking operation)."""
+        cap = cv2.VideoCapture(input_path)
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        original_duration = total_frames / fps
+        cap.release()
+
+        output_duration = original_duration - self.trim_start - self.trim_end
+
+        if output_duration <= 0:
+            raise ValueError(
+                f"Invalid trim parameters: output duration would be {output_duration:.2f}s. "
+                f"Video duration is {original_duration:.2f}s. "
+                f"Reduce trim_start ({self.trim_start}s) and/or trim_end ({self.trim_end}s)."
+            )
+
+        # Using -ss (start time) and -t (duration) for precise trimming.
+        # -c copy would be fastest but may have keyframe issues, so we re-encode with H.264.
+        ffmpeg_executable = resolve_ffmpeg_executable()
+        ffmpeg_cmd = [
+            ffmpeg_executable,
+            "-y",  # Overwrite output file if it exists
+            "-i", input_path,  # Input file
+            "-ss", str(self.trim_start),  # Start time
+            "-t", str(output_duration),  # Duration
+            "-c:v", "libx264",  # Video codec: H.264
+            "-preset", "medium",  # Encoding preset (balance speed/quality)
+            "-crf", "23",  # Constant Rate Factor (quality: 0-51, lower is better)
+            "-c:a", "aac",  # Audio codec
+            "-b:a", "128k",  # Audio bitrate
+            "-movflags", "+faststart",  # Enable fast start for web playback
+            "-pix_fmt", "yuv420p",  # Pixel format for compatibility
+            output_path
+        ]
+        
+        try:
+            subprocess.run(
+                ffmpeg_cmd,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                check=True
+            )
+        except subprocess.CalledProcessError as e:
+            raise RuntimeError(
+                f"ffmpeg failed to trim video. Error: {e.stderr}"
+            )
+        except FileNotFoundError:
+            raise RuntimeError(
+                "ffmpeg not found. Please install ffmpeg and ensure it's in your PATH. "
+                "Download from: https://ffmpeg.org/download.html"
+            )
+    
+    def _extract_first_frame(self, video_path: str, output_path: str) -> Optional[object]:
+        """Extract the first frame from video as thumbnail."""
+        cap = cv2.VideoCapture(video_path)
+        ret, frame = cap.read()
+        cap.release()
+        
+        if not ret:
+            return None
+
+        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        thumbnail_image = Image.fromarray(frame_rgb)
+        thumbnail_image.save(output_path)
+        
+        return thumbnail_image
+
+
+if __name__ == "__main__":
+    # Example: Trim only from start (trim_end defaults to 0.0)
+    tool = TrimVideo(
+        product_name="Test_Product",
+        input_video="test_video",  # Will look for test_video.mp4 in generated_videos
+        output_name="test_video_trimmed",
+        # trim_start=1.0,  # Trim 1 second from start
+        trim_end=1.0
+    )
+    
+    try:
+        result = asyncio.run(tool.run())
+        print("\nTrim complete!")
+        print(result)
+    except Exception as exc:
+        print(f"Video trim failed: {exc}")
+

package/video_generation_agent/tools/__init__.py

@@ -0,0 +1 @@
+"""Tools for the video generation agent."""

package/video_generation_agent/tools/utils/__init__.py

@@ -0,0 +1 @@
+"""Utility helpers for video generation tools."""

package/video_generation_agent/tools/utils/image_utils.py

@@ -0,0 +1,174 @@
+import os
+import io
+import base64
+import logging
+from pathlib import Path
+import asyncio
+from PIL import Image
+
+logger = logging.getLogger(__name__)
+
+MODEL_NAME = "gemini-2.5-flash-image"
+
+if os.path.isfile("/.dockerenv"):
+    MNT_DIR = Path("/app/mnt")
+else:
+    MNT_DIR = Path(__file__).parent.parent.parent.parent / "mnt"
+
+IMAGES_DIR = str(MNT_DIR / "generated_images")
+
+OUTPUT_FORMAT = "png"
+
+
+def get_images_dir(product_name: str) -> str:
+    """Return (and create) the images directory for a specific product."""
+    images_dir = MNT_DIR / product_name / "generated_images"
+    images_dir.mkdir(parents=True, exist_ok=True)
+    return str(images_dir)
+
+
+def create_filename(file_name, variant_num, num_variants, output_format):
+    if num_variants == 1:
+        image_name = file_name
+    else:
+        image_name = f"{file_name}_variant_{variant_num}"
+    filename = f"{image_name}.{output_format}"
+    return image_name, filename
+
+
+def load_image_by_name(image_name, images_dir, extensions=None):
+    """Load an image by name, trying common extensions in order."""
+    if extensions is None:
+        extensions = ['.png', '.jpg', '.jpeg', '.webp']
+
+    for ext in extensions:
+        potential_path = os.path.join(images_dir, f"{image_name}{ext}")
+        if os.path.exists(potential_path):
+            try:
+                image = Image.open(potential_path)
+                return image, potential_path, None
+            except Exception as e:
+                return None, None, f"Error loading image {potential_path}: {str(e)}"
+
+    return None, None, f"Error: Image file not found: {image_name} (tried {', '.join(extensions)})"
+
+
+def extract_image_from_response(response):
+    """Extract the first image and any text from a Gemini API response."""
+    image = None
+    text_output = ""
+
+    for part in response.candidates[0].content.parts:
+        if part.text is not None:
+            text_output += part.text
+        elif part.inline_data is not None:
+            image = Image.open(io.BytesIO(part.inline_data.data))
+
+    return image, text_output
+
+
+def extract_image_parts_from_response(response):
+    """Extract raw image bytes from a Gemini API response."""
+    return [
+        part.inline_data.data
+        for part in response.candidates[0].content.parts
+        if part.inline_data
+    ]
+
+
+def extract_usage_metadata(response):
+    usage_metadata = getattr(response, "usage_metadata", {}) or {}
+    if usage_metadata and not isinstance(usage_metadata, dict):
+        if hasattr(usage_metadata, "model_dump"):
+            usage_metadata = usage_metadata.model_dump()
+        elif hasattr(usage_metadata, "__dict__"):
+            usage_metadata = vars(usage_metadata)
+        else:
+            try:
+                usage_metadata = dict(usage_metadata)
+            except (TypeError, ValueError):
+                usage_metadata = {}
+    return usage_metadata
+
+
+def split_results_and_usage(raw_results):
+    results = []
+    total_prompt_tokens = 0.0
+    total_candidate_tokens = 0.0
+
+    for item in raw_results:
+        result = dict(item)
+        total_prompt_tokens += float(result.pop("prompt_tokens", 0.0) or 0.0)
+        total_candidate_tokens += float(result.pop("candidate_tokens", 0.0) or 0.0)
+        results.append(result)
+
+    usage_metadata = {
+        "prompt_token_count": total_prompt_tokens,
+        "candidates_token_count": total_candidate_tokens,
+    }
+    return results, usage_metadata
+
+
+def process_variant_result(variant_num, image, file_name, num_variants, compress_func, images_dir=None):
+    """Save a variant image and return its result dict."""
+    save_dir = images_dir if images_dir is not None else IMAGES_DIR
+    image_name, filename = create_filename(file_name, variant_num, num_variants, OUTPUT_FORMAT)
+    filepath = os.path.join(save_dir, filename)
+    image.save(filepath, OUTPUT_FORMAT)
+    compressed_b64 = compress_func(image)
+    return {
+        "variant": variant_num,
+        "file_path": filepath,
+        "image_name": image_name,
+        "base64": compressed_b64,
+    }
+
+
+async def run_parallel_variants(variant_func, num_variants):
+    """Run variant_func for each variant concurrently in a thread pool."""
+    loop = asyncio.get_event_loop()
+    tasks = [loop.run_in_executor(None, variant_func, i + 1) for i in range(num_variants)]
+    completed = await asyncio.gather(*tasks, return_exceptions=True)
+    return [r for r in completed if r is not None and not isinstance(r, Exception)]
+
+
+def compress_image_for_base64(image: Image.Image, max_size: int = 1024) -> str:
+    """Resize and JPEG-compress an image, returning a base64-encoded string."""
+    if image.mode in ("RGBA", "P"):
+        image = image.convert("RGB")
+
+    width, height = image.size
+    if max(width, height) > max_size:
+        if width > height:
+            new_width = max_size
+            new_height = int(height * (max_size / width))
+        else:
+            new_height = max_size
+            new_width = int(width * (max_size / height))
+        image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
+
+    buffer = io.BytesIO()
+    image.save(buffer, format="JPEG", quality=85)
+    return base64.b64encode(buffer.getvalue()).decode()
+
+
+def combine_image_parts(image_parts):
+    """Combine multiple raw image parts into a single horizontally-stitched image."""
+    if not image_parts:
+        return None
+
+    images = [Image.open(io.BytesIO(part)) for part in image_parts]
+
+    if len(images) == 1:
+        return images[0]
+
+    total_width = sum(img.width for img in images)
+    max_height = max(img.height for img in images)
+    combined = Image.new('RGB', (total_width, max_height))
+
+    x_offset = 0
+    for img in images:
+        combined.paste(img, (x_offset, 0))
+        x_offset += img.width
+
+    return combined