@_vrsen/openswarm 0.1.0

package/image_generation_agent/tools/EditImages.py

@@ -0,0 +1,200 @@
+"""Edit images with Gemini or OpenAI image models."""
+
+from typing import Literal
+
+import os
+from openai import OpenAI
+from pydantic import Field, field_validator, model_validator
+
+from agency_swarm import BaseTool
+
+from .utils.image_io import (
+    get_images_dir,
+    build_variant_output_name,
+    resolve_image_reference,
+    save_image,
+    image_to_base64_jpeg,
+    build_multimodal_outputs,
+    extract_gemini_image_and_usage,
+    extract_openai_images_and_usage,
+    run_parallel_variants_sync,
+    validate_aspect_ratio_for_model,
+    get_openai_size_for_aspect_ratio,
+)
+
+
+class EditImages(BaseTool):
+    """
+    Edit an existing image using a text instruction.
+
+    Supports local path, URL, or generated image name references.
+    Outputs are saved to: mnt/{product_name}/generated_images/
+    """
+
+    product_name: str = Field(..., description="Product namespace for output files.")
+    input_image_ref: str = Field(
+        ...,
+        description="Input image reference (URL, absolute path, or generated image name).",
+    )
+    edit_prompt: str = Field(..., description="Instruction describing how to edit the image.")
+    output_file_name: str = Field(
+        ...,
+        description=(
+            "Output image name (without extension) or output path. "
+            "If a path is provided, the image is saved at that path."
+        ),
+    )
+    model: Literal["gemini-2.5-flash-image", "gemini-3-pro-image-preview", "gpt-image-1.5"] = Field(
+        default="gemini-2.5-flash-image",
+        description="Image model to use.",
+    )
+    num_variants: int = Field(default=1, description="Number of variants to generate (1-4).")
+    aspect_ratio: Literal["1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9"] = Field(
+        default="1:1",
+        description="Target aspect ratio. Model compatibility is validated automatically.",
+    )
+
+    @field_validator("product_name", "input_image_ref", "edit_prompt", "output_file_name")
+    @classmethod
+    def _not_blank(cls, value: str) -> str:
+        if not value.strip():
+            raise ValueError("value must not be empty")
+        return value
+
+    @field_validator("num_variants")
+    @classmethod
+    def _validate_variants(cls, value: int) -> int:
+        if value < 1 or value > 4:
+            raise ValueError("num_variants must be between 1 and 4")
+        return value
+
+    @model_validator(mode="after")
+    def _validate_model_aspect_ratio(self) -> "EditImages":
+        validate_aspect_ratio_for_model(self.model, self.aspect_ratio)
+        return self
+
+    def run(self) -> list:
+        images_dir = get_images_dir(self.product_name)
+        input_image, source = resolve_image_reference(self.product_name, self.input_image_ref)
+
+        if self.model.startswith("gemini-"):
+            results, usage_metadata = self._run_gemini(images_dir, input_image)
+            return build_multimodal_outputs(results, "Image editing complete")
+
+        results, usage_metadata = self._run_openai(images_dir, input_image)
+        return build_multimodal_outputs(results, "Image editing complete")
+
+    def _run_gemini(self, images_dir, input_image):
+        from google import genai
+        from google.genai.types import GenerateContentConfig, ImageConfig
+
+        api_key = os.getenv("GOOGLE_API_KEY")
+        if not api_key:
+            raise ValueError("GOOGLE_API_KEY is not set. Add it to your .env to use image editing.")
+
+        client = genai.Client(api_key=api_key)
+        results: list[dict] = []
+        total_prompt_tokens = 0.0
+        total_candidate_tokens = 0.0
+
+        def edit_single_variant(idx: int):
+            response = client.models.generate_content(
+                model=self.model,
+                contents=[self.edit_prompt, input_image],
+                config=GenerateContentConfig(
+                    image_config=ImageConfig(aspect_ratio=self.aspect_ratio),
+                ),
+            )
+            image, usage = extract_gemini_image_and_usage(response)
+            if image is None:
+                return None
+
+            variant_name = build_variant_output_name(self.output_file_name, idx, self.num_variants)
+            image_name, file_path = save_image(image, variant_name, images_dir)
+            return {
+                "image_name": image_name,
+                "file_path": file_path,
+                "preview_b64": image_to_base64_jpeg(image),
+                "prompt_tokens": float(usage.get("prompt_token_count") or 0),
+                "candidate_tokens": float(usage.get("candidates_token_count") or 0),
+            }
+
+        raw_results = run_parallel_variants_sync(edit_single_variant, self.num_variants)
+        if not raw_results:
+            raise RuntimeError("Gemini did not return any edited images.")
+
+        for item in raw_results:
+            total_prompt_tokens += item.pop("prompt_tokens")
+            total_candidate_tokens += item.pop("candidate_tokens")
+            results.append(item)
+
+        usage_metadata = {
+            "prompt_token_count": total_prompt_tokens,
+            "candidates_token_count": total_candidate_tokens,
+        }
+        return results, usage_metadata
+
+    def _run_openai(self, images_dir, input_image):
+        api_key = os.getenv("OPENAI_API_KEY")
+        if not api_key:
+            raise RuntimeError("OPENAI_API_KEY is required for OpenAI image editing.")
+
+        size = get_openai_size_for_aspect_ratio(self.aspect_ratio)
+
+        from io import BytesIO
+
+        buffer = BytesIO()
+        input_image.save(buffer, format="PNG")
+        buffer.seek(0)
+        buffer.name = "input.png"
+
+        try:
+            client = OpenAI(api_key=api_key)
+            response = client.images.edit(
+                model=self.model,
+                image=buffer,
+                prompt=self.edit_prompt,
+                size=size,
+                n=self.num_variants,
+            )
+        finally:
+            buffer.close()
+
+        images, usage_metadata = extract_openai_images_and_usage(response)
+        if not images:
+            raise RuntimeError("OpenAI image edit API did not return images.")
+
+        results: list[dict] = []
+        for idx, image in enumerate(images, start=1):
+            variant_name = build_variant_output_name(self.output_file_name, idx, len(images))
+            image_name, file_path = save_image(image, variant_name, images_dir)
+            results.append(
+                {
+                    "image_name": image_name,
+                    "file_path": file_path,
+                    "preview_b64": image_to_base64_jpeg(image),
+                }
+            )
+
+        return results, usage_metadata
+
+
+if __name__ == "__main__":
+    # Example test scenario
+    tool = EditImages(
+        product_name="Test_Product",
+        input_image_ref="hero_image_example_oai_hero",
+        edit_prompt=(
+            "Replace the background with a black wall, slightly illuminated with a white neon light"
+        ),
+        output_file_name="edited_image_example",
+        model="gpt-image-1.5",
+        aspect_ratio="1:1",
+        num_variants=1,
+    )
+    try:
+        result = tool.run()
+        print(result)
+    except Exception as exc:
+        print(f"Image editing failed: {exc}")
+

package/image_generation_agent/tools/GenerateImages.py

@@ -0,0 +1,184 @@
+"""Generate images with Gemini or OpenAI image models."""
+
+from typing import Literal
+
+import os
+from openai import OpenAI
+from pydantic import Field, field_validator, model_validator
+
+from agency_swarm import BaseTool
+
+from .utils.image_io import (
+    get_images_dir,
+    build_variant_output_name,
+    save_image,
+    image_to_base64_jpeg,
+    build_multimodal_outputs,
+    extract_gemini_image_and_usage,
+    extract_openai_images_and_usage,
+    run_parallel_variants_sync,
+    validate_aspect_ratio_for_model,
+    get_openai_size_for_aspect_ratio,
+)
+
+
+class GenerateImages(BaseTool):
+    """
+    Generate one or more images from a prompt using Gemini or OpenAI image models.
+
+    Outputs are saved to: mnt/{product_name}/generated_images/
+    """
+
+    product_name: str = Field(..., description="Product namespace for output files.")
+    prompt: str = Field(..., description="Detailed prompt describing the desired image.")
+    file_name: str = Field(
+        ...,
+        description=(
+            "Output file name (without extension) or output path. "
+            "If a path is provided, the image is saved at that path."
+        ),
+    )
+    model: Literal["gemini-2.5-flash-image", "gemini-3-pro-image-preview", "gpt-image-1.5"] = Field(
+        default="gemini-2.5-flash-image",
+        description="Image model to use.",
+    )
+    num_variants: int = Field(default=1, description="Number of variants to generate (1-4).")
+    aspect_ratio: Literal["1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9"] = Field(
+        default="1:1",
+        description="Target aspect ratio.",
+    )
+
+    @field_validator("prompt", "product_name", "file_name")
+    @classmethod
+    def _not_blank(cls, value: str) -> str:
+        if not value.strip():
+            raise ValueError("value must not be empty")
+        return value
+
+    @field_validator("num_variants")
+    @classmethod
+    def _validate_variants(cls, value: int) -> int:
+        if value < 1 or value > 4:
+            raise ValueError("num_variants must be between 1 and 4")
+        return value
+
+    @model_validator(mode="after")
+    def _validate_model_aspect_ratio(self) -> "GenerateImages":
+        validate_aspect_ratio_for_model(self.model, self.aspect_ratio)
+        return self
+
+    def run(self) -> list:
+        images_dir = get_images_dir(self.product_name)
+
+        if self.model.startswith("gemini-"):
+            results, usage_metadata = self._run_gemini(images_dir)
+            return build_multimodal_outputs(results, "Image generation complete")
+
+        results, usage_metadata = self._run_openai(images_dir)
+        return build_multimodal_outputs(results, "Image generation complete")
+
+    def _run_gemini(self, images_dir):
+        from google import genai
+        from google.genai.types import GenerateContentConfig, ImageConfig
+
+        api_key = os.getenv("GOOGLE_API_KEY")
+        if not api_key:
+            raise ValueError("GOOGLE_API_KEY is not set. Add it to your .env to use image generation.")
+
+        client = genai.Client(api_key=api_key)
+        results: list[dict] = []
+        total_prompt_tokens = 0.0
+        total_candidate_tokens = 0.0
+
+        def generate_single_variant(idx: int):
+            response = client.models.generate_content(
+                model=self.model,
+                contents=self.prompt,
+                config=GenerateContentConfig(
+                    image_config=ImageConfig(aspect_ratio=self.aspect_ratio),
+                ),
+            )
+            image, usage = extract_gemini_image_and_usage(response)
+            if image is None:
+                return None
+
+            variant_name = build_variant_output_name(self.file_name, idx, self.num_variants)
+            image_name, file_path = save_image(image, variant_name, images_dir)
+            return {
+                "image_name": image_name,
+                "file_path": file_path,
+                "preview_b64": image_to_base64_jpeg(image),
+                "prompt_tokens": float(usage.get("prompt_token_count") or 0),
+                "candidate_tokens": float(usage.get("candidates_token_count") or 0),
+            }
+
+        raw_results = run_parallel_variants_sync(generate_single_variant, self.num_variants)
+        if not raw_results:
+            raise RuntimeError("Gemini did not return any images.")
+
+        for item in raw_results:
+            total_prompt_tokens += item.pop("prompt_tokens")
+            total_candidate_tokens += item.pop("candidate_tokens")
+            results.append(item)
+
+        usage_metadata = {
+            "prompt_token_count": total_prompt_tokens,
+            "candidates_token_count": total_candidate_tokens,
+        }
+        return results, usage_metadata
+
+    def _run_openai(self, images_dir):
+        api_key = os.getenv("OPENAI_API_KEY")
+        if not api_key:
+            raise RuntimeError("OPENAI_API_KEY is required for OpenAI image generation.")
+
+        size = get_openai_size_for_aspect_ratio(self.aspect_ratio)
+
+        client = OpenAI(api_key=api_key)
+        response = client.images.generate(
+            model=self.model,
+            prompt=self.prompt,
+            size=size,
+            n=self.num_variants,
+        )
+        images, usage_metadata = extract_openai_images_and_usage(response)
+        if not images:
+            raise RuntimeError("OpenAI image API did not return images.")
+
+        results: list[dict] = []
+        for idx, image in enumerate(images, start=1):
+            variant_name = build_variant_output_name(self.file_name, idx, len(images))
+            image_name, file_path = save_image(image, variant_name, images_dir)
+            results.append(
+                {
+                    "image_name": image_name,
+                    "file_path": file_path,
+                    "preview_b64": image_to_base64_jpeg(image),
+                }
+            )
+        return results, usage_metadata
+
+
+if __name__ == "__main__":
+    # Example test scenario
+    tool = GenerateImages(
+        product_name="Test_Product",
+        prompt=(
+            "A premium product hero shot of a minimalist glass skincare bottle on a "
+            "clean marble surface, soft studio lighting, subtle reflections, "
+            "high-end commercial photography style."
+        ),
+        # prompt=(
+        #     "Create a logo for a skincare product that includes a leaf and a drop of water"
+        # ),
+        file_name="hero_image_example_oai_hero",
+        model="gpt-image-1.5",
+        aspect_ratio="1:1",
+        num_variants=1,
+    )
+    try:
+        result = tool.run()
+        print(result)
+    except Exception as exc:
+        print(f"Image generation failed: {exc}")
+

package/image_generation_agent/tools/RemoveBackground.py

@@ -0,0 +1,108 @@
+"""Remove image backgrounds using the Pixelcut model via fal.ai."""
+
+import io
+import os
+from pathlib import Path
+from urllib.parse import urlparse
+
+import fal_client
+import requests
+from PIL import Image
+from pydantic import Field, field_validator
+
+from agency_swarm import BaseTool
+
+from .utils.image_io import (
+    find_image_path_from_name,
+    get_images_dir,
+    save_image,
+)
+
+FAL_ENDPOINT = "pixelcut/background-removal"
+
+
+class RemoveBackground(BaseTool):
+    """
+    Remove the background from an image using Pixelcut via fal.ai.
+
+    The output is a transparent PNG saved to: mnt/{product_name}/generated_images/
+    Supports local path, URL, or generated image name as input.
+    """
+
+    product_name: str = Field(..., description="Product namespace for output files.")
+    input_image_ref: str = Field(
+        ...,
+        description="Input image reference (URL, absolute path, or generated image name).",
+    )
+    output_file_name: str = Field(
+        ...,
+        description=(
+            "Output image name (without extension) or output path. "
+            "If a path is provided, the image is saved at that path."
+        ),
+    )
+
+    @field_validator("product_name", "input_image_ref", "output_file_name")
+    @classmethod
+    def _not_blank(cls, value: str) -> str:
+        if not value.strip():
+            raise ValueError("value must not be empty")
+        return value
+
+    def run(self) -> list:
+        api_key = os.getenv("FAL_KEY")
+        if not api_key:
+            raise ValueError("FAL_KEY is not set. Add it to your .env to use background removal.")
+        fal = fal_client.SyncClient(key=api_key)
+
+        images_dir = get_images_dir(self.product_name)
+        image_url = self._resolve_to_upload_url(images_dir, fal)
+        result = fal.subscribe(
+            FAL_ENDPOINT,
+            arguments={"image_url": image_url, "output_format": "rgba", "sync_mode": False},
+        )
+
+        result_url = (result.get("image") or {}).get("url")
+        if not result_url:
+            raise RuntimeError("fal.ai background removal returned no image URL.")
+
+        rgba_image = self._download_rgba(result_url)
+        image_name, file_path = save_image(rgba_image, self.output_file_name, images_dir)
+
+        return f"Background removal complete.\nImage name: {image_name}\nPath: {file_path}"
+
+    def _resolve_to_upload_url(self, images_dir: Path, fal: fal_client.SyncClient) -> str:
+        ref = self.input_image_ref.strip()
+
+        parsed = urlparse(ref)
+        if parsed.scheme in ("http", "https"):
+            return ref
+
+        candidate = Path(ref).expanduser().resolve()
+        if candidate.exists():
+            return fal.upload_file(str(candidate))
+
+        by_name = find_image_path_from_name(images_dir, ref)
+        if by_name is not None:
+            return fal.upload_file(str(by_name))
+
+        raise FileNotFoundError(
+            f"Could not resolve image reference '{ref}' as URL, path, or name in {images_dir}."
+        )
+
+    def _download_rgba(self, url: str) -> Image.Image:
+        response = requests.get(url, timeout=30)
+        response.raise_for_status()
+        return Image.open(io.BytesIO(response.content)).convert("RGBA")
+
+if __name__ == "__main__":
+    tool = RemoveBackground(
+        product_name="Test_Product",
+        input_image_ref="test_image.jpg",
+        output_file_name="hero_no_bg",
+    )
+    try:
+        result = tool.run()
+        print(result)
+    except Exception as exc:
+        print(f"Background removal failed: {exc}")

package/image_generation_agent/tools/__init__.py

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

package/image_generation_agent/tools/utils/__init__.py

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